From 765a7af6db9ede660dbae65a665ff8896da4ba49 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 14 Dec 2022 20:23:12 +0000 Subject: [PATCH 01/80] chore: add license and readme Change-Id: I7a5a716eac371b3ab765743500bb109ea5f47de1 --- LICENSE | 202 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ README.md | 3 + 2 files changed, 205 insertions(+) create mode 100644 LICENSE create mode 100644 README.md diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000000..7a4a3ea2424 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 00000000000..3b3fd5761d9 --- /dev/null +++ b/README.md @@ -0,0 +1,3 @@ +# BigFrames + +BigFrames supports scalable DataFrame APIs on top of BigQuery. From b6caad63cbae63123645a7c560e3242b47f2ca4c Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 20 Jul 2023 16:31:09 -0500 Subject: [PATCH 02/80] feat: prepare for initial release Includes code from all commits in internal repo up to 6a7b313e1d71c523ab039bced341518b4a6928bc Prepared in the following way: ``` git fetch github main git checkout -b prepare-0.1.0 github/main git merge --squash 6a7b313e1d71c523ab039bced341518b4a6928bc ``` --- .flake8 | 35 + .gitignore | 70 + .isort.cfg | 7 + .kokoro/build.sh | 47 + .kokoro/continuous/common.cfg | 10 + .kokoro/continuous/continuous.cfg | 1 + .kokoro/continuous/e2e.cfg | 7 + .kokoro/continuous/nightly.cfg | 8 + .kokoro/presubmit/common.cfg | 10 + .kokoro/presubmit/e2e.cfg | 7 + .kokoro/presubmit/presubmit.cfg | 1 + .kokoro/release-nightly.sh | 213 + .kokoro/requirements.in | 13 + .kokoro/requirements.txt | 567 +++ .pre-commit-config.yaml | 41 + CHANGELOG.md | 22 + CODE_OF_CONDUCT.md | 94 + CONTRIBUTING.md | 33 + LICENSE | 2 +- OWNERS | 10 + README.md | 3 - README.rst | 8 + bigframes/__init__.py | 28 + bigframes/_config/__init__.py | 49 + bigframes/_config/bigquery_options.py | 114 + bigframes/_config/display_options.py | 55 + bigframes/core/__init__.py | 1005 ++++ bigframes/core/block_transforms.py | 199 + bigframes/core/blocks.py | 917 ++++ bigframes/core/groupby/__init__.py | 315 ++ bigframes/core/guid.py | 21 + bigframes/core/indexers.py | 337 ++ bigframes/core/indexes/__init__.py | 23 + bigframes/core/indexes/index.py | 250 + bigframes/core/joins/__init__.py | 23 + bigframes/core/joins/row_identity.py | 201 + bigframes/core/joins/single_column.py | 281 ++ bigframes/core/ordering.py | 181 + bigframes/core/reshape/__init__.py | 61 + bigframes/core/scalar.py | 48 + bigframes/core/window/__init__.py | 85 + bigframes/dataframe.py | 1626 +++++++ bigframes/dtypes.py | 304 ++ bigframes/formatting_helpers.py | 120 + bigframes/ml/__init__.py | 29 + bigframes/ml/base.py | 135 + bigframes/ml/cluster.py | 97 + bigframes/ml/compose.py | 123 + bigframes/ml/core.py | 334 ++ bigframes/ml/decomposition.py | 104 + bigframes/ml/ensemble.py | 672 +++ bigframes/ml/forecasting.py | 125 + bigframes/ml/imported.py | 101 + bigframes/ml/linear_model.py | 247 + bigframes/ml/llm.py | 168 + bigframes/ml/loader.py | 74 + bigframes/ml/metrics.py | 371 ++ bigframes/ml/model_selection.py | 86 + bigframes/ml/pipeline.py | 102 + bigframes/ml/preprocessing.py | 140 + bigframes/ml/sql.py | 173 + bigframes/operations/__init__.py | 546 +++ bigframes/operations/aggregations.py | 288 ++ bigframes/operations/base.py | 198 + bigframes/operations/datetimes.py | 66 + bigframes/operations/strings.py | 81 + bigframes/pandas/__init__.py | 339 ++ bigframes/remote_function.py | 697 +++ bigframes/series.py | 964 ++++ bigframes/session.py | 1055 +++++ bigframes/version.py | 15 + docs/Makefile | 20 + docs/README.rst | 1 + docs/_static/custom.css | 15 + docs/_templates/layout.html | 45 + docs/changelog.md | 1 + docs/conf.py | 382 ++ docs/getting_started/index.rst | 27 + docs/index.rst | 18 + docs/reference/bigframes.ml/README.rst | 125 + docs/reference/bigframes.ml/cluster.rst | 5 + docs/reference/bigframes.ml/compose.rst | 5 + docs/reference/bigframes.ml/decomposition.rst | 5 + docs/reference/bigframes.ml/forecasting.rst | 5 + docs/reference/bigframes.ml/imported.rst | 5 + docs/reference/bigframes.ml/index.rst | 30 + docs/reference/bigframes.ml/linear_model.rst | 5 + docs/reference/bigframes.ml/llm.rst | 5 + docs/reference/bigframes.ml/metrics.rst | 5 + .../bigframes.ml/model_selection.rst | 5 + docs/reference/bigframes.ml/pipeline.rst | 5 + docs/reference/bigframes.ml/preprocessing.rst | 5 + docs/reference/bigframes.pandas/frame.rst | 9 + .../bigframes.pandas/general_functions.rst | 8 + docs/reference/bigframes.pandas/groupby.rst | 20 + docs/reference/bigframes.pandas/index.rst | 15 + docs/reference/bigframes.pandas/indexing.rst | 9 + docs/reference/bigframes.pandas/options.rst | 6 + docs/reference/bigframes.pandas/series.rst | 36 + docs/reference/bigframes.pandas/window.rst | 9 + docs/reference/bigframes/index.rst | 19 + docs/reference/bigframes/options.rst | 12 + docs/reference/index.rst | 12 + docs/samples | 1 + .../bigframes.pandas/remote_functions.rst | 40 + docs/user_guide/index.rst | 9 + mypy.ini | 25 + notebooks/00 - Summary.ipynb | 2060 ++++++++ notebooks/01 - Getting Started.ipynb | 965 ++++ notebooks/02 - DataFrame.ipynb | 4158 +++++++++++++++++ .../03 - Using ML - ML fundamentals.ipynb | 2250 +++++++++ ...Using ML - SKLearn linear regression.ipynb | 977 ++++ ... - Using ML - Easy linear regression.ipynb | 920 ++++ ...6 - Using ML - Large Language Models.ipynb | 201 + notebooks/50 - Remote Function.ipynb | 839 ++++ notebooks/99 - Longer ML demo.ipynb | 1925 ++++++++ noxfile.py | 636 +++ pytest.ini | 3 + samples/snippets/__init__.py | 13 + samples/snippets/quickstart.py | 89 + samples/snippets/quickstart_test.py | 31 + scripts/generate_third_party_notices.py | 332 ++ scripts/publish_api_coverage.py | 179 + scripts/test_publish_api_coverage.py | 45 + scripts/update_firebase_docs_site.sh | 105 + scripts/update_x20_docs_site.sh | 106 + setup.cfg | 33 + setup.py | 109 + testing/constraints-3.10.txt | 15 + testing/constraints-3.11.txt | 0 testing/constraints-3.9.txt | 118 + tests/__init__.py | 13 + tests/data/penguins.jsonl | 344 ++ tests/data/penguins_schema.json | 37 + tests/data/scalars.jsonl | 9 + tests/data/scalars_schema.json | 75 + tests/data/time_series.jsonl | 366 ++ tests/data/time_series_schema.json | 12 + tests/system/__init__.py | 13 + tests/system/conftest.py | 667 +++ tests/system/large/__init__.py | 13 + tests/system/large/ml/test_cluster.py | 103 + tests/system/large/ml/test_compose.py | 71 + tests/system/large/ml/test_core.py | 183 + tests/system/large/ml/test_decomposition.py | 70 + tests/system/large/ml/test_ensemble.py | 458 ++ tests/system/large/ml/test_forecasting.py | 48 + tests/system/large/ml/test_linear_model.py | 196 + tests/system/large/ml/test_pipeline.py | 524 +++ tests/system/large/test_remote_function.py | 776 +++ tests/system/large/test_session.py | 52 + tests/system/small/__init__.py | 13 + tests/system/small/ml/__init__.py | 13 + tests/system/small/ml/conftest.py | 294 ++ tests/system/small/ml/test_cluster.py | 79 + tests/system/small/ml/test_core.py | 193 + tests/system/small/ml/test_decomposition.py | 55 + tests/system/small/ml/test_ensemble.py | 427 ++ tests/system/small/ml/test_forecasting.py | 44 + tests/system/small/ml/test_imported.py | 68 + tests/system/small/ml/test_linear_model.py | 209 + tests/system/small/ml/test_llm.py | 72 + tests/system/small/ml/test_metrics.py | 615 +++ tests/system/small/ml/test_model_selection.py | 214 + tests/system/small/ml/test_preprocessing.py | 89 + tests/system/small/ml/test_register.py | 87 + tests/system/small/operations/__init__.py | 13 + .../system/small/operations/test_datetimes.py | 198 + tests/system/small/operations/test_strings.py | 230 + tests/system/small/test_dataframe.py | 1694 +++++++ tests/system/small/test_dataframe_io.py | 325 ++ tests/system/small/test_groupby.py | 129 + tests/system/small/test_index.py | 46 + tests/system/small/test_multiindex.py | 273 ++ tests/system/small/test_pandas.py | 64 + tests/system/small/test_pandas_options.py | 145 + tests/system/small/test_progress_bar.py | 74 + tests/system/small/test_remote_function.py | 278 ++ tests/system/small/test_scalar.py | 25 + tests/system/small/test_series.py | 1934 ++++++++ tests/system/small/test_session.py | 673 +++ tests/system/small/test_window.py | 66 + tests/system/utils.py | 141 + tests/unit/__init__.py | 13 + tests/unit/_config/__init__.py | 13 + tests/unit/_config/test_bigquery_options.py | 73 + tests/unit/conftest.py | 223 + tests/unit/ml/__init__.py | 13 + tests/unit/ml/test_api_primitives.py | 53 + tests/unit/ml/test_compose.py | 95 + tests/unit/ml/test_pipeline.py | 110 + tests/unit/ml/test_sql.py | 145 + tests/unit/test_core.py | 99 + tests/unit/test_dtypes.py | 184 + tests/unit/test_pandas.py | 81 + tests/unit/test_session.py | 48 + third_party/__init__.py | 0 third_party/bigframes_vendored/__init__.py | 0 .../bigframes_vendored/cpython/LICENSE | 10 + .../bigframes_vendored/cpython/__init__.py | 0 .../bigframes_vendored/cpython/_pprint.py | 470 ++ .../bigframes_vendored/pandas/AUTHORS.md | 56 + third_party/bigframes_vendored/pandas/LICENSE | 31 + .../bigframes_vendored/pandas/README.md | 174 + .../bigframes_vendored/pandas/__init__.py | 0 .../pandas/core/config_init.py | 25 + .../bigframes_vendored/pandas/core/frame.py | 1124 +++++ .../bigframes_vendored/pandas/core/generic.py | 438 ++ .../pandas/core/groupby/__init__.py | 318 ++ .../pandas/core/indexes/__init__.py | 0 .../pandas/core/indexes/accessor.py | 94 + .../pandas/core/indexes/base.py | 20 + .../pandas/core/indexing.py | 72 + .../pandas/core/reshape/concat.py | 135 + .../bigframes_vendored/pandas/core/series.py | 1553 ++++++ .../pandas/core/strings/accessor.py | 187 + .../pandas/core/window/__init__.py | 0 .../pandas/core/window/rolling.py | 37 + .../bigframes_vendored/pandas/io/__init__.py | 0 .../bigframes_vendored/pandas/io/common.py | 44 + .../bigframes_vendored/pandas/io/gbq.py | 37 + .../bigframes_vendored/pandas/io/parquet.py | 22 + .../pandas/io/parsers/__init__.py | 0 .../pandas/io/parsers/readers.py | 110 + .../pandas/pandas/_typing.py | 425 ++ .../bigframes_vendored/sklearn/COPYING | 29 + .../bigframes_vendored/sklearn/base.py | 128 + .../sklearn/cluster/_kmeans.py | 71 + .../sklearn/compose/_column_transformer.py | 61 + .../sklearn/decomposition/_pca.py | 53 + .../sklearn/ensemble/_forest.py | 191 + .../sklearn/linear_model/_base.py | 98 + .../sklearn/linear_model/_logistic.py | 83 + .../sklearn/metrics/_classification.py | 173 + .../sklearn/metrics/_ranking.py | 100 + .../sklearn/metrics/_regression.py | 53 + .../bigframes_vendored/sklearn/pipeline.py | 83 + .../sklearn/preprocessing/_data.py | 53 + .../sklearn/preprocessing/_encoder.py | 41 + .../bigframes_vendored/xgboost/LICENSE | 201 + .../bigframes_vendored/xgboost/sklearn.py | 145 + 241 files changed, 52741 insertions(+), 4 deletions(-) create mode 100644 .flake8 create mode 100644 .gitignore create mode 100644 .isort.cfg create mode 100644 .kokoro/build.sh create mode 100644 .kokoro/continuous/common.cfg create mode 100644 .kokoro/continuous/continuous.cfg create mode 100644 .kokoro/continuous/e2e.cfg create mode 100644 .kokoro/continuous/nightly.cfg create mode 100644 .kokoro/presubmit/common.cfg create mode 100644 .kokoro/presubmit/e2e.cfg create mode 100644 .kokoro/presubmit/presubmit.cfg create mode 100755 .kokoro/release-nightly.sh create mode 100644 .kokoro/requirements.in create mode 100644 .kokoro/requirements.txt create mode 100644 .pre-commit-config.yaml create mode 100644 CHANGELOG.md create mode 100644 CODE_OF_CONDUCT.md create mode 100644 CONTRIBUTING.md create mode 100644 OWNERS delete mode 100644 README.md create mode 100644 README.rst create mode 100644 bigframes/__init__.py create mode 100644 bigframes/_config/__init__.py create mode 100644 bigframes/_config/bigquery_options.py create mode 100644 bigframes/_config/display_options.py create mode 100644 bigframes/core/__init__.py create mode 100644 bigframes/core/block_transforms.py create mode 100644 bigframes/core/blocks.py create mode 100644 bigframes/core/groupby/__init__.py create mode 100644 bigframes/core/guid.py create mode 100644 bigframes/core/indexers.py create mode 100644 bigframes/core/indexes/__init__.py create mode 100644 bigframes/core/indexes/index.py create mode 100644 bigframes/core/joins/__init__.py create mode 100644 bigframes/core/joins/row_identity.py create mode 100644 bigframes/core/joins/single_column.py create mode 100644 bigframes/core/ordering.py create mode 100644 bigframes/core/reshape/__init__.py create mode 100644 bigframes/core/scalar.py create mode 100644 bigframes/core/window/__init__.py create mode 100644 bigframes/dataframe.py create mode 100644 bigframes/dtypes.py create mode 100644 bigframes/formatting_helpers.py create mode 100644 bigframes/ml/__init__.py create mode 100644 bigframes/ml/base.py create mode 100644 bigframes/ml/cluster.py create mode 100644 bigframes/ml/compose.py create mode 100644 bigframes/ml/core.py create mode 100644 bigframes/ml/decomposition.py create mode 100644 bigframes/ml/ensemble.py create mode 100644 bigframes/ml/forecasting.py create mode 100644 bigframes/ml/imported.py create mode 100644 bigframes/ml/linear_model.py create mode 100644 bigframes/ml/llm.py create mode 100644 bigframes/ml/loader.py create mode 100644 bigframes/ml/metrics.py create mode 100644 bigframes/ml/model_selection.py create mode 100644 bigframes/ml/pipeline.py create mode 100644 bigframes/ml/preprocessing.py create mode 100644 bigframes/ml/sql.py create mode 100644 bigframes/operations/__init__.py create mode 100644 bigframes/operations/aggregations.py create mode 100644 bigframes/operations/base.py create mode 100644 bigframes/operations/datetimes.py create mode 100644 bigframes/operations/strings.py create mode 100644 bigframes/pandas/__init__.py create mode 100644 bigframes/remote_function.py create mode 100644 bigframes/series.py create mode 100644 bigframes/session.py create mode 100644 bigframes/version.py create mode 100644 docs/Makefile create mode 120000 docs/README.rst create mode 100644 docs/_static/custom.css create mode 100644 docs/_templates/layout.html create mode 120000 docs/changelog.md create mode 100644 docs/conf.py create mode 100644 docs/getting_started/index.rst create mode 100644 docs/index.rst create mode 100644 docs/reference/bigframes.ml/README.rst create mode 100644 docs/reference/bigframes.ml/cluster.rst create mode 100644 docs/reference/bigframes.ml/compose.rst create mode 100644 docs/reference/bigframes.ml/decomposition.rst create mode 100644 docs/reference/bigframes.ml/forecasting.rst create mode 100644 docs/reference/bigframes.ml/imported.rst create mode 100644 docs/reference/bigframes.ml/index.rst create mode 100644 docs/reference/bigframes.ml/linear_model.rst create mode 100644 docs/reference/bigframes.ml/llm.rst create mode 100644 docs/reference/bigframes.ml/metrics.rst create mode 100644 docs/reference/bigframes.ml/model_selection.rst create mode 100644 docs/reference/bigframes.ml/pipeline.rst create mode 100644 docs/reference/bigframes.ml/preprocessing.rst create mode 100644 docs/reference/bigframes.pandas/frame.rst create mode 100644 docs/reference/bigframes.pandas/general_functions.rst create mode 100644 docs/reference/bigframes.pandas/groupby.rst create mode 100644 docs/reference/bigframes.pandas/index.rst create mode 100644 docs/reference/bigframes.pandas/indexing.rst create mode 100644 docs/reference/bigframes.pandas/options.rst create mode 100644 docs/reference/bigframes.pandas/series.rst create mode 100644 docs/reference/bigframes.pandas/window.rst create mode 100644 docs/reference/bigframes/index.rst create mode 100644 docs/reference/bigframes/options.rst create mode 100644 docs/reference/index.rst create mode 120000 docs/samples create mode 100644 docs/user_guide/bigframes.pandas/remote_functions.rst create mode 100644 docs/user_guide/index.rst create mode 100644 mypy.ini create mode 100644 notebooks/00 - Summary.ipynb create mode 100644 notebooks/01 - Getting Started.ipynb create mode 100644 notebooks/02 - DataFrame.ipynb create mode 100644 notebooks/03 - Using ML - ML fundamentals.ipynb create mode 100644 notebooks/04 - Using ML - SKLearn linear regression.ipynb create mode 100644 notebooks/05 - Using ML - Easy linear regression.ipynb create mode 100644 notebooks/06 - Using ML - Large Language Models.ipynb create mode 100644 notebooks/50 - Remote Function.ipynb create mode 100644 notebooks/99 - Longer ML demo.ipynb create mode 100644 noxfile.py create mode 100644 pytest.ini create mode 100644 samples/snippets/__init__.py create mode 100644 samples/snippets/quickstart.py create mode 100644 samples/snippets/quickstart_test.py create mode 100644 scripts/generate_third_party_notices.py create mode 100644 scripts/publish_api_coverage.py create mode 100644 scripts/test_publish_api_coverage.py create mode 100644 scripts/update_firebase_docs_site.sh create mode 100644 scripts/update_x20_docs_site.sh create mode 100644 setup.cfg create mode 100644 setup.py create mode 100644 testing/constraints-3.10.txt create mode 100644 testing/constraints-3.11.txt create mode 100644 testing/constraints-3.9.txt create mode 100644 tests/__init__.py create mode 100644 tests/data/penguins.jsonl create mode 100644 tests/data/penguins_schema.json create mode 100644 tests/data/scalars.jsonl create mode 100644 tests/data/scalars_schema.json create mode 100644 tests/data/time_series.jsonl create mode 100644 tests/data/time_series_schema.json create mode 100644 tests/system/__init__.py create mode 100644 tests/system/conftest.py create mode 100644 tests/system/large/__init__.py create mode 100644 tests/system/large/ml/test_cluster.py create mode 100644 tests/system/large/ml/test_compose.py create mode 100644 tests/system/large/ml/test_core.py create mode 100644 tests/system/large/ml/test_decomposition.py create mode 100644 tests/system/large/ml/test_ensemble.py create mode 100644 tests/system/large/ml/test_forecasting.py create mode 100644 tests/system/large/ml/test_linear_model.py create mode 100644 tests/system/large/ml/test_pipeline.py create mode 100644 tests/system/large/test_remote_function.py create mode 100644 tests/system/large/test_session.py create mode 100644 tests/system/small/__init__.py create mode 100644 tests/system/small/ml/__init__.py create mode 100644 tests/system/small/ml/conftest.py create mode 100644 tests/system/small/ml/test_cluster.py create mode 100644 tests/system/small/ml/test_core.py create mode 100644 tests/system/small/ml/test_decomposition.py create mode 100644 tests/system/small/ml/test_ensemble.py create mode 100644 tests/system/small/ml/test_forecasting.py create mode 100644 tests/system/small/ml/test_imported.py create mode 100644 tests/system/small/ml/test_linear_model.py create mode 100644 tests/system/small/ml/test_llm.py create mode 100644 tests/system/small/ml/test_metrics.py create mode 100644 tests/system/small/ml/test_model_selection.py create mode 100644 tests/system/small/ml/test_preprocessing.py create mode 100644 tests/system/small/ml/test_register.py create mode 100644 tests/system/small/operations/__init__.py create mode 100644 tests/system/small/operations/test_datetimes.py create mode 100644 tests/system/small/operations/test_strings.py create mode 100644 tests/system/small/test_dataframe.py create mode 100644 tests/system/small/test_dataframe_io.py create mode 100644 tests/system/small/test_groupby.py create mode 100644 tests/system/small/test_index.py create mode 100644 tests/system/small/test_multiindex.py create mode 100644 tests/system/small/test_pandas.py create mode 100644 tests/system/small/test_pandas_options.py create mode 100644 tests/system/small/test_progress_bar.py create mode 100644 tests/system/small/test_remote_function.py create mode 100644 tests/system/small/test_scalar.py create mode 100644 tests/system/small/test_series.py create mode 100644 tests/system/small/test_session.py create mode 100644 tests/system/small/test_window.py create mode 100644 tests/system/utils.py create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/_config/__init__.py create mode 100644 tests/unit/_config/test_bigquery_options.py create mode 100644 tests/unit/conftest.py create mode 100644 tests/unit/ml/__init__.py create mode 100644 tests/unit/ml/test_api_primitives.py create mode 100644 tests/unit/ml/test_compose.py create mode 100644 tests/unit/ml/test_pipeline.py create mode 100644 tests/unit/ml/test_sql.py create mode 100644 tests/unit/test_core.py create mode 100644 tests/unit/test_dtypes.py create mode 100644 tests/unit/test_pandas.py create mode 100644 tests/unit/test_session.py create mode 100644 third_party/__init__.py create mode 100644 third_party/bigframes_vendored/__init__.py create mode 100644 third_party/bigframes_vendored/cpython/LICENSE create mode 100644 third_party/bigframes_vendored/cpython/__init__.py create mode 100644 third_party/bigframes_vendored/cpython/_pprint.py create mode 100644 third_party/bigframes_vendored/pandas/AUTHORS.md create mode 100644 third_party/bigframes_vendored/pandas/LICENSE create mode 100644 third_party/bigframes_vendored/pandas/README.md create mode 100644 third_party/bigframes_vendored/pandas/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/config_init.py create mode 100644 third_party/bigframes_vendored/pandas/core/frame.py create mode 100644 third_party/bigframes_vendored/pandas/core/generic.py create mode 100644 third_party/bigframes_vendored/pandas/core/groupby/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/indexes/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/indexes/accessor.py create mode 100644 third_party/bigframes_vendored/pandas/core/indexes/base.py create mode 100644 third_party/bigframes_vendored/pandas/core/indexing.py create mode 100644 third_party/bigframes_vendored/pandas/core/reshape/concat.py create mode 100644 third_party/bigframes_vendored/pandas/core/series.py create mode 100644 third_party/bigframes_vendored/pandas/core/strings/accessor.py create mode 100644 third_party/bigframes_vendored/pandas/core/window/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/window/rolling.py create mode 100644 third_party/bigframes_vendored/pandas/io/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/io/common.py create mode 100644 third_party/bigframes_vendored/pandas/io/gbq.py create mode 100644 third_party/bigframes_vendored/pandas/io/parquet.py create mode 100644 third_party/bigframes_vendored/pandas/io/parsers/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/io/parsers/readers.py create mode 100644 third_party/bigframes_vendored/pandas/pandas/_typing.py create mode 100644 third_party/bigframes_vendored/sklearn/COPYING create mode 100644 third_party/bigframes_vendored/sklearn/base.py create mode 100644 third_party/bigframes_vendored/sklearn/cluster/_kmeans.py create mode 100644 third_party/bigframes_vendored/sklearn/compose/_column_transformer.py create mode 100644 third_party/bigframes_vendored/sklearn/decomposition/_pca.py create mode 100644 third_party/bigframes_vendored/sklearn/ensemble/_forest.py create mode 100644 third_party/bigframes_vendored/sklearn/linear_model/_base.py create mode 100644 third_party/bigframes_vendored/sklearn/linear_model/_logistic.py create mode 100644 third_party/bigframes_vendored/sklearn/metrics/_classification.py create mode 100644 third_party/bigframes_vendored/sklearn/metrics/_ranking.py create mode 100644 third_party/bigframes_vendored/sklearn/metrics/_regression.py create mode 100644 third_party/bigframes_vendored/sklearn/pipeline.py create mode 100644 third_party/bigframes_vendored/sklearn/preprocessing/_data.py create mode 100644 third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py create mode 100644 third_party/bigframes_vendored/xgboost/LICENSE create mode 100644 third_party/bigframes_vendored/xgboost/sklearn.py diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000000..7826eceab6a --- /dev/null +++ b/.flake8 @@ -0,0 +1,35 @@ +# -*- coding: utf-8 -*- +# +# 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. + +# Generated by synthtool. DO NOT EDIT! +[flake8] +ignore = E203, E231, E266, E501, W503 +exclude = + third_party/** + + # Exclude generated code. + **/proto/** + **/gapic/** + **/services/** + **/types/** + *_pb2.py + + # Standard linting exemptions. + **/.nox/** + __pycache__, + .git, + *.pyc, + conf.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000000..2cb117c23ac --- /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 +.coverage.* +.nox +.cache +.mypy_cache +.pytest_cache +.pytype + + +# 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 + +# Make sure a generated file isn't accidentally committed. +pylintrc +pylintrc.test + +# Notebook scratch +.ipynb_checkpoints diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 00000000000..56ce2981190 --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,7 @@ +[settings] +profile=black +force_sort_within_sections=True +lexicographical=True +single_line_exclusions=('typing',) +order_by_type=False +group_by_package=True diff --git a/.kokoro/build.sh b/.kokoro/build.sh new file mode 100644 index 00000000000..f80cf6eab97 --- /dev/null +++ b/.kokoro/build.sh @@ -0,0 +1,47 @@ +#!/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 + +if [[ -z "${PROJECT_ROOT:-}" ]]; then + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes" +fi + +cd "${PROJECT_ROOT}" + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=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 diff --git a/.kokoro/continuous/common.cfg b/.kokoro/continuous/common.cfg new file mode 100644 index 00000000000..5d40578ac79 --- /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: "bigframes/.kokoro/build.sh" diff --git a/.kokoro/continuous/continuous.cfg b/.kokoro/continuous/continuous.cfg new file mode 100644 index 00000000000..18a4c35325b --- /dev/null +++ b/.kokoro/continuous/continuous.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto diff --git a/.kokoro/continuous/e2e.cfg b/.kokoro/continuous/e2e.cfg new file mode 100644 index 00000000000..e2ca8bc78db --- /dev/null +++ b/.kokoro/continuous/e2e.cfg @@ -0,0 +1,7 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "system_prerelease system_noextras e2e notebook samples" +} diff --git a/.kokoro/continuous/nightly.cfg b/.kokoro/continuous/nightly.cfg new file mode 100644 index 00000000000..5d1ec3a8f75 --- /dev/null +++ b/.kokoro/continuous/nightly.cfg @@ -0,0 +1,8 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "NOX_SESSION" + value: "unit unit_prerelease system system_prerelease cover lint lint_setup_py mypy format docs e2e notebook" +} + +build_file: "bigframes/.kokoro/release-nightly.sh" diff --git a/.kokoro/presubmit/common.cfg b/.kokoro/presubmit/common.cfg new file mode 100644 index 00000000000..5d40578ac79 --- /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: "bigframes/.kokoro/build.sh" diff --git a/.kokoro/presubmit/e2e.cfg b/.kokoro/presubmit/e2e.cfg new file mode 100644 index 00000000000..e2ca8bc78db --- /dev/null +++ b/.kokoro/presubmit/e2e.cfg @@ -0,0 +1,7 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "system_prerelease system_noextras e2e notebook samples" +} diff --git a/.kokoro/presubmit/presubmit.cfg b/.kokoro/presubmit/presubmit.cfg new file mode 100644 index 00000000000..18a4c35325b --- /dev/null +++ b/.kokoro/presubmit/presubmit.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh new file mode 100755 index 00000000000..8742a0fc790 --- /dev/null +++ b/.kokoro/release-nightly.sh @@ -0,0 +1,213 @@ +#!/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 "${PROJECT_ROOT:-}" ]; then + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes" +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}" + +python3.10 -m pip install --require-hashes -r .kokoro/requirements.txt + +# 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] + +# If NOX_SESSION is set, it only runs the specified session, +# otherwise run all the sessions. +if ! [ ${DRY_RUN} ]; then + if [ -n "${NOX_SESSION:-}" ]; then + python3.10 -m nox -s ${NOX_SESSION:-} + else + python3.10 -m nox + fi +fi + +# Generate third party notices and include it in the licenses in setup.cfg +# TODO(shobs): Don't include it in the package once vertex colab can pick it +# from elsewhere +THIRD_PARTY_NOTICES_FILE=THIRD_PARTY_NOTICES +python3.10 -m pip install pip-licenses +python3.10 scripts/generate_third_party_notices.py --output-file ${THIRD_PARTY_NOTICES_FILE} +if ! [ -s ${THIRD_PARTY_NOTICES_FILE} ]; then + echo "${THIRD_PARTY_NOTICES_FILE} was generated with zero size" + exit -1 +fi +SETUP_CFG_BKP=`mktemp` +cp -f setup.cfg ${SETUP_CFG_BKP} +cat >> setup.cfg << EOF + +[metadata] +license_files = + LICENSE + ${THIRD_PARTY_NOTICES_FILE} +EOF + +# 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 + +# Make sure the wheel file has the third party notices included +# TODO(shobs): An utimate validation would be to create a virtual environment +# and install the wheel file, then verify that +# site-packages/bigframes-*.dist-info/ includes third party notices +python3.10 -c " +from zipfile import ZipFile +with ZipFile('$VERSION_WHEEL') as myzip: + third_party_licenses_info = [ + info + for info in myzip.infolist() + if info.filename.endswith('.dist-info/${THIRD_PARTY_NOTICES_FILE}') + ] + assert ( + len(third_party_licenses_info) == 1 + ), f'Found {len(third_party_licenses_info)} third party licenses' + assert ( + third_party_licenses_info[0].file_size > 0 + ), 'Package contains third party license of size 0' +" + +# 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 + gsutil cp -v dist/* ${gcs_path} + gsutil cp -v LICENSE ${gcs_path} + gsutil cp -v ${THIRD_PARTY_NOTICES_FILE} ${gcs_path} + gsutil -m cp -v "notebooks/00 - Summary.ipynb" \ + "notebooks/01 - Getting Started.ipynb" \ + "notebooks/02 - DataFrame.ipynb" \ + "notebooks/03 - Using ML - ML fundamentals.ipynb" \ + "notebooks/04 - Using ML - SKLearn linear regression.ipynb" \ + "notebooks/05 - Using ML - Easy linear regression.ipynb" \ + "notebooks/06 - Using ML - Large Language Models.ipynb" \ + "notebooks/50 - Remote Function.ipynb" \ + ${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 \ + --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 +mv -f ${SETUP_CFG_BKP} setup.cfg +rm -f ${THIRD_PARTY_NOTICES_FILE} + +# Keep this last so as not to block the release on PDF docs build. +pdf_docs () { + sudo apt update + sudo apt install -y texlive texlive-latex-extra latexmk + + pushd "${PROJECT_ROOT}/docs" + make latexpdf + + cp "_build/latex/bigframes.pdf" "_build/latex/bigframes-${RELEASE_VERSION}.pdf" + cp "_build/latex/bigframes.pdf" "_build/latex/bigframes-latest.pdf" + + if ! [ ${DRY_RUN} ]; then + for gcs_path in gs://vertex_sdk_private_releases/bigframe/ \ + gs://dl-platform-colab/bigframes/ \ + gs://bigframes-wheels/; + do + gsutil cp -v "_build/latex/bigframes-*.pdf" ${gcs_path} + done + fi + + popd +} + +pdf_docs + +# Copy html docs to GCS from where it can be deployed to anywhere else +gcs_docs () { + docs_gcs_bucket=gs://bigframes-docs + docs_local_html_folder=docs/_build/html + if [ ! -d ${docs_local_html_folder} ]; then + python3.10 -m nox -s docs + fi + + if ! [ ${DRY_RUN} ]; then + gsutil -m cp -v -r ${docs_local_html_folder} ${docs_gcs_bucket}/${GIT_HASH} + + # Copy the script to refresh firebase docs website from GCS to GCS itself + gsutil -m cp -v scripts/update_firebase_docs_site.sh ${docs_gcs_bucket} + fi +} + +gcs_docs diff --git a/.kokoro/requirements.in b/.kokoro/requirements.in new file mode 100644 index 00000000000..9fd9e127baf --- /dev/null +++ b/.kokoro/requirements.in @@ -0,0 +1,13 @@ +gcp-docuploader +gcp-releasetool>=1.10.5 # required for compatibility with cryptography>=39.x +importlib-metadata +typing-extensions +twine +wheel +setuptools +nox>=2022.11.21 # required to remove dependency on py +charset-normalizer<3 +click<8.1.0 +sphinx==4.5.0 # docs PDF see noxfile for HTML docs build +alabaster # docs PDF +recommonmark # docs PDF diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt new file mode 100644 index 00000000000..37a007667f1 --- /dev/null +++ b/.kokoro/requirements.txt @@ -0,0 +1,567 @@ +# +# This file is autogenerated by pip-compile with Python 3.10 +# by the following command: +# +# pip-compile --allow-unsafe --generate-hashes requirements.in +# +alabaster==0.7.13 \ + --hash=sha256:1ee19aca801bbabb5ba3f5f258e4422dfa86f82f3e9cefb0859b283cdd7f62a3 \ + --hash=sha256:a27a4a084d5e690e16e01e03ad2b2e552c61a65469419b907243193de1a84ae2 + # via + # -r requirements.in + # sphinx +argcomplete==3.1.1 \ + --hash=sha256:35fa893a88deea85ea7b20d241100e64516d6af6d7b0ae2bed1d263d26f70948 \ + --hash=sha256:6c4c563f14f01440aaffa3eae13441c5db2357b5eec639abe7c0b15334627dff + # via nox +attrs==23.1.0 \ + --hash=sha256:1f28b4522cdc2fb4256ac1a020c78acf9cba2c6b461ccd2c126f3aa8e8335d04 \ + --hash=sha256:6279836d581513a26f1bf235f9acd333bc9115683f14f7e8fae46c98fc50e015 + # via gcp-releasetool +babel==2.12.1 \ + --hash=sha256:b4246fb7677d3b98f501a39d43396d3cafdc8eadb045f4a31be01863f655c610 \ + --hash=sha256:cc2d99999cd01d44420ae725a21c9e3711b3aadc7976d6147f622d8581963455 + # via sphinx +bleach==6.0.0 \ + --hash=sha256:1a1a85c1595e07d8db14c5f09f09e6433502c51c595970edc090551f0db99414 \ + --hash=sha256:33c16e3353dbd13028ab4799a0f89a83f113405c766e9c122df8a06f5b85b3f4 + # via readme-renderer +cachetools==5.3.1 \ + --hash=sha256:95ef631eeaea14ba2e36f06437f36463aac3a096799e876ee55e5cdccb102590 \ + --hash=sha256:dce83f2d9b4e1f732a8cd44af8e8fab2dbe46201467fc98b3ef8f269092bf62b + # via google-auth +certifi==2023.5.7 \ + --hash=sha256:0f0d56dc5a6ad56fd4ba36484d6cc34451e1c6548c61daad8c320169f91eddc7 \ + --hash=sha256:c6c2e98f5c7869efca1f8916fed228dd91539f9f1b444c314c06eef02980c716 + # via requests +cffi==1.15.1 \ + --hash=sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5 \ + --hash=sha256:03425bdae262c76aad70202debd780501fabeaca237cdfddc008987c0e0f59ef \ + --hash=sha256:04ed324bda3cda42b9b695d51bb7d54b680b9719cfab04227cdd1e04e5de3104 \ + --hash=sha256:0e2642fe3142e4cc4af0799748233ad6da94c62a8bec3a6648bf8ee68b1c7426 \ + --hash=sha256:173379135477dc8cac4bc58f45db08ab45d228b3363adb7af79436135d028405 \ + --hash=sha256:198caafb44239b60e252492445da556afafc7d1e3ab7a1fb3f0584ef6d742375 \ + --hash=sha256:1e74c6b51a9ed6589199c787bf5f9875612ca4a8a0785fb2d4a84429badaf22a \ + --hash=sha256:2012c72d854c2d03e45d06ae57f40d78e5770d252f195b93f581acf3ba44496e \ + --hash=sha256:21157295583fe8943475029ed5abdcf71eb3911894724e360acff1d61c1d54bc \ + --hash=sha256:2470043b93ff09bf8fb1d46d1cb756ce6132c54826661a32d4e4d132e1977adf \ + --hash=sha256:285d29981935eb726a4399badae8f0ffdff4f5050eaa6d0cfc3f64b857b77185 \ + --hash=sha256:30d78fbc8ebf9c92c9b7823ee18eb92f2e6ef79b45ac84db507f52fbe3ec4497 \ + --hash=sha256:320dab6e7cb2eacdf0e658569d2575c4dad258c0fcc794f46215e1e39f90f2c3 \ + --hash=sha256:33ab79603146aace82c2427da5ca6e58f2b3f2fb5da893ceac0c42218a40be35 \ + --hash=sha256:3548db281cd7d2561c9ad9984681c95f7b0e38881201e157833a2342c30d5e8c \ + --hash=sha256:3799aecf2e17cf585d977b780ce79ff0dc9b78d799fc694221ce814c2c19db83 \ + --hash=sha256:39d39875251ca8f612b6f33e6b1195af86d1b3e60086068be9cc053aa4376e21 \ + --hash=sha256:3b926aa83d1edb5aa5b427b4053dc420ec295a08e40911296b9eb1b6170f6cca \ + --hash=sha256:3bcde07039e586f91b45c88f8583ea7cf7a0770df3a1649627bf598332cb6984 \ + --hash=sha256:3d08afd128ddaa624a48cf2b859afef385b720bb4b43df214f85616922e6a5ac \ + --hash=sha256:3eb6971dcff08619f8d91607cfc726518b6fa2a9eba42856be181c6d0d9515fd \ + --hash=sha256:40f4774f5a9d4f5e344f31a32b5096977b5d48560c5592e2f3d2c4374bd543ee \ + --hash=sha256:4289fc34b2f5316fbb762d75362931e351941fa95fa18789191b33fc4cf9504a \ + --hash=sha256:470c103ae716238bbe698d67ad020e1db9d9dba34fa5a899b5e21577e6d52ed2 \ + --hash=sha256:4f2c9f67e9821cad2e5f480bc8d83b8742896f1242dba247911072d4fa94c192 \ + --hash=sha256:50a74364d85fd319352182ef59c5c790484a336f6db772c1a9231f1c3ed0cbd7 \ + --hash=sha256:54a2db7b78338edd780e7ef7f9f6c442500fb0d41a5a4ea24fff1c929d5af585 \ + --hash=sha256:5635bd9cb9731e6d4a1132a498dd34f764034a8ce60cef4f5319c0541159392f \ + --hash=sha256:59c0b02d0a6c384d453fece7566d1c7e6b7bae4fc5874ef2ef46d56776d61c9e \ + --hash=sha256:5d598b938678ebf3c67377cdd45e09d431369c3b1a5b331058c338e201f12b27 \ + --hash=sha256:5df2768244d19ab7f60546d0c7c63ce1581f7af8b5de3eb3004b9b6fc8a9f84b \ + --hash=sha256:5ef34d190326c3b1f822a5b7a45f6c4535e2f47ed06fec77d3d799c450b2651e \ + --hash=sha256:6975a3fac6bc83c4a65c9f9fcab9e47019a11d3d2cf7f3c0d03431bf145a941e \ + --hash=sha256:6c9a799e985904922a4d207a94eae35c78ebae90e128f0c4e521ce339396be9d \ + --hash=sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c \ + --hash=sha256:7473e861101c9e72452f9bf8acb984947aa1661a7704553a9f6e4baa5ba64415 \ + --hash=sha256:8102eaf27e1e448db915d08afa8b41d6c7ca7a04b7d73af6514df10a3e74bd82 \ + --hash=sha256:87c450779d0914f2861b8526e035c5e6da0a3199d8f1add1a665e1cbc6fc6d02 \ + --hash=sha256:8b7ee99e510d7b66cdb6c593f21c043c248537a32e0bedf02e01e9553a172314 \ + --hash=sha256:91fc98adde3d7881af9b59ed0294046f3806221863722ba7d8d120c575314325 \ + --hash=sha256:94411f22c3985acaec6f83c6df553f2dbe17b698cc7f8ae751ff2237d96b9e3c \ + --hash=sha256:98d85c6a2bef81588d9227dde12db8a7f47f639f4a17c9ae08e773aa9c697bf3 \ + --hash=sha256:9ad5db27f9cabae298d151c85cf2bad1d359a1b9c686a275df03385758e2f914 \ + --hash=sha256:a0b71b1b8fbf2b96e41c4d990244165e2c9be83d54962a9a1d118fd8657d2045 \ + --hash=sha256:a0f100c8912c114ff53e1202d0078b425bee3649ae34d7b070e9697f93c5d52d \ + --hash=sha256:a591fe9e525846e4d154205572a029f653ada1a78b93697f3b5a8f1f2bc055b9 \ + --hash=sha256:a5c84c68147988265e60416b57fc83425a78058853509c1b0629c180094904a5 \ + --hash=sha256:a66d3508133af6e8548451b25058d5812812ec3798c886bf38ed24a98216fab2 \ + --hash=sha256:a8c4917bd7ad33e8eb21e9a5bbba979b49d9a97acb3a803092cbc1133e20343c \ + --hash=sha256:b3bbeb01c2b273cca1e1e0c5df57f12dce9a4dd331b4fa1635b8bec26350bde3 \ + --hash=sha256:cba9d6b9a7d64d4bd46167096fc9d2f835e25d7e4c121fb2ddfc6528fb0413b2 \ + --hash=sha256:cc4d65aeeaa04136a12677d3dd0b1c0c94dc43abac5860ab33cceb42b801c1e8 \ + --hash=sha256:ce4bcc037df4fc5e3d184794f27bdaab018943698f4ca31630bc7f84a7b69c6d \ + --hash=sha256:cec7d9412a9102bdc577382c3929b337320c4c4c4849f2c5cdd14d7368c5562d \ + --hash=sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9 \ + --hash=sha256:d61f4695e6c866a23a21acab0509af1cdfd2c013cf256bbf5b6b5e2695827162 \ + --hash=sha256:db0fbb9c62743ce59a9ff687eb5f4afbe77e5e8403d6697f7446e5f609976f76 \ + --hash=sha256:dd86c085fae2efd48ac91dd7ccffcfc0571387fe1193d33b6394db7ef31fe2a4 \ + --hash=sha256:e00b098126fd45523dd056d2efba6c5a63b71ffe9f2bbe1a4fe1716e1d0c331e \ + --hash=sha256:e229a521186c75c8ad9490854fd8bbdd9a0c9aa3a524326b55be83b54d4e0ad9 \ + --hash=sha256:e263d77ee3dd201c3a142934a086a4450861778baaeeb45db4591ef65550b0a6 \ + --hash=sha256:ed9cb427ba5504c1dc15ede7d516b84757c3e3d7868ccc85121d9310d27eed0b \ + --hash=sha256:fa6693661a4c91757f4412306191b6dc88c1703f780c8234035eac011922bc01 \ + --hash=sha256:fcd131dd944808b5bdb38e6f5b53013c5aa4f334c5cad0c72742f6eba4b73db0 + # via cryptography +charset-normalizer==2.1.1 \ + --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845 \ + --hash=sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f + # via + # -r requirements.in + # requests +click==8.0.4 \ + --hash=sha256:6a7a62563bbfabfda3a38f3023a1db4a35978c0abd76f6c9605ecd6554d6d9b1 \ + --hash=sha256:8458d7b1287c5fb128c90e23381cf99dcde74beaf6c7ff6384ce84d6fe090adb + # via + # -r requirements.in + # gcp-docuploader + # gcp-releasetool +colorlog==6.7.0 \ + --hash=sha256:0d33ca236784a1ba3ff9c532d4964126d8a2c44f1f0cb1d2b0728196f512f662 \ + --hash=sha256:bd94bd21c1e13fac7bd3153f4bc3a7dc0eb0974b8bc2fdf1a989e474f6e582e5 + # via + # gcp-docuploader + # nox +commonmark==0.9.1 \ + --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ + --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 + # via recommonmark +cryptography==41.0.1 \ + --hash=sha256:059e348f9a3c1950937e1b5d7ba1f8e968508ab181e75fc32b879452f08356db \ + --hash=sha256:1a5472d40c8f8e91ff7a3d8ac6dfa363d8e3138b961529c996f3e2df0c7a411a \ + --hash=sha256:1a8e6c2de6fbbcc5e14fd27fb24414507cb3333198ea9ab1258d916f00bc3039 \ + --hash=sha256:1fee5aacc7367487b4e22484d3c7e547992ed726d14864ee33c0176ae43b0d7c \ + --hash=sha256:5d092fdfedaec4cbbffbf98cddc915ba145313a6fdaab83c6e67f4e6c218e6f3 \ + --hash=sha256:5f0ff6e18d13a3de56f609dd1fd11470918f770c6bd5d00d632076c727d35485 \ + --hash=sha256:7bfc55a5eae8b86a287747053140ba221afc65eb06207bedf6e019b8934b477c \ + --hash=sha256:7fa01527046ca5facdf973eef2535a27fec4cb651e4daec4d043ef63f6ecd4ca \ + --hash=sha256:8dde71c4169ec5ccc1087bb7521d54251c016f126f922ab2dfe6649170a3b8c5 \ + --hash=sha256:8f4ab7021127a9b4323537300a2acfb450124b2def3756f64dc3a3d2160ee4b5 \ + --hash=sha256:948224d76c4b6457349d47c0c98657557f429b4e93057cf5a2f71d603e2fc3a3 \ + --hash=sha256:9a6c7a3c87d595608a39980ebaa04d5a37f94024c9f24eb7d10262b92f739ddb \ + --hash=sha256:b46e37db3cc267b4dea1f56da7346c9727e1209aa98487179ee8ebed09d21e43 \ + --hash=sha256:b4ceb5324b998ce2003bc17d519080b4ec8d5b7b70794cbd2836101406a9be31 \ + --hash=sha256:cb33ccf15e89f7ed89b235cff9d49e2e62c6c981a6061c9c8bb47ed7951190bc \ + --hash=sha256:d198820aba55660b4d74f7b5fd1f17db3aa5eb3e6893b0a41b75e84e4f9e0e4b \ + --hash=sha256:d34579085401d3f49762d2f7d6634d6b6c2ae1242202e860f4d26b046e3a1006 \ + --hash=sha256:eb8163f5e549a22888c18b0d53d6bb62a20510060a22fd5a995ec8a05268df8a \ + --hash=sha256:f73bff05db2a3e5974a6fd248af2566134d8981fd7ab012e5dd4ddb1d9a70699 + # via + # gcp-releasetool + # secretstorage +distlib==0.3.6 \ + --hash=sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46 \ + --hash=sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e + # via virtualenv +docutils==0.17.1 \ + --hash=sha256:686577d2e4c32380bb50cbb22f575ed742d58168cee37e99117a854bcd88f125 \ + --hash=sha256:cf316c8370a737a022b72b56874f6602acf974a37a9fba42ec2876387549fc61 + # via + # readme-renderer + # recommonmark + # sphinx +filelock==3.12.2 \ + --hash=sha256:002740518d8aa59a26b0c76e10fb8c6e15eae825d34b6fdf670333fd7b938d81 \ + --hash=sha256:cbb791cdea2a72f23da6ac5b5269ab0a0d161e9ef0100e653b69049a7706d1ec + # via virtualenv +gcp-docuploader==0.6.5 \ + --hash=sha256:30221d4ac3e5a2b9c69aa52fdbef68cc3f27d0e6d0d90e220fc024584b8d2318 \ + --hash=sha256:b7458ef93f605b9d46a4bf3a8dc1755dad1f31d030c8679edf304e343b347eea + # via -r requirements.in +gcp-releasetool==1.14.0 \ + --hash=sha256:a39708206838477bf42804a9dd7b549131ff3fc08c6abb0dd6939e274bd6dfd4 \ + --hash=sha256:aa163ec8322045c9d46e03f8f7500974568bd5d8d322b6a2d492068a5bdf57bf + # via -r requirements.in +google-api-core==2.11.1 \ + --hash=sha256:25d29e05a0058ed5f19c61c0a78b1b53adea4d9364b464d014fbda941f6d1c9a \ + --hash=sha256:d92a5a92dc36dd4f4b9ee4e55528a90e432b059f93aee6ad857f9de8cc7ae94a + # via + # google-cloud-core + # google-cloud-storage +google-auth==2.20.0 \ + --hash=sha256:030af34138909ccde0fbce611afc178f1d65d32fbff281f25738b1fe1c6f3eaa \ + --hash=sha256:23b7b0950fcda519bfb6692bf0d5289d2ea49fc143717cc7188458ec620e63fa + # via + # gcp-releasetool + # google-api-core + # google-cloud-core + # google-cloud-storage +google-cloud-core==2.3.2 \ + --hash=sha256:8417acf6466be2fa85123441696c4badda48db314c607cf1e5d543fa8bdc22fe \ + --hash=sha256:b9529ee7047fd8d4bf4a2182de619154240df17fbe60ead399078c1ae152af9a + # via google-cloud-storage +google-cloud-storage==2.9.0 \ + --hash=sha256:83a90447f23d5edd045e0037982c270302e3aeb45fc1288d2c2ca713d27bad94 \ + --hash=sha256:9b6ae7b509fc294bdacb84d0f3ea8e20e2c54a8b4bbe39c5707635fec214eff3 + # via gcp-docuploader +google-crc32c==1.5.0 \ + --hash=sha256:024894d9d3cfbc5943f8f230e23950cd4906b2fe004c72e29b209420a1e6b05a \ + --hash=sha256:02c65b9817512edc6a4ae7c7e987fea799d2e0ee40c53ec573a692bee24de876 \ + --hash=sha256:02ebb8bf46c13e36998aeaad1de9b48f4caf545e91d14041270d9dca767b780c \ + --hash=sha256:07eb3c611ce363c51a933bf6bd7f8e3878a51d124acfc89452a75120bc436289 \ + --hash=sha256:1034d91442ead5a95b5aaef90dbfaca8633b0247d1e41621d1e9f9db88c36298 \ + --hash=sha256:116a7c3c616dd14a3de8c64a965828b197e5f2d121fedd2f8c5585c547e87b02 \ + --hash=sha256:19e0a019d2c4dcc5e598cd4a4bc7b008546b0358bd322537c74ad47a5386884f \ + --hash=sha256:1c7abdac90433b09bad6c43a43af253e688c9cfc1c86d332aed13f9a7c7f65e2 \ + --hash=sha256:1e986b206dae4476f41bcec1faa057851f3889503a70e1bdb2378d406223994a \ + --hash=sha256:272d3892a1e1a2dbc39cc5cde96834c236d5327e2122d3aaa19f6614531bb6eb \ + --hash=sha256:278d2ed7c16cfc075c91378c4f47924c0625f5fc84b2d50d921b18b7975bd210 \ + --hash=sha256:2ad40e31093a4af319dadf503b2467ccdc8f67c72e4bcba97f8c10cb078207b5 \ + --hash=sha256:2e920d506ec85eb4ba50cd4228c2bec05642894d4c73c59b3a2fe20346bd00ee \ + --hash=sha256:3359fc442a743e870f4588fcf5dcbc1bf929df1fad8fb9905cd94e5edb02e84c \ + --hash=sha256:37933ec6e693e51a5b07505bd05de57eee12f3e8c32b07da7e73669398e6630a \ + --hash=sha256:398af5e3ba9cf768787eef45c803ff9614cc3e22a5b2f7d7ae116df8b11e3314 \ + --hash=sha256:3b747a674c20a67343cb61d43fdd9207ce5da6a99f629c6e2541aa0e89215bcd \ + --hash=sha256:461665ff58895f508e2866824a47bdee72497b091c730071f2b7575d5762ab65 \ + --hash=sha256:4c6fdd4fccbec90cc8a01fc00773fcd5fa28db683c116ee3cb35cd5da9ef6c37 \ + --hash=sha256:5829b792bf5822fd0a6f6eb34c5f81dd074f01d570ed7f36aa101d6fc7a0a6e4 \ + --hash=sha256:596d1f98fc70232fcb6590c439f43b350cb762fb5d61ce7b0e9db4539654cc13 \ + --hash=sha256:5ae44e10a8e3407dbe138984f21e536583f2bba1be9491239f942c2464ac0894 \ + --hash=sha256:635f5d4dd18758a1fbd1049a8e8d2fee4ffed124462d837d1a02a0e009c3ab31 \ + --hash=sha256:64e52e2b3970bd891309c113b54cf0e4384762c934d5ae56e283f9a0afcd953e \ + --hash=sha256:66741ef4ee08ea0b2cc3c86916ab66b6aef03768525627fd6a1b34968b4e3709 \ + --hash=sha256:67b741654b851abafb7bc625b6d1cdd520a379074e64b6a128e3b688c3c04740 \ + --hash=sha256:6ac08d24c1f16bd2bf5eca8eaf8304812f44af5cfe5062006ec676e7e1d50afc \ + --hash=sha256:6f998db4e71b645350b9ac28a2167e6632c239963ca9da411523bb439c5c514d \ + --hash=sha256:72218785ce41b9cfd2fc1d6a017dc1ff7acfc4c17d01053265c41a2c0cc39b8c \ + --hash=sha256:74dea7751d98034887dbd821b7aae3e1d36eda111d6ca36c206c44478035709c \ + --hash=sha256:759ce4851a4bb15ecabae28f4d2e18983c244eddd767f560165563bf9aefbc8d \ + --hash=sha256:77e2fd3057c9d78e225fa0a2160f96b64a824de17840351b26825b0848022906 \ + --hash=sha256:7c074fece789b5034b9b1404a1f8208fc2d4c6ce9decdd16e8220c5a793e6f61 \ + --hash=sha256:7c42c70cd1d362284289c6273adda4c6af8039a8ae12dc451dcd61cdabb8ab57 \ + --hash=sha256:7f57f14606cd1dd0f0de396e1e53824c371e9544a822648cd76c034d209b559c \ + --hash=sha256:83c681c526a3439b5cf94f7420471705bbf96262f49a6fe546a6db5f687a3d4a \ + --hash=sha256:8485b340a6a9e76c62a7dce3c98e5f102c9219f4cfbf896a00cf48caf078d438 \ + --hash=sha256:84e6e8cd997930fc66d5bb4fde61e2b62ba19d62b7abd7a69920406f9ecca946 \ + --hash=sha256:89284716bc6a5a415d4eaa11b1726d2d60a0cd12aadf5439828353662ede9dd7 \ + --hash=sha256:8b87e1a59c38f275c0e3676fc2ab6d59eccecfd460be267ac360cc31f7bcde96 \ + --hash=sha256:8f24ed114432de109aa9fd317278518a5af2d31ac2ea6b952b2f7782b43da091 \ + --hash=sha256:98cb4d057f285bd80d8778ebc4fde6b4d509ac3f331758fb1528b733215443ae \ + --hash=sha256:998679bf62b7fb599d2878aa3ed06b9ce688b8974893e7223c60db155f26bd8d \ + --hash=sha256:9ba053c5f50430a3fcfd36f75aff9caeba0440b2d076afdb79a318d6ca245f88 \ + --hash=sha256:9c99616c853bb585301df6de07ca2cadad344fd1ada6d62bb30aec05219c45d2 \ + --hash=sha256:a1fd716e7a01f8e717490fbe2e431d2905ab8aa598b9b12f8d10abebb36b04dd \ + --hash=sha256:a2355cba1f4ad8b6988a4ca3feed5bff33f6af2d7f134852cf279c2aebfde541 \ + --hash=sha256:b1f8133c9a275df5613a451e73f36c2aea4fe13c5c8997e22cf355ebd7bd0728 \ + --hash=sha256:b8667b48e7a7ef66afba2c81e1094ef526388d35b873966d8a9a447974ed9178 \ + --hash=sha256:ba1eb1843304b1e5537e1fca632fa894d6f6deca8d6389636ee5b4797affb968 \ + --hash=sha256:be82c3c8cfb15b30f36768797a640e800513793d6ae1724aaaafe5bf86f8f346 \ + --hash=sha256:c02ec1c5856179f171e032a31d6f8bf84e5a75c45c33b2e20a3de353b266ebd8 \ + --hash=sha256:c672d99a345849301784604bfeaeba4db0c7aae50b95be04dd651fd2a7310b93 \ + --hash=sha256:c6c777a480337ac14f38564ac88ae82d4cd238bf293f0a22295b66eb89ffced7 \ + --hash=sha256:cae0274952c079886567f3f4f685bcaf5708f0a23a5f5216fdab71f81a6c0273 \ + --hash=sha256:cd67cf24a553339d5062eff51013780a00d6f97a39ca062781d06b3a73b15462 \ + --hash=sha256:d3515f198eaa2f0ed49f8819d5732d70698c3fa37384146079b3799b97667a94 \ + --hash=sha256:d5280312b9af0976231f9e317c20e4a61cd2f9629b7bfea6a693d1878a264ebd \ + --hash=sha256:de06adc872bcd8c2a4e0dc51250e9e65ef2ca91be023b9d13ebd67c2ba552e1e \ + --hash=sha256:e1674e4307fa3024fc897ca774e9c7562c957af85df55efe2988ed9056dc4e57 \ + --hash=sha256:e2096eddb4e7c7bdae4bd69ad364e55e07b8316653234a56552d9c988bd2d61b \ + --hash=sha256:e560628513ed34759456a416bf86b54b2476c59144a9138165c9a1575801d0d9 \ + --hash=sha256:edfedb64740750e1a3b16152620220f51d58ff1b4abceb339ca92e934775c27a \ + --hash=sha256:f13cae8cc389a440def0c8c52057f37359014ccbc9dc1f0827936bcd367c6100 \ + --hash=sha256:f314013e7dcd5cf45ab1945d92e713eec788166262ae8deb2cfacd53def27325 \ + --hash=sha256:f583edb943cf2e09c60441b910d6a20b4d9d626c75a36c8fcac01a6c96c01183 \ + --hash=sha256:fd8536e902db7e365f49e7d9029283403974ccf29b13fc7028b97e2295b33556 \ + --hash=sha256:fe70e325aa68fa4b5edf7d1a4b6f691eb04bbccac0ace68e34820d283b5f80d4 + # via google-resumable-media +google-resumable-media==2.5.0 \ + --hash=sha256:218931e8e2b2a73a58eb354a288e03a0fd5fb1c4583261ac6e4c078666468c93 \ + --hash=sha256:da1bd943e2e114a56d85d6848497ebf9be6a14d3db23e9fc57581e7c3e8170ec + # via google-cloud-storage +googleapis-common-protos==1.59.1 \ + --hash=sha256:0cbedb6fb68f1c07e18eb4c48256320777707e7d0c55063ae56c15db3224a61e \ + --hash=sha256:b35d530fe825fb4227857bc47ad84c33c809ac96f312e13182bdeaa2abe1178a + # via google-api-core +idna==3.4 \ + --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \ + --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2 + # via requests +imagesize==1.4.1 \ + --hash=sha256:0d8d18d08f840c19d0ee7ca1fd82490fdc3729b7ac93f49870406ddde8ef8d8b \ + --hash=sha256:69150444affb9cb0d5cc5a92b3676f0b2fb7cd9ae39e947a5e11a36b4497cd4a + # via sphinx +importlib-metadata==6.7.0 \ + --hash=sha256:1aaf550d4f73e5d6783e7acb77aec43d49da8017410afae93822cc9cca98c4d4 \ + --hash=sha256:cb52082e659e97afc5dac71e79de97d8681de3aa07ff18578330904a9d18e5b5 + # via + # -r requirements.in + # keyring + # twine +jaraco-classes==3.2.3 \ + --hash=sha256:2353de3288bc6b82120752201c6b1c1a14b058267fa424ed5ce5984e3b922158 \ + --hash=sha256:89559fa5c1d3c34eff6f631ad80bb21f378dbcbb35dd161fd2c6b93f5be2f98a + # via keyring +jeepney==0.8.0 \ + --hash=sha256:5efe48d255973902f6badc3ce55e2aa6c5c3b3bc642059ef3a91247bcfcc5806 \ + --hash=sha256:c0a454ad016ca575060802ee4d590dd912e35c122fa04e70306de3d076cce755 + # via + # keyring + # secretstorage +jinja2==3.1.2 \ + --hash=sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852 \ + --hash=sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61 + # via + # gcp-releasetool + # sphinx +keyring==24.0.0 \ + --hash=sha256:4e87665a19c514c7edada8b15015cf89bd99b8d7edabc5c43cca77166fa8dfad \ + --hash=sha256:770f609eed2a16c65a6349f3ba1545d00c73f9fed4254c13766c674fe6d0d22b + # via + # gcp-releasetool + # twine +markdown-it-py==3.0.0 \ + --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ + --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb + # via rich +markupsafe==2.1.3 \ + --hash=sha256:05fb21170423db021895e1ea1e1f3ab3adb85d1c2333cbc2310f2a26bc77272e \ + --hash=sha256:0a4e4a1aff6c7ac4cd55792abf96c915634c2b97e3cc1c7129578aa68ebd754e \ + --hash=sha256:10bbfe99883db80bdbaff2dcf681dfc6533a614f700da1287707e8a5d78a8431 \ + --hash=sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686 \ + --hash=sha256:1577735524cdad32f9f694208aa75e422adba74f1baee7551620e43a3141f559 \ + --hash=sha256:1b40069d487e7edb2676d3fbdb2b0829ffa2cd63a2ec26c4938b2d34391b4ecc \ + --hash=sha256:282c2cb35b5b673bbcadb33a585408104df04f14b2d9b01d4c345a3b92861c2c \ + --hash=sha256:2c1b19b3aaacc6e57b7e25710ff571c24d6c3613a45e905b1fde04d691b98ee0 \ + --hash=sha256:2ef12179d3a291be237280175b542c07a36e7f60718296278d8593d21ca937d4 \ + --hash=sha256:338ae27d6b8745585f87218a3f23f1512dbf52c26c28e322dbe54bcede54ccb9 \ + --hash=sha256:3c0fae6c3be832a0a0473ac912810b2877c8cb9d76ca48de1ed31e1c68386575 \ + --hash=sha256:3fd4abcb888d15a94f32b75d8fd18ee162ca0c064f35b11134be77050296d6ba \ + --hash=sha256:42de32b22b6b804f42c5d98be4f7e5e977ecdd9ee9b660fda1a3edf03b11792d \ + --hash=sha256:504b320cd4b7eff6f968eddf81127112db685e81f7e36e75f9f84f0df46041c3 \ + --hash=sha256:525808b8019e36eb524b8c68acdd63a37e75714eac50e988180b169d64480a00 \ + --hash=sha256:56d9f2ecac662ca1611d183feb03a3fa4406469dafe241673d521dd5ae92a155 \ + --hash=sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac \ + --hash=sha256:65c1a9bcdadc6c28eecee2c119465aebff8f7a584dd719facdd9e825ec61ab52 \ + --hash=sha256:68e78619a61ecf91e76aa3e6e8e33fc4894a2bebe93410754bd28fce0a8a4f9f \ + --hash=sha256:69c0f17e9f5a7afdf2cc9fb2d1ce6aabdb3bafb7f38017c0b77862bcec2bbad8 \ + --hash=sha256:6b2b56950d93e41f33b4223ead100ea0fe11f8e6ee5f641eb753ce4b77a7042b \ + --hash=sha256:787003c0ddb00500e49a10f2844fac87aa6ce977b90b0feaaf9de23c22508b24 \ + --hash=sha256:7ef3cb2ebbf91e330e3bb937efada0edd9003683db6b57bb108c4001f37a02ea \ + --hash=sha256:8023faf4e01efadfa183e863fefde0046de576c6f14659e8782065bcece22198 \ + --hash=sha256:8758846a7e80910096950b67071243da3e5a20ed2546e6392603c096778d48e0 \ + --hash=sha256:8afafd99945ead6e075b973fefa56379c5b5c53fd8937dad92c662da5d8fd5ee \ + --hash=sha256:8c41976a29d078bb235fea9b2ecd3da465df42a562910f9022f1a03107bd02be \ + --hash=sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2 \ + --hash=sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707 \ + --hash=sha256:962f82a3086483f5e5f64dbad880d31038b698494799b097bc59c2edf392fce6 \ + --hash=sha256:9dcdfd0eaf283af041973bff14a2e143b8bd64e069f4c383416ecd79a81aab58 \ + --hash=sha256:aa7bd130efab1c280bed0f45501b7c8795f9fdbeb02e965371bbef3523627779 \ + --hash=sha256:ab4a0df41e7c16a1392727727e7998a467472d0ad65f3ad5e6e765015df08636 \ + --hash=sha256:ad9e82fb8f09ade1c3e1b996a6337afac2b8b9e365f926f5a61aacc71adc5b3c \ + --hash=sha256:af598ed32d6ae86f1b747b82783958b1a4ab8f617b06fe68795c7f026abbdcad \ + --hash=sha256:b076b6226fb84157e3f7c971a47ff3a679d837cf338547532ab866c57930dbee \ + --hash=sha256:b7ff0f54cb4ff66dd38bebd335a38e2c22c41a8ee45aa608efc890ac3e3931bc \ + --hash=sha256:bfce63a9e7834b12b87c64d6b155fdd9b3b96191b6bd334bf37db7ff1fe457f2 \ + --hash=sha256:c011a4149cfbcf9f03994ec2edffcb8b1dc2d2aede7ca243746df97a5d41ce48 \ + --hash=sha256:c9c804664ebe8f83a211cace637506669e7890fec1b4195b505c214e50dd4eb7 \ + --hash=sha256:ca379055a47383d02a5400cb0d110cef0a776fc644cda797db0c5696cfd7e18e \ + --hash=sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b \ + --hash=sha256:cd0f502fe016460680cd20aaa5a76d241d6f35a1c3350c474bac1273803893fa \ + --hash=sha256:ceb01949af7121f9fc39f7d27f91be8546f3fb112c608bc4029aef0bab86a2a5 \ + --hash=sha256:d080e0a5eb2529460b30190fcfcc4199bd7f827663f858a226a81bc27beaa97e \ + --hash=sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb \ + --hash=sha256:df0be2b576a7abbf737b1575f048c23fb1d769f267ec4358296f31c2479db8f9 \ + --hash=sha256:e09031c87a1e51556fdcb46e5bd4f59dfb743061cf93c4d6831bf894f125eb57 \ + --hash=sha256:e4dd52d80b8c83fdce44e12478ad2e85c64ea965e75d66dbeafb0a3e77308fcc \ + --hash=sha256:fec21693218efe39aa7f8599346e90c705afa52c5b31ae019b2e57e8f6542bb2 + # via jinja2 +mdurl==0.1.2 \ + --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \ + --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba + # via markdown-it-py +more-itertools==9.1.0 \ + --hash=sha256:cabaa341ad0389ea83c17a94566a53ae4c9d07349861ecb14dc6d0345cf9ac5d \ + --hash=sha256:d2bc7f02446e86a68911e58ded76d6561eea00cddfb2a91e7019bbb586c799f3 + # via jaraco-classes +nox==2023.4.22 \ + --hash=sha256:0b1adc619c58ab4fa57d6ab2e7823fe47a32e70202f287d78474adcc7bda1891 \ + --hash=sha256:46c0560b0dc609d7d967dc99e22cb463d3c4caf54a5fda735d6c11b5177e3a9f + # via -r requirements.in +packaging==23.1 \ + --hash=sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61 \ + --hash=sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f + # via + # gcp-releasetool + # nox + # sphinx +pkginfo==1.9.6 \ + --hash=sha256:4b7a555a6d5a22169fcc9cf7bfd78d296b0361adad412a346c1226849af5e546 \ + --hash=sha256:8fd5896e8718a4372f0ea9cc9d96f6417c9b986e23a4d116dda26b62cc29d046 + # via twine +platformdirs==3.7.0 \ + --hash=sha256:87fbf6473e87c078d536980ba970a472422e94f17b752cfad17024c18876d481 \ + --hash=sha256:cfd065ba43133ff103ab3bd10aecb095c2a0035fcd1f07217c9376900d94ba07 + # via virtualenv +protobuf==3.20.3 \ + --hash=sha256:03038ac1cfbc41aa21f6afcbcd357281d7521b4157926f30ebecc8d4ea59dcb7 \ + --hash=sha256:28545383d61f55b57cf4df63eebd9827754fd2dc25f80c5253f9184235db242c \ + --hash=sha256:2e3427429c9cffebf259491be0af70189607f365c2f41c7c3764af6f337105f2 \ + --hash=sha256:398a9e0c3eaceb34ec1aee71894ca3299605fa8e761544934378bbc6c97de23b \ + --hash=sha256:44246bab5dd4b7fbd3c0c80b6f16686808fab0e4aca819ade6e8d294a29c7050 \ + --hash=sha256:447d43819997825d4e71bf5769d869b968ce96848b6479397e29fc24c4a5dfe9 \ + --hash=sha256:67a3598f0a2dcbc58d02dd1928544e7d88f764b47d4a286202913f0b2801c2e7 \ + --hash=sha256:74480f79a023f90dc6e18febbf7b8bac7508420f2006fabd512013c0c238f454 \ + --hash=sha256:819559cafa1a373b7096a482b504ae8a857c89593cf3a25af743ac9ecbd23480 \ + --hash=sha256:899dc660cd599d7352d6f10d83c95df430a38b410c1b66b407a6b29265d66469 \ + --hash=sha256:8c0c984a1b8fef4086329ff8dd19ac77576b384079247c770f29cc8ce3afa06c \ + --hash=sha256:9aae4406ea63d825636cc11ffb34ad3379335803216ee3a856787bcf5ccc751e \ + --hash=sha256:a7ca6d488aa8ff7f329d4c545b2dbad8ac31464f1d8b1c87ad1346717731e4db \ + --hash=sha256:b6cc7ba72a8850621bfec987cb72623e703b7fe2b9127a161ce61e61558ad905 \ + --hash=sha256:bf01b5720be110540be4286e791db73f84a2b721072a3711efff6c324cdf074b \ + --hash=sha256:c02ce36ec760252242a33967d51c289fd0e1c0e6e5cc9397e2279177716add86 \ + --hash=sha256:d9e4432ff660d67d775c66ac42a67cf2453c27cb4d738fc22cb53b5d84c135d4 \ + --hash=sha256:daa564862dd0d39c00f8086f88700fdbe8bc717e993a21e90711acfed02f2402 \ + --hash=sha256:de78575669dddf6099a8a0f46a27e82a1783c557ccc38ee620ed8cc96d3be7d7 \ + --hash=sha256:e64857f395505ebf3d2569935506ae0dfc4a15cb80dc25261176c784662cdcc4 \ + --hash=sha256:f4bd856d702e5b0d96a00ec6b307b0f51c1982c2bf9c0052cf9019e9a544ba99 \ + --hash=sha256:f4c42102bc82a51108e449cbb32b19b180022941c727bac0cfd50170341f16ee + # via + # gcp-docuploader + # gcp-releasetool + # google-api-core +pyasn1==0.5.0 \ + --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \ + --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.3.0 \ + --hash=sha256:5bd01446b736eb9d31512a30d46c1ac3395d676c6f3cafa4c03eb54b9925631c \ + --hash=sha256:d3ccd6ed470d9ffbc716be08bd90efbd44d0734bc9303818f7336070984a162d + # via google-auth +pycparser==2.21 \ + --hash=sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9 \ + --hash=sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206 + # via cffi +pygments==2.15.1 \ + --hash=sha256:8ace4d3c1dd481894b2005f560ead0f9f19ee64fe983366be1a21e171d12775c \ + --hash=sha256:db2db3deb4b4179f399a09054b023b6a586b76499d36965813c71aa8ed7b5fd1 + # via + # readme-renderer + # rich + # sphinx +pyjwt==2.7.0 \ + --hash=sha256:ba2b425b15ad5ef12f200dc67dd56af4e26de2331f965c5439994dad075876e1 \ + --hash=sha256:bd6ca4a3c4285c1a2d4349e5a035fdf8fb94e04ccd0fcbe6ba289dae9cc3e074 + # via gcp-releasetool +pyperclip==1.8.2 \ + --hash=sha256:105254a8b04934f0bc84e9c24eb360a591aaf6535c9def5f29d92af107a9bf57 + # via gcp-releasetool +python-dateutil==2.8.2 \ + --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \ + --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9 + # via gcp-releasetool +readme-renderer==40.0 \ + --hash=sha256:9f77b519d96d03d7d7dce44977ba543090a14397c4f60de5b6eb5b8048110aa4 \ + --hash=sha256:e18feb2a1e7706f2865b81ebb460056d93fb29d69daa10b223c00faa7bd9a00a + # via twine +recommonmark==0.7.1 \ + --hash=sha256:1b1db69af0231efce3fa21b94ff627ea33dee7079a01dd0a7f8482c3da148b3f \ + --hash=sha256:bdb4db649f2222dcd8d2d844f0006b958d627f732415d399791ee436a3686d67 + # via -r requirements.in +requests==2.31.0 \ + --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ + --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1 + # via + # gcp-releasetool + # google-api-core + # google-cloud-storage + # requests-toolbelt + # sphinx + # twine +requests-toolbelt==1.0.0 \ + --hash=sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6 \ + --hash=sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06 + # via twine +rfc3986==2.0.0 \ + --hash=sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd \ + --hash=sha256:97aacf9dbd4bfd829baad6e6309fa6573aaf1be3f6fa735c8ab05e46cecb261c + # via twine +rich==13.4.2 \ + --hash=sha256:8f87bc7ee54675732fa66a05ebfe489e27264caeeff3728c945d25971b6485ec \ + --hash=sha256:d653d6bccede5844304c605d5aac802c7cf9621efd700b46c7ec2b51ea914898 + # via twine +rsa==4.9 \ + --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ + --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21 + # via google-auth +secretstorage==3.3.3 \ + --hash=sha256:2403533ef369eca6d2ba81718576c5e0f564d5cca1b58f73a8b23e7d4eeebd77 \ + --hash=sha256:f356e6628222568e3af06f2eba8df495efa13b3b63081dafd4f7d9a7b7bc9f99 + # via keyring +six==1.16.0 \ + --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ + --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 + # via + # bleach + # gcp-docuploader + # google-auth + # python-dateutil +snowballstemmer==2.2.0 \ + --hash=sha256:09b16deb8547d3412ad7b590689584cd0fe25ec8db3be37788be3810cbf19cb1 \ + --hash=sha256:c8e1716e83cc398ae16824e5572ae04e0d9fc2c6b985fb0f900f5f0c96ecba1a + # via sphinx +sphinx==4.5.0 \ + --hash=sha256:7bf8ca9637a4ee15af412d1a1d9689fec70523a68ca9bb9127c2f3eeb344e2e6 \ + --hash=sha256:ebf612653238bcc8f4359627a9b7ce44ede6fdd75d9d30f68255c7383d3a6226 + # via + # -r requirements.in + # recommonmark +sphinxcontrib-applehelp==1.0.4 \ + --hash=sha256:29d341f67fb0f6f586b23ad80e072c8e6ad0b48417db2bde114a4c9746feb228 \ + --hash=sha256:828f867945bbe39817c210a1abfd1bc4895c8b73fcaade56d45357a348a07d7e + # via sphinx +sphinxcontrib-devhelp==1.0.2 \ + --hash=sha256:8165223f9a335cc1af7ffe1ed31d2871f325254c0423bc0c4c7cd1c1e4734a2e \ + --hash=sha256:ff7f1afa7b9642e7060379360a67e9c41e8f3121f2ce9164266f61b9f4b338e4 + # via sphinx +sphinxcontrib-htmlhelp==2.0.1 \ + --hash=sha256:0cbdd302815330058422b98a113195c9249825d681e18f11e8b1f78a2f11efff \ + --hash=sha256:c38cb46dccf316c79de6e5515e1770414b797162b23cd3d06e67020e1d2a6903 + # via sphinx +sphinxcontrib-jsmath==1.0.1 \ + --hash=sha256:2ec2eaebfb78f3f2078e73666b1415417a116cc848b72e5172e596c871103178 \ + --hash=sha256:a9925e4a4587247ed2191a22df5f6970656cb8ca2bd6284309578f2153e0c4b8 + # via sphinx +sphinxcontrib-qthelp==1.0.3 \ + --hash=sha256:4c33767ee058b70dba89a6fc5c1892c0d57a54be67ddd3e7875a18d14cba5a72 \ + --hash=sha256:bd9fc24bcb748a8d51fd4ecaade681350aa63009a347a8c14e637895444dfab6 + # via sphinx +sphinxcontrib-serializinghtml==1.1.5 \ + --hash=sha256:352a9a00ae864471d3a7ead8d7d79f5fc0b57e8b3f95e9867eb9eb28999b92fd \ + --hash=sha256:aa5f6de5dfdf809ef505c4895e51ef5c9eac17d0f287933eb49ec495280b6952 + # via sphinx +twine==4.0.2 \ + --hash=sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8 \ + --hash=sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8 + # via -r requirements.in +typing-extensions==4.6.3 \ + --hash=sha256:88a4153d8505aabbb4e13aacb7c486c2b4a33ca3b3f807914a9b4c844c471c26 \ + --hash=sha256:d91d5919357fe7f681a9f2b5b4cb2a5f1ef0a1e9f59c4d8ff0d3491e05c0ffd5 + # via -r requirements.in +urllib3==1.26.16 \ + --hash=sha256:8d36afa7616d8ab714608411b4a3b13e58f463aee519024578e062e141dce20f \ + --hash=sha256:8f135f6502756bde6b2a9b28989df5fbe87c9970cecaa69041edcce7f0589b14 + # via + # google-auth + # requests + # twine +virtualenv==20.23.1 \ + --hash=sha256:34da10f14fea9be20e0fd7f04aba9732f84e593dac291b757ce42e3368a39419 \ + --hash=sha256:8ff19a38c1021c742148edc4f81cb43d7f8c6816d2ede2ab72af5b84c749ade1 + # via nox +webencodings==0.5.1 \ + --hash=sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78 \ + --hash=sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923 + # via bleach +wheel==0.40.0 \ + --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \ + --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247 + # via -r requirements.in +zipp==3.15.0 \ + --hash=sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b \ + --hash=sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556 + # via importlib-metadata + +# The following packages are considered to be unsafe in a requirements file: +setuptools==68.0.0 \ + --hash=sha256:11e52c67415a381d10d6b462ced9cfb97066179f0e871399e006c4ab101fc85f \ + --hash=sha256:baf1fdb41c6da4cd2eae722e135500da913332ab3f2f5c7d33af9b492acb5235 + # via -r requirements.in diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000000..07d7e6b7445 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,41 @@ +# Copyright 2021 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 +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.0.1 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - 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: 22.3.0 + hooks: + - id: black +- repo: https://github.com/pycqa/flake8 + rev: 3.9.2 + hooks: + - id: flake8 +- repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.1.1 + hooks: + - id: mypy + additional_dependencies: [types-requests] diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000000..39353a3f1ba --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,22 @@ +# Changelog + +[PyPI History][1] + +[1]: https://pypi.org/project/bigframes/#history + +## 0.1.0 (TBD) + +### Features + +* Add `bigframes.pandas` package with an API compatible with + [pandas](https://pandas.pydata.org/). Supported data sources include: + BigQuery SQL queries, BigQuery tables, CSV (local and GCS), Parquet (local + and GCS), and more. +* Add `bigframes.ml` package with an API inspired by + [scikit-learn](https://scikit-learn.org/stable/). Train machine learning + models and run batch predicition, powered by [BigQuery + ML](https://cloud.google.com/bigquery/docs/bqml-introduction). + +## [0.0.0](https://pypi.org/project/bigframes/0.0.0/) (2023-02-22) + +* Empty package to reserve package name. diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 00000000000..51d6abc4d3d --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,94 @@ +# 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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000000..b16bd944285 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,33 @@ +# How to contribute + +We'd love to accept your patches and contributions to this project. + +## Before you begin + +### Sign our Contributor License Agreement + +Contributions to this project must be accompanied by a +[Contributor License Agreement](https://cla.developers.google.com/about) (CLA). +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. + +If you or your current employer have already signed the Google CLA (even if it +was for a different project), you probably don't need to do it again. + +Visit to see your current agreements or to +sign a new one. + +### Review our community guidelines + +This project follows +[Google's Open Source Community Guidelines](https://opensource.google/conduct/). + +## Contribution process + +### 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. diff --git a/LICENSE b/LICENSE index 7a4a3ea2424..d6456956733 100644 --- a/LICENSE +++ b/LICENSE @@ -199,4 +199,4 @@ 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. \ No newline at end of file + limitations under the License. diff --git a/OWNERS b/OWNERS new file mode 100644 index 00000000000..f2a0b4383d1 --- /dev/null +++ b/OWNERS @@ -0,0 +1,10 @@ +bmil@google.com +chelsealin@google.com +garrettwu@google.com +henryjsolberg@google.com +hormati@google.com +huanc@google.com +jiaxun@google.com +shobs@google.com +swast@google.com +tbergeron@google.com diff --git a/README.md b/README.md deleted file mode 100644 index 3b3fd5761d9..00000000000 --- a/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# BigFrames - -BigFrames supports scalable DataFrame APIs on top of BigQuery. diff --git a/README.rst b/README.rst new file mode 100644 index 00000000000..5f473e11890 --- /dev/null +++ b/README.rst @@ -0,0 +1,8 @@ +BigQuery DataFrames +=================== + +BigQuery DataFrames provides a Pythonic DataFrame and machine learning (ML) API +powered by the BigQuery engine. + +* ``bigframes.pandas`` provides a pandas-like API for analytics. +* ``bigframes.ml`` provides a Scikit-Learn-like API for ML. diff --git a/bigframes/__init__.py b/bigframes/__init__.py new file mode 100644 index 00000000000..2ee745bc52e --- /dev/null +++ b/bigframes/__init__.py @@ -0,0 +1,28 @@ +# 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. + +"""BigQuery DataFrames provides a DataFrame API scaled by the BigQuery engine.""" + +from bigframes._config import options +from bigframes._config.bigquery_options import BigQueryOptions +from bigframes.session import connect, Session +from bigframes.version import __version__ + +__all__ = [ + "BigQueryOptions", + "connect", + "options", + "Session", + "__version__", +] diff --git a/bigframes/_config/__init__.py b/bigframes/_config/__init__.py new file mode 100644 index 00000000000..54c529fb0ca --- /dev/null +++ b/bigframes/_config/__init__.py @@ -0,0 +1,49 @@ +# 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. + +""" +Configuration for BigQuery DataFrames. Do not depend on other parts of BigQuery +DataFrames from this package. +""" + +import bigframes._config.bigquery_options as bigquery_options +import bigframes._config.display_options as display_options + + +class Options: + """Global options affecting BigQuery DataFrames behavior.""" + + def __init__(self): + self._bigquery_options = bigquery_options.BigQueryOptions() + self._display_options = display_options.DisplayOptions() + + @property + def bigquery(self) -> bigquery_options.BigQueryOptions: + """Options to use with the BigQuery engine.""" + return self._bigquery_options + + @property + def display(self) -> display_options.DisplayOptions: + """Options controlling object representation.""" + return self._display_options + + +options = Options() +"""Global options for default session.""" + + +__all__ = ( + "Options", + "options", +) diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py new file mode 100644 index 00000000000..5eb5ba7bc85 --- /dev/null +++ b/bigframes/_config/bigquery_options.py @@ -0,0 +1,114 @@ +# 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. + +"""Options for BigQuery DataFrames.""" + +from __future__ import annotations + +from typing import Optional + +import google.api_core.exceptions +import google.auth.credentials + +SESSION_STARTED_MESSAGE = "Cannot change '{attribute}' once a session has started." + + +class BigQueryOptions: + """Encapsulates configuration for working with an Session.""" + + def __init__( + self, + credentials: Optional[google.auth.credentials.Credentials] = None, + project: Optional[str] = None, + location: Optional[str] = None, + remote_udf_connection: Optional[str] = None, + use_regional_endpoints: bool = False, + ): + self._credentials = credentials + self._project = project + self._location = location + self._remote_udf_connection = remote_udf_connection + self._use_regional_endpoints = use_regional_endpoints + self._session_started = False + + @property + def credentials(self) -> Optional[google.auth.credentials.Credentials]: + """The OAuth2 Credentials to use for this client.""" + return self._credentials + + @credentials.setter + def credentials(self, value: Optional[google.auth.credentials.Credentials]): + if self._session_started and self._credentials is not value: + raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="credentials")) + self._credentials = value + + @property + def location(self) -> Optional[str]: + """Default location for jobs / datasets / tables. + + See: https://cloud.google.com/bigquery/docs/locations + """ + return self._location + + @location.setter + def location(self, value: Optional[str]): + if self._session_started and self._location != value: + raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="location")) + self._location = value + + @property + def project(self) -> Optional[str]: + """Google Cloud project ID to use for billing and default data project.""" + return self._project + + @project.setter + def project(self, value: Optional[str]): + if self._session_started and self._project != value: + raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="project")) + self._project = value + + @property + def remote_udf_connection(self) -> Optional[str]: + """Name of the BigQuery connection for the purpose of remote UDFs. + + It should be either pre created in `location`, or the user should have + privilege to create one. + """ + return self._remote_udf_connection + + @remote_udf_connection.setter + def remote_udf_connection(self, value: Optional[str]): + if self._session_started and self._remote_udf_connection != value: + raise ValueError( + SESSION_STARTED_MESSAGE.format(attribute="remote_udf_connection") + ) + self._remote_udf_connection = value + + @property + def use_regional_endpoints(self) -> bool: + """In preview. Flag to connect to regional API endpoints. + + Requires ``location`` to also be set. For example, set + ``location='asia-northeast1'`` and ``use_regional_endpoints=True`` to + connect to asia-northeast1-bigquery.googleapis.com. + """ + return self._use_regional_endpoints + + @use_regional_endpoints.setter + def use_regional_endpoints(self, value: bool): + if self._session_started and self._use_regional_endpoints != value: + raise ValueError( + SESSION_STARTED_MESSAGE.format(attribute="use_regional_endpoints") + ) + self._use_regional_endpoints = value diff --git a/bigframes/_config/display_options.py b/bigframes/_config/display_options.py new file mode 100644 index 00000000000..43faad0e785 --- /dev/null +++ b/bigframes/_config/display_options.py @@ -0,0 +1,55 @@ +# 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. + +"""Options for displaying objects.""" + +import contextlib +import dataclasses +from typing import Optional + +import pandas as pd + +import third_party.bigframes_vendored.pandas.core.config_init as vendored_pandas_config + + +@dataclasses.dataclass +class DisplayOptions: + __doc__ = vendored_pandas_config.display_options_doc + + max_columns: int = 20 + max_rows: int = 25 + progress_bar: Optional[str] = "auto" + + +@contextlib.contextmanager +def pandas_repr(display_options: DisplayOptions): + """Use this when visualizing with pandas. + + This context manager makes sure we reset the pandas options when we're done + so that we don't override pandas behavior. + """ + original_max_cols = pd.options.display.max_columns + original_max_rows = pd.options.display.max_rows + original_show_dimensions = pd.options.display.show_dimensions + + pd.options.display.max_columns = display_options.max_columns + pd.options.display.max_rows = display_options.max_rows + pd.options.display.show_dimensions = True # type: ignore + + try: + yield + finally: + pd.options.display.max_columns = original_max_cols + pd.options.display.max_rows = original_max_rows + pd.options.display.show_dimensions = original_show_dimensions diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py new file mode 100644 index 00000000000..70565fddf56 --- /dev/null +++ b/bigframes/core/__init__.py @@ -0,0 +1,1005 @@ +# 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. +from __future__ import annotations + +from dataclasses import dataclass +import functools +import math +import typing +from typing import Collection, Dict, Iterable, Literal, Optional, Sequence, Tuple + +from google.cloud import bigquery +import ibis +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types +import pandas + +import bigframes.core.guid +from bigframes.core.ordering import ( + ExpressionOrdering, + OrderingColumnReference, + stringify_order_id, +) +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops + +if typing.TYPE_CHECKING: + from bigframes.session import Session + +ORDER_ID_COLUMN = "bigframes_ordering_id" +PREDICATE_COLUMN = "bigframes_predicate" + + +@dataclass(frozen=True) +class WindowSpec: + """ + Specifies a window over which aggregate and analytic function may be applied. + grouping_keys: set of column ids to group on + preceding: Number of preceding rows in the window + following: Number of preceding rows in the window + ordering: List of columns ids and ordering direction to override base ordering + """ + + grouping_keys: typing.Sequence[str] = tuple() + ordering: typing.Sequence[OrderingColumnReference] = tuple() + preceding: typing.Optional[int] = None + following: typing.Optional[int] = None + min_periods: int = 0 + + +# TODO(swast): We might want to move this to it's own sub-module. +class ArrayValue: + """Immutable BigQuery DataFrames expression tree. + + Note: Usage of this class is considered to be private and subject to change + at any time. + + This class is a wrapper around Ibis expressions. Its purpose is to defer + Ibis projection operations to keep generated SQL small and correct when + mixing and matching columns from different versions of a DataFrame. + + Args: + session: + A BigQuery DataFrames session to allow more flexibility in running + queries. + table: An Ibis table expression. + columns: Ibis value expressions that can be projected as columns. + hidden_ordering_columns: Ibis value expressions to store ordering. + ordering: An ordering property of the data frame. + predicates: A list of filters on the data frame. + """ + + def __init__( + self, + session: Session, + table: ibis_types.Table, + columns: Optional[Sequence[ibis_types.Value]] = None, + hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None, + ordering: Optional[ExpressionOrdering] = None, + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + self._session = session + self._table = table + self._predicates = tuple(predicates) if predicates is not None else () + # TODO: Validate ordering + self._ordering = ordering or ExpressionOrdering() + # Allow creating a DataFrame directly from an Ibis table expression. + if columns is None: + self._columns = tuple( + table[key] + for key in table.columns + if ordering is None or key != ordering.ordering_id + ) + else: + # TODO(swast): Validate that each column references the same table (or + # no table for literal values). + self._columns = tuple(columns) + + # Meta columns store ordering, or other data that doesn't correspond to dataframe columns + self._hidden_ordering_columns = ( + tuple(hidden_ordering_columns) + if hidden_ordering_columns is not None + else () + ) + + # To allow for more efficient lookup by column name, create a + # dictionary mapping names to column values. + self._column_names = {column.get_name(): column for column in self._columns} + self._hidden_ordering_column_names = { + column.get_name(): column for column in self._hidden_ordering_columns + } + + @classmethod + def mem_expr_from_pandas( + cls, + pd_df: pandas.DataFrame, + session: Optional[Session], + ) -> ArrayValue: + """ + Builds an in-memory only (SQL only) expr from a pandas dataframe. + + Caution: If session is None, only a subset of expr functionality will be available (null Session is usually not supported). + """ + # must set non-null column labels. these are not the user-facing labels + pd_df = pd_df.set_axis( + [column or bigframes.core.guid.generate_guid() for column in pd_df.columns], + axis="columns", + ) + pd_df = pd_df.assign(**{ORDER_ID_COLUMN: range(len(pd_df))}) + # ibis memtable cannot handle NA, must convert to None + pd_df = pd_df.astype("object") # type: ignore + pd_df = pd_df.where(pandas.notnull(pd_df), None) + keys_memtable = ibis.memtable(pd_df) + return cls( + session, # type: ignore # Session cannot normally be none, see "caution" above + keys_memtable, + ordering=ExpressionOrdering( + ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN) + ), + hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), + ) + + @property + def table(self) -> ibis_types.Table: + return self._table + + @property + def predicates(self) -> typing.Tuple[ibis_types.BooleanValue, ...]: + return self._predicates + + @property + def reduced_predicate(self) -> typing.Optional[ibis_types.BooleanValue]: + """Returns the frame's predicates as an equivalent boolean value, useful where a single predicate value is preferred.""" + return ( + _reduce_predicate_list(self._predicates).name(PREDICATE_COLUMN) + if self._predicates + else None + ) + + @property + def columns(self) -> typing.Tuple[ibis_types.Value, ...]: + return self._columns + + @property + def column_names(self) -> Dict[str, ibis_types.Value]: + return self._column_names + + @property + def hidden_ordering_columns(self) -> typing.Tuple[ibis_types.Value, ...]: + return self._hidden_ordering_columns + + @property + def ordering(self) -> Sequence[ibis_types.Value]: + """Returns a sequence of ibis values which can be directly used to order a table expression. Has direction modifiers applied.""" + if not self._ordering: + return [] + else: + # TODO(swast): When we assign literals / scalars, we might not + # have a true Column. Do we need to check this before trying to + # sort by such a column? + return _convert_ordering_to_table_values( + {**self._column_names, **self._hidden_ordering_column_names}, + self._ordering.all_ordering_columns, + ) + + def builder(self) -> ArrayValueBuilder: + """Creates a mutable builder for expressions.""" + # Since ArrayValue is intended to be immutable (immutability offers + # potential opportunities for caching, though we might need to introduce + # more node types for that to be useful), we create a builder class. + return ArrayValueBuilder( + self._session, + self._table, + self._columns, + self._hidden_ordering_columns, + ordering=self._ordering, + predicates=self._predicates, + ) + + def insert_column(self, index: int, column: ibis_types.Value) -> ArrayValue: + expr = self.builder() + expr.columns.insert(index, column) + return expr.build() + + def drop_columns(self, columns: Iterable[str]) -> ArrayValue: + # Must generate offsets if we are dropping a column that ordering depends on + expr = self + for ordering_column in set(columns).intersection( + [col.column_id for col in self._ordering.ordering_value_columns] + ): + expr = self._hide_column(ordering_column) + + expr_builder = expr.builder() + remain_cols = [ + column for column in expr.columns if column.get_name() not in columns + ] + expr_builder.columns = remain_cols + return expr_builder.build() + + def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: + ibis_type = typing.cast(bigframes.dtypes.IbisDtype, self.get_column(key).type()) + return typing.cast( + bigframes.dtypes.Dtype, + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), + ) + + def get_column(self, key: str) -> ibis_types.Value: + """Gets the Ibis expression for a given column.""" + if key not in self._column_names.keys(): + raise ValueError( + "Column name {} not in set of values: {}".format( + key, self._column_names.keys() + ) + ) + return typing.cast(ibis_types.Value, self._column_names[key]) + + def get_any_column(self, key: str) -> ibis_types.Value: + """Gets the Ibis expression for a given column. Will also get hidden columns.""" + all_columns = {**self._column_names, **self._hidden_ordering_column_names} + if key not in all_columns.keys(): + raise ValueError( + "Column name {} not in set of values: {}".format( + key, all_columns.keys() + ) + ) + return typing.cast(ibis_types.Value, all_columns[key]) + + def _get_hidden_ordering_column(self, key: str) -> ibis_types.Column: + """Gets the Ibis expression for a given hidden column.""" + if key not in self._hidden_ordering_column_names.keys(): + raise ValueError( + "Column name {} not in set of values: {}".format( + key, self._hidden_ordering_column_names.keys() + ) + ) + return typing.cast(ibis_types.Column, self._hidden_ordering_column_names[key]) + + def apply_limit(self, max_results: int) -> ArrayValue: + table = self.to_ibis_expr( + ordering_mode="order_by", + expose_hidden_cols=True, + ).limit(max_results) + columns = [table[column_name] for column_name in self._column_names] + hidden_ordering_columns = [ + table[column_name] for column_name in self._hidden_ordering_column_names + ] + return ArrayValue( + self._session, + table, + columns=columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=self._ordering, + ) + + def filter(self, predicate: ibis_types.BooleanValue) -> ArrayValue: + """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" + expr = self.builder() + if expr.ordering: + expr.ordering = expr.ordering.with_is_sequential(False) + expr.predicates = [*self._predicates, predicate] + return expr.build() + + def order_by( + self, by: Sequence[OrderingColumnReference], stable: bool = False + ) -> ArrayValue: + expr_builder = self.builder() + expr_builder.ordering = self._ordering.with_ordering_columns(by, stable=stable) + return expr_builder.build() + + def reversed(self) -> ArrayValue: + expr_builder = self.builder() + expr_builder.ordering = self._ordering.with_reverse() + return expr_builder.build() + + @property + def offsets(self): + if not self._ordering.is_sequential: + raise ValueError( + "Expression does not have offsets. Generate them first using project_offsets." + ) + return self._get_hidden_ordering_column(self._ordering.ordering_id) + + def project_offsets(self) -> ArrayValue: + """Create a new expression that contains offsets. Should only be executed when offsets are needed for an operations. Has no effect on expression semantics.""" + if self._ordering.is_sequential: + return self + # TODO(tbergeron): Enforce total ordering + table = self.to_ibis_expr( + ordering_mode="offset_col", order_col_name=ORDER_ID_COLUMN + ) + columns = [table[column_name] for column_name in self._column_names] + ordering = ExpressionOrdering( + ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN), + is_sequential=True, + ) + return ArrayValue( + self._session, + table, + columns=columns, + hidden_ordering_columns=[table[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def _hide_column(self, column_id) -> ArrayValue: + """Pushes columns to hidden columns list. Used to hide ordering columns that have been dropped or destructively mutated.""" + expr_builder = self.builder() + # Need to rename column as caller might be creating a new row with the same name but different values. + # Can avoid this if don't allow callers to determine ids and instead generate unique ones in this class. + new_name = bigframes.core.guid.generate_guid(prefix="bigframes_hidden_") + expr_builder.hidden_ordering_columns = [ + *self._hidden_ordering_columns, + self.get_column(column_id).name(new_name), + ] + + ordering_columns = [ + col if col.column_id != column_id else col.with_name(new_name) + for col in self._ordering.ordering_value_columns + ] + + expr_builder.ordering = self._ordering.with_ordering_columns(ordering_columns) + return expr_builder.build() + + def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: + """ + Convenience function to promote copy of column offsets to a value column. Can be used to reset index. + """ + # Special case: offsets already exist + ordering = self._ordering + + if (not ordering.is_sequential) or (not ordering.ordering_id): + return self.project_offsets().promote_offsets() + col_id = bigframes.core.guid.generate_guid() + expr_builder = self.builder() + expr_builder.columns = [ + self._get_hidden_ordering_column(ordering.ordering_id).name(col_id), + *self.columns, + ] + return expr_builder.build(), col_id + + def select_columns(self, column_ids: typing.Sequence[str]): + return self.projection([self.get_column(col_id) for col_id in column_ids]) + + def projection(self, columns: Iterable[ibis_types.Value]) -> ArrayValue: + """Creates a new expression based on this expression with new columns.""" + # TODO(swast): We might want to do validation here that columns derive + # from the same table expression instead of (in addition to?) at + # construction time. + + expr = self + for ordering_column in set(self.column_names.keys()).intersection( + [col_ref.column_id for col_ref in self._ordering.ordering_value_columns] + ): + # Need to hide ordering columns that are being dropped. Alternatively, could project offsets + expr = expr._hide_column(ordering_column) + builder = expr.builder() + builder.columns = list(columns) + new_expr = builder.build() + return new_expr + + def shape(self) -> typing.Tuple[int, int]: + """Returns dimensions as (length, width) tuple.""" + width = len(self.columns) + count_expr = self.to_ibis_expr(ordering_mode="unordered").count() + sql = self._session.ibis_client.compile(count_expr) + row_iterator, _ = self._session._start_query( + sql=sql, + max_results=1, + ) + length = next(row_iterator)[0] + return (length, width) + + def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: + """Append together multiple ArrayValue objects.""" + if len(other) == 0: + return self + tables = [] + prefix_base = 10 + prefix_size = math.ceil(math.log(len(other) + 1, prefix_base)) + # Must normalize all ids to the same encoding size + max_encoding_size = max( + self._ordering.ordering_encoding_size, + *[expression._ordering.ordering_encoding_size for expression in other], + ) + for i, expr in enumerate([self, *other]): + ordering_prefix = str(i).zfill(prefix_size) + table = expr.to_ibis_expr( + ordering_mode="ordered_col", order_col_name=ORDER_ID_COLUMN + ) + # Rename the value columns based on horizontal offset before applying union. + table = table.select( + [ + table[col].name(f"column_{i}") + if col != ORDER_ID_COLUMN + else ( + ordering_prefix + + stringify_order_id(table[ORDER_ID_COLUMN], max_encoding_size) + ).name(ORDER_ID_COLUMN) + for i, col in enumerate(table.columns) + ] + ) + tables.append(table) + combined_table = ibis.union(*tables) + ordering = ExpressionOrdering( + ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN), + ordering_encoding_size=prefix_size + max_encoding_size, + ) + return ArrayValue( + self._session, + combined_table, + columns=[ + combined_table[col] + for col in combined_table.columns + if col != ORDER_ID_COLUMN + ], + hidden_ordering_columns=[combined_table[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def project_unary_op( + self, column_name: str, op: ops.UnaryOp, output_name=None + ) -> ArrayValue: + """Creates a new expression based on this expression with unary operation applied to one column.""" + value = op._as_ibis(self.get_column(column_name)).name( + output_name or column_name + ) + return self._set_or_replace_by_id(output_name or column_name, value) + + def project_binary_op( + self, + left_column_id: str, + right_column_id: str, + op: ops.BinaryOp, + output_column_id: str, + ) -> ArrayValue: + """Creates a new expression based on this expression with binary operation applied to two columns.""" + value = op( + self.get_column(left_column_id), self.get_column(right_column_id) + ).name(output_column_id) + return self._set_or_replace_by_id(output_column_id, value) + + def project_ternary_op( + self, + col_id_1: str, + col_id_2: str, + col_id_3: str, + op: ops.TernaryOp, + output_column_id: str, + ) -> ArrayValue: + """Creates a new expression based on this expression with ternary operation applied to three columns.""" + value = op( + self.get_column(col_id_1), + self.get_column(col_id_2), + self.get_column(col_id_3), + ).name(output_column_id) + return self._set_or_replace_by_id(output_column_id, value) + + def aggregate( + self, + aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp, str]], + by_column_ids: typing.Sequence[str] = (), + dropna: bool = True, + ) -> ArrayValue: + """ + Apply aggregations to the expression. + Arguments: + by_column_id: column id of the aggregation key, this is preserved through the transform + aggregations: input_column_id, operation, output_column_id tuples + dropna: whether null keys should be dropped + """ + table = self.to_ibis_expr() + stats = { + col_out: agg_op._as_ibis(table[col_in]) + for col_in, agg_op, col_out in aggregations + } + if by_column_ids: + result = table.group_by(by_column_ids).aggregate(**stats) + # Must have deterministic ordering, so order by the unique "by" column + ordering = ExpressionOrdering( + [ + OrderingColumnReference(column_id=column_id) + for column_id in by_column_ids + ] + ) + expr = ArrayValue(self._session, result, ordering=ordering) + if dropna: + for column_id in by_column_ids: + expr = expr.filter( + ops.notnull_op._as_ibis(expr.get_column(column_id)) + ) + # Can maybe remove this as Ordering id is redundant as by_column is unique after aggregation + return expr.project_offsets() + else: + aggregates = {**stats, ORDER_ID_COLUMN: ibis_types.literal(0)} + result = table.aggregate(**aggregates) + # Ordering is irrelevant for single-row output, but set ordering id regardless as other ops(join etc.) expect it. + ordering = ExpressionOrdering( + ordering_id_column=OrderingColumnReference(column_id=ORDER_ID_COLUMN), + is_sequential=True, + ) + return ArrayValue( + self._session, + result, + columns=[result[col_id] for col_id in [*stats.keys()]], + hidden_ordering_columns=[result[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def project_window_op( + self, + column_name: str, + op: agg_ops.WindowOp, + window_spec: WindowSpec, + output_name=None, + *, + skip_null_groups=False, + skip_reproject_unsafe: bool = False, + ) -> ArrayValue: + """ + Creates a new expression based on this expression with unary operation applied to one column. + column_name: the id of the input column present in the expression + op: the windowable operator to apply to the input column + window_spec: a specification of the window over which to apply the operator + output_name: the id to assign to the output of the operator, by default will replace input col if distinct output id not provided + skip_null_groups: will filter out any rows where any of the grouping keys is null + skip_reproject_unsafe: skips the reprojection step, can be used when performing many non-dependent window operations, user responsible for not nesting window expressions, or using outputs as join, filter or aggregation keys before a reprojection + """ + column = typing.cast(ibis_types.Column, self.get_column(column_name)) + window = self._ibis_window_from_spec(window_spec, allow_ties=op.handles_ties) + + window_op = op._as_ibis(column, window) + + clauses = [] + if op.skips_nulls: + clauses.append((column.isnull(), ibis.NA)) + if skip_null_groups: + for key in window_spec.grouping_keys: + clauses.append((self.get_column(key).isnull(), ibis.NA)) + if window_spec.min_periods: + clauses.append( + ( + agg_ops.count_op._as_ibis(column, window) + < ibis_types.literal(window_spec.min_periods), + ibis.NA, + ) + ) + + if clauses: + case_statement = ibis.case() + for clause in clauses: + case_statement = case_statement.when(clause[0], clause[1]) + case_statement = case_statement.else_(window_op).end() + window_op = case_statement + + result = self._set_or_replace_by_id(output_name or column_name, window_op) + # TODO(tbergeron): Automatically track analytic expression usage and defer reprojection until required for valid query generation. + return result._reproject_to_table() if not skip_reproject_unsafe else result + + def to_ibis_expr( + self, + ordering_mode: Literal[ + "order_by", "ordered_col", "offset_col", "unordered" + ] = "order_by", + order_col_name: Optional[str] = ORDER_ID_COLUMN, + expose_hidden_cols: bool = False, + ): + """ + Creates an Ibis table expression representing the DataFrame. + + ArrayValue objects are sorted, so the following options are available + to reflect this in the ibis expression. + + * "order_by" (Default): The output table will not have an ordering + column, however there will be an order_by clause applied to the ouput. + * "offset_col": Zero-based offsets are generated as a column, this will + not sort the rows however. + * "ordered_col": An ordered column is provided in output table, without + guarantee that the values are sequential + * "unordered": No ordering information will be provided in output. Only + value columns are projected. + + For offset or ordered column, order_col_name can be used to assign the + output label for the ordering column. If none is specified, the default + column name will be 'bigframes_ordering_id' + + Args: + ordering_mode: + How to construct the Ibis expression from the ArrayValue. See + above for details. + order_col_name: + If the ordering mode outputs a single ordering or offsets + column, use this as the column name. + expose_hidden_cols: + If True, include the hidden ordering columns in the results. + Only compatible with `order_by` and `unordered` + ``ordering_mode``. + Returns: + An ibis expression representing the data help by the ArrayValue object. + """ + assert ordering_mode in ( + "order_by", + "ordered_col", + "offset_col", + "unordered", + ) + if expose_hidden_cols and ordering_mode in ("ordered_col", "offset_col"): + raise ValueError( + f"Cannot expose hidden ordering columns with ordering_mode {ordering_mode}" + ) + + table = self._table + columns = list(self._columns) + hidden_ordering_columns = [ + col.column_id + for col in self._ordering.all_ordering_columns + if col.column_id not in self._column_names.keys() + ] + + if self.reduced_predicate is not None: + columns.append(self.reduced_predicate) + if ordering_mode in ("offset_col", "ordered_col"): + # Generate offsets if current ordering id semantics are not sufficiently strict + if (ordering_mode == "offset_col" and not self._ordering.is_sequential) or ( + ordering_mode == "ordered_col" and not self._ordering.order_id_defined + ): + window = ibis.window(order_by=self.ordering) + if self._predicates: + window = window.group_by(self.reduced_predicate) + columns.append(ibis.row_number().name(order_col_name).over(window)) + elif self._ordering.ordering_id: + columns.append( + self._get_hidden_ordering_column(self._ordering.ordering_id).name( + order_col_name + ) + ) + else: + # Should not be possible. + raise ValueError( + "Expression does not have ordering id and none was generated." + ) + elif ordering_mode == "order_by": + columns.extend( + [ + self._get_hidden_ordering_column(name) + for name in hidden_ordering_columns + ] + ) + + # We already need to add the hidden ordering columns for "order_by" so + # we can order by them. + if expose_hidden_cols and ordering_mode != "order_by": + columns.extend( + [ + self._get_hidden_ordering_column(name) + for name in hidden_ordering_columns + ] + ) + + # Special case for empty tables, since we can't create an empty + # projection. + if not columns: + return ibis.memtable([]) + table = table.select(columns) + # Make sure all dtypes are the "canonical" ones for BigFrames. This is + # important for operations like UNION where the schema must match. + table = bigframes.dtypes.ibis_table_to_canonical_types(table) + + if self.reduced_predicate is not None: + table = table.filter(table[PREDICATE_COLUMN]) + # Drop predicate as it is will be all TRUE after filtering + table = table.drop(PREDICATE_COLUMN) + if ordering_mode == "order_by": + # Some ordering columns are value columns, while other are used purely for ordering. + # We drop the non-value columns after the ordering + table = table.order_by( + _convert_ordering_to_table_values( + {col: table[col] for col in table.columns}, + self._ordering.all_ordering_columns, + ) # type: ignore + ) + # TODO(swast): We should be able to avoid this subquery by ordering + # by columns that don't have to be in the SELECT clause. + if not expose_hidden_cols: + table = table.drop(*hidden_ordering_columns) + + return table + + def start_query( + self, + job_config: Optional[bigquery.job.QueryJobConfig] = None, + max_results: Optional[int] = None, + ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: + """Execute a query and return metadata about the results.""" + # TODO(swast): Cache the job ID so we can look it up again if they ask + # for the results? We'd need a way to invalidate the cache if DataFrame + # becomes mutable, though. Or move this method to the immutable + # expression class. + # TODO(swast): We might want to move this method to Session and/or + # provide our own minimal metadata class. Tight coupling to the + # BigQuery client library isn't ideal, especially if we want to support + # a LocalSession for unit testing. + # TODO(swast): Add a timeout here? If the query is taking a long time, + # maybe we just print the job metadata that we have so far? + table = self.to_ibis_expr() + sql = self._session.ibis_client.compile(table) # type:ignore + return self._session._start_query( + sql=sql, + job_config=job_config, + max_results=max_results, + ) + + def _reproject_to_table(self) -> ArrayValue: + """ + Internal operators that projects the internal representation into a + new ibis table expression where each value column is a direct + reference to a column in that table expression. Needed after + some operations such as window operations that cannot be used + recursively in projections. + """ + table = self.to_ibis_expr( + ordering_mode="unordered", + order_col_name=self._ordering.ordering_id, + expose_hidden_cols=True, + ) + columns = [table[column_name] for column_name in self._column_names] + hidden_ordering_columns = [ + table[column_name] for column_name in self._hidden_ordering_column_names + ] + return ArrayValue( + self._session, + table, + columns=columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=self._ordering, + ) + + def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = False): + group_by: typing.List[ibis_types.Value] = ( + [ + typing.cast(ibis_types.Column, _as_identity(self.get_column(column))) + for column in window_spec.grouping_keys + ] + if window_spec.grouping_keys + else [] + ) + if self.reduced_predicate is not None: + group_by.append(self.reduced_predicate) + if window_spec.ordering: + order_by = _convert_ordering_to_table_values( + {**self._column_names, **self._hidden_ordering_column_names}, + window_spec.ordering, + ) + if not allow_ties: + # Most operator need an unambiguous ordering, so the table's total ordering is appended + order_by = tuple([*order_by, *self.ordering]) + elif (window_spec.following is not None) or (window_spec.preceding is not None): + # If window spec has following or preceding bounds, we need to apply an unambiguous ordering. + order_by = tuple(self.ordering) + else: + # Unbound grouping window. Suitable for aggregations but not for analytic function application. + order_by = None + return ibis.window( + preceding=window_spec.preceding, + following=window_spec.following, + order_by=order_by, + group_by=group_by, + ) + + def transpose_single_row( + self, + labels, + *, + index_col_id: str = "index", + value_col_id: str = "values", + dtype=pandas.Float64Dtype(), + ) -> ArrayValue: + """Pivot a single row into a 3 column expression with index, values and offsets. Only works if all values can be cast to a common type.""" + table = self.to_ibis_expr(ordering_mode="unordered") + sub_expressions = [] + for i, col_id in enumerate(self._column_names.keys()): + sub_expr = table.select( + ibis_types.literal(labels[i]).name(index_col_id), + ops.AsTypeOp(dtype)._as_ibis(table[col_id]).name(value_col_id), + ibis_types.literal(i).name(ORDER_ID_COLUMN), + ) + sub_expressions.append(sub_expr) + rotated_table = ibis.union(*sub_expressions) + return ArrayValue( + session=self._session, + table=rotated_table, + columns=[rotated_table[index_col_id], rotated_table[value_col_id]], + hidden_ordering_columns=[rotated_table[ORDER_ID_COLUMN]], + ordering=ExpressionOrdering( + ordering_id_column=OrderingColumnReference(column_id=ORDER_ID_COLUMN), + ), + ) + + # TODO(b/282041134) Remove deprecate_rename_column once label/id separation in dataframe + def deprecated_rename_column(self, old_id, new_id) -> ArrayValue: + """ + Don't use this, temporary measure until dataframe supports sqlid!=dataframe col id. + In future, caller shouldn't need to control internal column id strings. + """ + if new_id == old_id: + return self + return self._set_or_replace_by_id(new_id, self.get_column(old_id)).drop_columns( + [old_id] + ) + + def assign(self, source_id: str, destination_id: str) -> ArrayValue: + return self._set_or_replace_by_id(destination_id, self.get_column(source_id)) + + def assign_constant( + self, + destination_id: str, + value: typing.Any, + dtype: typing.Optional[bigframes.dtypes.Dtype], + ) -> ArrayValue: + # TODO(b/281587571): Solve scalar constant aggregation problem w/Ibis. + ibis_value = bigframes.dtypes.literal_to_ibis_scalar(value, dtype) + if ibis_value is None: + raise NotImplementedError( + f"Type not supported as scalar value {type(value)}" + ) + expr = self._set_or_replace_by_id(destination_id, ibis_value) + return expr._reproject_to_table() + + def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value): + builder = self.builder() + if id in self.column_names: + builder.columns = [ + val if (col_id != id) else new_value.name(id) + for col_id, val in self.column_names.items() + ] + else: + builder.columns = [*self.columns, new_value.name(id)] + return builder.build() + + def slice( + self, + start: typing.Optional[int] = None, + stop: typing.Optional[int] = None, + step: typing.Optional[int] = None, + ) -> ArrayValue: + if step == 0: + raise ValueError("slice step cannot be zero") + + if not step: + step = 1 + + # Special cases for head() and tail(), where we don't need to project + # offsets. LIMIT clause is much more efficient in BigQuery than a + # filter on row_number(). + if ( + (start is None or start == 0) + and step == 1 + and stop is not None + and stop > 0 + ): + return self.apply_limit(stop) + + if start is not None and start < 0 and step == 1 and stop is None: + return self.reversed().apply_limit(abs(start)).reversed() + + expr_with_offsets = self.project_offsets() + + # start with True and reduce with start, stop, and step conditions + cond_list = [expr_with_offsets.offsets == expr_with_offsets.offsets] + + last_offset = expr_with_offsets.offsets.max() + + # Convert negative indexes to positive indexes + if start and start < 0: + start = last_offset + start + 1 + if stop and stop < 0: + stop = last_offset + stop + 1 + + if start is not None: + if step >= 1: + cond_list.append(expr_with_offsets.offsets >= start) + else: + cond_list.append(expr_with_offsets.offsets <= start) + if stop is not None: + if step >= 1: + cond_list.append(expr_with_offsets.offsets < stop) + else: + cond_list.append(expr_with_offsets.offsets > stop) + if step > 1: + start = start if (start is not None) else 0 + cond_list.append((expr_with_offsets.offsets - start) % step == 0) + if step < 0: + start = start if (start is not None) else last_offset + cond_list.append((start - expr_with_offsets.offsets) % (-step) == 0) + + sliced_expr = expr_with_offsets.filter( + functools.reduce(lambda x, y: x & y, cond_list) + ) + return sliced_expr if step > 0 else sliced_expr.reversed() + + +class ArrayValueBuilder: + """Mutable expression class. + Use ArrayValue.builder() to create from a ArrayValue object. + """ + + def __init__( + self, + session: Session, + table: ibis_types.Table, + columns: Collection[ibis_types.Value] = (), + hidden_ordering_columns: Collection[ibis_types.Value] = (), + ordering: Optional[ExpressionOrdering] = None, + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + self.session = session + self.table = table + self.columns = list(columns) + self.hidden_ordering_columns = list(hidden_ordering_columns) + self.ordering = ordering + self.predicates = list(predicates) if predicates is not None else None + + def build(self) -> ArrayValue: + return ArrayValue( + session=self.session, + table=self.table, + columns=self.columns, + hidden_ordering_columns=self.hidden_ordering_columns, + ordering=self.ordering, + predicates=self.predicates, + ) + + +def _reduce_predicate_list( + predicate_list: typing.Collection[ibis_types.BooleanValue], +) -> ibis_types.BooleanValue: + """Converts a list of predicates BooleanValues into a single BooleanValue.""" + if len(predicate_list) == 0: + raise ValueError("Cannot reduce empty list of predicates") + if len(predicate_list) == 1: + (item,) = predicate_list + return item + return functools.reduce(lambda acc, pred: acc.__and__(pred), predicate_list) + + +def _convert_ordering_to_table_values( + value_lookup: typing.Mapping[str, ibis_types.Value], + ordering_columns: typing.Sequence[OrderingColumnReference], +) -> typing.Sequence[ibis_types.Value]: + column_refs = ordering_columns + ordering_values = [] + for ordering_col in column_refs: + column = typing.cast(ibis_types.Column, value_lookup[ordering_col.column_id]) + ordering_value = ( + ibis.asc(column) + if ordering_col.direction.is_ascending + else ibis.desc(column) + ) + # Bigquery SQL considers NULLS to be "smallest" values, but we need to override in these cases. + if (not ordering_col.na_last) and (not ordering_col.direction.is_ascending): + # Force nulls to be first + is_null_val = typing.cast(ibis_types.Column, column.isnull()) + ordering_values.append(ibis.desc(is_null_val)) + elif (ordering_col.na_last) and (ordering_col.direction.is_ascending): + # Force nulls to be last + is_null_val = typing.cast(ibis_types.Column, column.isnull()) + ordering_values.append(ibis.asc(is_null_val)) + ordering_values.append(ordering_value) + return ordering_values + + +def _as_identity(value: ibis_types.Value): + # Some types need to be converted to string to enable groupby + if value.type().is_float64() or value.type().is_geospatial(): + return value.cast(ibis_dtypes.str) + return value diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py new file mode 100644 index 00000000000..b13d7bf2d3a --- /dev/null +++ b/bigframes/core/block_transforms.py @@ -0,0 +1,199 @@ +# 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. +from __future__ import annotations + +import typing + +import pandas as pd + +import bigframes.core as core +import bigframes.core.blocks as blocks +import bigframes.core.ordering as ordering +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops + + +def indicate_duplicates( + block: blocks.Block, columns: typing.Sequence[str], keep: str = "first" +) -> typing.Tuple[blocks.Block, str]: + """Create a boolean column where True indicates a duplicate value""" + if keep not in ["first", "last", False]: + raise ValueError("keep must be one of 'first', 'last', or False'") + + if keep == "first": + # Count how many copies occur up to current copy of value + # Discard this value if there are copies BEFORE + window_spec = core.WindowSpec( + grouping_keys=tuple(columns), + following=0, + ) + elif keep == "last": + # Count how many copies occur up to current copy of values + # Discard this value if there are copies AFTER + window_spec = core.WindowSpec( + grouping_keys=tuple(columns), + preceding=0, + ) + else: # keep == False + # Count how many copies of the value occur in entire series. + # Discard this value if there are copies ANYWHERE + window_spec = core.WindowSpec(grouping_keys=tuple(columns)) + block, dummy = block.create_constant(1) + block, val_count_col_id = block.apply_window_op( + dummy, + agg_ops.count_op, + window_spec=window_spec, + ) + block, duplicate_indicator = block.apply_unary_op( + val_count_col_id, + ops.partial_right(ops.gt_op, 1), + ) + return ( + block.drop_columns( + ( + dummy, + val_count_col_id, + ) + ), + duplicate_indicator, + ) + + +def drop_duplicates( + block: blocks.Block, columns: typing.Sequence[str], keep: str = "first" +) -> blocks.Block: + block, dupe_indicator_id = indicate_duplicates(block, columns, keep) + block, keep_indicator_id = block.apply_unary_op(dupe_indicator_id, ops.invert_op) + return block.filter(keep_indicator_id).drop_columns( + (dupe_indicator_id, keep_indicator_id) + ) + + +def value_counts( + block: blocks.Block, + columns: typing.Sequence[str], + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + dropna: bool = True, +): + block, dummy = block.create_constant(1) + block, agg_ids = block.aggregate( + by_column_ids=columns, + aggregations=[(dummy, agg_ops.count_op)], + dropna=dropna, + as_index=True, + ) + count_id = agg_ids[0] + if normalize: + unbound_window = core.WindowSpec() + block, total_count_id = block.apply_window_op( + count_id, agg_ops.sum_op, unbound_window + ) + block, count_id = block.apply_binary_op(count_id, total_count_id, ops.div_op) + + if sort: + block = block.order_by( + [ + ordering.OrderingColumnReference( + count_id, + direction=ordering.OrderingDirection.ASC + if ascending + else ordering.OrderingDirection.DESC, + ) + ] + ) + return block.select_column(count_id).with_column_labels(["count"]) + + +def rank( + block: blocks.Block, + method: str = "average", + na_option: str = "keep", + ascending: bool = True, +): + if method not in ["average", "min", "max", "first", "dense"]: + raise ValueError( + "method must be one of 'average', 'min', 'max', 'first', or 'dense'" + ) + if na_option not in ["keep", "top", "bottom"]: + raise ValueError("na_option must be one of 'keep', 'top', or 'bottom'") + + columns = block.value_columns + labels = block.column_labels + # Step 1: Calculate row numbers for each row + # Identify null values to be treated according to na_option param + rownum_col_ids = [] + nullity_col_ids = [] + for col in columns: + block, nullity_col_id = block.apply_unary_op( + col, + ops.isnull_op, + ) + nullity_col_ids.append(nullity_col_id) + window = core.WindowSpec( + # BigQuery has syntax to reorder nulls with "NULLS FIRST/LAST", but that is unavailable through ibis presently, so must order on a separate nullity expression first. + ordering=( + ordering.OrderingColumnReference( + col, + ordering.OrderingDirection.ASC + if ascending + else ordering.OrderingDirection.DESC, + na_last=(na_option in ["bottom", "keep"]), + ), + ), + ) + # Count_op ignores nulls, so if na_option is "top" or "bottom", we instead count the nullity columns, where nulls have been mapped to bools + block, rownum_id = block.apply_window_op( + col if na_option == "keep" else nullity_col_id, + agg_ops.dense_rank_op if method == "dense" else agg_ops.count_op, + window_spec=window, + skip_reproject_unsafe=(col != columns[-1]), + ) + rownum_col_ids.append(rownum_id) + + # Step 2: Apply aggregate to groups of like input values. + # This step is skipped for method=='first' or 'dense' + if method in ["average", "min", "max"]: + agg_op = { + "average": agg_ops.mean_op, + "min": agg_ops.min_op, + "max": agg_ops.max_op, + }[method] + post_agg_rownum_col_ids = [] + for i in range(len(columns)): + block, result_id = block.apply_window_op( + rownum_col_ids[i], + agg_op, + window_spec=core.WindowSpec(grouping_keys=[columns[i]]), + skip_reproject_unsafe=(i < (len(columns) - 1)), + ) + post_agg_rownum_col_ids.append(result_id) + rownum_col_ids = post_agg_rownum_col_ids + + # Step 3: post processing: mask null values and cast to float + if method in ["min", "max", "first", "dense"]: + # Pandas rank always produces Float64, so must cast for aggregation types that produce ints + block = block.multi_apply_unary_op( + rownum_col_ids, ops.AsTypeOp(pd.Float64Dtype()) + ) + if na_option == "keep": + # For na_option "keep", null inputs must produce null outputs + for i in range(len(columns)): + block, null_const = block.create_constant(pd.NA, dtype=pd.Float64Dtype()) + block, rownum_col_ids[i] = block.apply_ternary_op( + null_const, nullity_col_ids[i], rownum_col_ids[i], ops.where_op + ) + + return block.select_columns(rownum_col_ids).with_column_labels(labels) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py new file mode 100644 index 00000000000..fe7261522db --- /dev/null +++ b/bigframes/core/blocks.py @@ -0,0 +1,917 @@ +# 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. + +"""Block is a 2D data structure that supports data mutability and views. + +These data structures are shared by DataFrame and Series. This allows views to +link in both directions (DataFrame to Series and vice versa) and prevents +circular dependencies. +""" + +from __future__ import annotations + +import functools +import itertools +import typing +from typing import Iterable, List, Optional, Sequence, Tuple + +import geopandas as gpd # type: ignore +import google.cloud.bigquery as bigquery +import ibis.expr.schema as ibis_schema +import ibis.expr.types as ibis_types +import numpy +import pandas as pd +import pyarrow as pa # type: ignore + +import bigframes.core as core +import bigframes.core.guid as guid +import bigframes.core.indexes as indexes +import bigframes.core.ordering as ordering +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops + +# Type constraint for wherever column labels are used +Label = typing.Optional[str] + + +class BlockHolder(typing.Protocol): + """Interface for mutable objects with state represented by a block value object.""" + + def _set_block(self, block: Block): + """Set the underlying block value of the object""" + + def _get_block(self) -> Block: + """Get the underlying block value of the object""" + + +class Block: + """A immutable 2D data structure.""" + + def __init__( + self, + expr: core.ArrayValue, + index_columns: Iterable[str] = (), + column_labels: Optional[Sequence[Label]] = None, + index_labels: Optional[Sequence[Label]] = None, + ): + """Construct a block object, will create default index if no index columns specified.""" + if index_labels and (len(index_labels) != len(list(index_columns))): + raise ValueError( + "'index_columns' and 'index_labels' must have equal length" + ) + if len(list(index_columns)) == 0: + expr, new_index_col_id = expr.promote_offsets() + index_columns = [new_index_col_id] + self._index_columns = tuple(index_columns) + self._index_labels = ( + tuple(index_labels) + if index_labels + else tuple([None for _ in index_columns]) + ) + self._expr = self._normalize_expression(expr, self._index_columns) + # TODO(tbergeron): Force callers to provide column labels + self._column_labels = ( + tuple(column_labels) if column_labels else tuple(self.value_columns) + ) + if len(self.value_columns) != len(self._column_labels): + raise ValueError( + f"'value_columns' (size {len(self.value_columns)}) and 'column_labels' (size {len(self._column_labels)}) must have equal length" + ) + + @property + def index(self) -> indexes.IndexValue: + """Row identities for values in the Block.""" + return indexes.IndexValue(self) + + @functools.cached_property + def shape(self) -> typing.Tuple[int, int]: + """Returns dimensions as (length, width) tuple.""" + impl_length, _ = self._expr.shape() + return (impl_length, len(self.value_columns)) + + @property + def index_columns(self) -> Sequence[str]: + """Column(s) to use as row labels.""" + return self._index_columns + + @property + def index_labels(self) -> Sequence[Label]: + """Name of column(s) to use as row labels.""" + return self._index_labels + + @property + def value_columns(self) -> Sequence[str]: + """All value columns, mutually exclusive with index columns.""" + return [ + column + for column in self._expr.column_names + if column not in self.index_columns + ] + + @property + def column_labels(self) -> List[Label]: + return list(self._column_labels) + + @property + def expr(self) -> core.ArrayValue: + """Expression representing all columns, including index columns.""" + return self._expr + + @property + def dtypes( + self, + ) -> Sequence[bigframes.dtypes.Dtype]: + """Returns the dtypes of the value columns.""" + return [self.expr.get_column_type(col) for col in self.value_columns] + + @property + def index_dtypes( + self, + ) -> Sequence[bigframes.dtypes.Dtype]: + """Returns the dtypes of the index columns.""" + return [self.expr.get_column_type(col) for col in self.index_columns] + + @functools.cached_property + def col_id_to_label(self) -> typing.Mapping[str, Label]: + """Get column label for value columns, or index name for index columns""" + return { + col_id: label + for col_id, label in zip(self.value_columns, self._column_labels) + } + + @functools.cached_property + def label_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]: + """Get column label for value columns, or index name for index columns""" + mapping: typing.Dict[Label, typing.Sequence[str]] = {} + for id, label in self.col_id_to_label.items(): + mapping[label] = (*mapping.get(label, ()), id) + return mapping + + @functools.cached_property + def col_id_to_index_name(self) -> typing.Mapping[str, Label]: + """Get column label for value columns, or index name for index columns""" + return { + col_id: label + for col_id, label in zip(self.index_columns, self._index_labels) + } + + @functools.cached_property + def index_name_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]: + """Get column label for value columns, or index name for index columns""" + mapping: typing.Dict[Label, typing.Sequence[str]] = {} + for id, label in self.col_id_to_index_name.items(): + mapping[label] = (*mapping.get(label, ()), id) + return mapping + + def order_by( + self, + by: typing.Sequence[ordering.OrderingColumnReference], + stable: bool = False, + ) -> Block: + return Block( + self._expr.order_by(by, stable=stable), + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + + def reversed(self) -> Block: + return Block( + self._expr.reversed(), + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + + def reset_index(self, drop: bool = True) -> Block: + """Reset the index of the block, promoting the old index to a value column. + + Arguments: + name: this is the column id for the new value id derived from the old index + + Returns: + A new Block because dropping index columns can break references + from Index classes that point to this block. + """ + block = self + expr, new_index_col_id = self._expr.promote_offsets() + if drop: + # Even though the index might be part of the ordering, keep that + # ordering expression as reset_index shouldn't change the row + # order. + expr = expr.drop_columns(self.index_columns) + block = Block( + expr, + index_columns=[new_index_col_id], + column_labels=self.column_labels, + index_labels=[None], + ) + else: + index_labels = self.index.names + index_labels_rewritten = [] + for level, label in enumerate(index_labels): + if label is None: + if "index" not in self.column_labels: + label = "index" + else: + label = f"level_{level}" + + if label in self.column_labels: + raise ValueError(f"cannot insert {label}, already exists") + index_labels_rewritten.append(label) + + block = Block( + expr, + index_columns=[new_index_col_id], + column_labels=[*index_labels_rewritten, *self.column_labels], + index_labels=[None], + ) + return block + + def set_index( + self, col_ids: typing.Sequence[str], drop: bool = True, append: bool = False + ) -> Block: + """Set the index of the block to + + Arguments: + ids: columns to be converted to index columns + drop: whether to drop the new index columns as value columns + append: whether to discard the existing index or add on to it + + Returns: + Block with new index + """ + expr = self._expr + + new_index_columns = [] + new_index_labels = [] + for col_id in col_ids: + col_copy_id = guid.generate_guid() + expr = expr.assign(col_id, col_copy_id) + new_index_columns.append(col_copy_id) + new_index_labels.append(self.col_id_to_label[col_id]) + + if append: + new_index_columns = [*self.index_columns, *new_index_columns] + new_index_labels = [*self._index_labels, *new_index_labels] + else: + expr = expr.drop_columns(self.index_columns) + + block = Block( + expr, + index_columns=new_index_columns, + column_labels=self.column_labels, + index_labels=new_index_labels, + ) + if drop: + # These are the value columns, new index uses the copies, so this is safe + block = block.drop_columns(col_ids) + return block + + def drop_levels(self, ids: typing.Sequence[str]): + for id in ids: + if id not in self.index_columns: + raise ValueError(f"{id} is not an index column") + expr = self._expr.drop_columns(ids) + remaining_index_col_ids = [ + col_id for col_id in self.index_columns if col_id not in ids + ] + if len(remaining_index_col_ids) == 0: + raise ValueError("Cannot drop all index levels, at least 1 must remain.") + level_names = [ + self.col_id_to_index_name[index_id] for index_id in remaining_index_col_ids + ] + return Block(expr, remaining_index_col_ids, self.column_labels, level_names) + + def reorder_levels(self, ids: typing.Sequence[str]): + if sorted(self.index_columns) != sorted(ids): + raise ValueError("Cannot drop or duplicate levels using reorder_levels.") + level_names = [self.col_id_to_index_name[index_id] for index_id in ids] + return Block(self.expr, ids, self.column_labels, level_names) + + def _to_dataframe(self, result, schema: ibis_schema.Schema) -> pd.DataFrame: + """Convert BigQuery data to pandas DataFrame with specific dtypes.""" + df = result.to_dataframe( + bool_dtype=pd.BooleanDtype(), + int_dtype=pd.Int64Dtype(), + float_dtype=pd.Float64Dtype(), + string_dtype=pd.StringDtype(storage="pyarrow"), + date_dtype=pd.ArrowDtype(pa.date32()), + datetime_dtype=pd.ArrowDtype(pa.timestamp("us")), + time_dtype=pd.ArrowDtype(pa.time64("us")), + timestamp_dtype=pd.ArrowDtype(pa.timestamp("us", tz="UTC")), + ) + + # Convert Geography column from StringDType to GeometryDtype. + for column_name, ibis_dtype in schema.items(): + if ibis_dtype.is_geospatial(): + df[column_name] = gpd.GeoSeries.from_wkt( + # https://github.com/geopandas/geopandas/issues/1879 + df[column_name].replace({numpy.nan: None}), + # BigQuery geography type is based on the WGS84 reference ellipsoid. + crs="EPSG:4326", + ) + return df + + def compute( + self, value_keys: Optional[Iterable[str]] = None, max_results=None + ) -> Tuple[pd.DataFrame, bigquery.QueryJob]: + """Run query and download results as a pandas DataFrame.""" + df, _, query_job = self._compute_and_count( + value_keys=value_keys, max_results=max_results + ) + return df, query_job + + def _compute_and_count( + self, value_keys: Optional[Iterable[str]] = None, max_results=None + ) -> Tuple[pd.DataFrame, int, bigquery.QueryJob]: + """Run query and download results as a pandas DataFrame. Return the total number of results as well.""" + # TODO(swast): Allow for dry run and timeout. + expr = self._expr + + value_column_names = value_keys or self.value_columns + if value_keys is not None: + index_columns = ( + expr.get_column(column_name) for column_name in self._index_columns + ) + value_columns = (expr.get_column(column_name) for column_name in value_keys) + expr = expr.projection(itertools.chain(index_columns, value_columns)) + + results_iterator, query_job = expr.start_query(max_results=max_results) + df = self._to_dataframe( + results_iterator, + expr.to_ibis_expr().schema(), + ) + + df = df.loc[:, [*self.index_columns, *value_column_names]] + if self.index_columns: + df = df.set_index(list(self.index_columns)) + df.index.names = self.index.names # type: ignore + + return df, results_iterator.total_rows, query_job + + def with_column_labels(self, value: typing.Iterable[Label]) -> Block: + label_list = tuple(value) + if len(label_list) != len(self.value_columns): + raise ValueError( + f"The column labels size `{len(label_list)} ` should equal to the value" + + f"columns size: {len(self.value_columns)}." + ) + return Block( + self._expr, + index_columns=self.index_columns, + column_labels=label_list, + index_labels=self.index.names, + ) + + def with_index_labels(self, value: typing.Sequence[Label]) -> Block: + if len(value) != len(self.index_columns): + raise ValueError( + f"The index labels size `{len(value)} ` should equal to the index" + + f"columns size: {len(self.value_columns)}." + ) + return Block( + self._expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=tuple(value), + ) + + def get_value_col_exprs( + self, column_names: Optional[Sequence[str]] = None + ) -> List[ibis_types.Value]: + """Retrive value column expressions.""" + column_names = self.value_columns if column_names is None else column_names + return [self._expr.get_column(column_name) for column_name in column_names] + + def apply_unary_op( + self, column: str, op: ops.UnaryOp, result_label: Label = None + ) -> typing.Tuple[Block, str]: + """ + Apply a unary op to the block. Creates a new column to store the result. + """ + # TODO(tbergeron): handle labels safely so callers don't need to + result_id = guid.generate_guid() + expr = self._expr.project_unary_op(column, op, result_id) + block = Block( + expr, + index_columns=self.index_columns, + column_labels=[*self.column_labels, result_label], + index_labels=self.index.names, + ) + return (block, result_id) + + def apply_binary_op( + self, + left_column_id: str, + right_column_id: str, + op: ops.BinaryOp, + result_label: Label = None, + ) -> typing.Tuple[Block, str]: + result_id = guid.generate_guid() + expr = self._expr.project_binary_op( + left_column_id, right_column_id, op, result_id + ) + block = Block( + expr, + index_columns=self.index_columns, + column_labels=[*self.column_labels, result_label], + index_labels=self.index.names, + ) + return (block, result_id) + + def apply_ternary_op( + self, + col_id_1: str, + col_id_2: str, + col_id_3: str, + op: ops.TernaryOp, + result_label: Label = None, + ) -> typing.Tuple[Block, str]: + result_id = guid.generate_guid() + expr = self._expr.project_ternary_op( + col_id_1, col_id_2, col_id_3, op, result_id + ) + block = Block( + expr, + index_columns=self.index_columns, + column_labels=[*self.column_labels, result_label], + index_labels=self.index.names, + ) + return (block, result_id) + + def multi_apply_window_op( + self, + columns: typing.Sequence[str], + op: agg_ops.WindowOp, + window_spec: core.WindowSpec, + *, + skip_null_groups: bool = False, + ) -> Block: + block = self + for i, col_id in enumerate(columns): + label = self.col_id_to_label[col_id] + block, result_id = block.apply_window_op( + col_id, + op, + window_spec=window_spec, + skip_reproject_unsafe=(i + 1) < len(columns), + result_label=label, + skip_null_groups=skip_null_groups, + ) + block = block.copy_values(result_id, col_id) + block = block.drop_columns([result_id]) + return block + + def multi_apply_unary_op( + self, + columns: typing.Sequence[str], + op: ops.UnaryOp, + ) -> Block: + block = self + for i, col_id in enumerate(columns): + label = self.col_id_to_label[col_id] + block, result_id = block.apply_unary_op( + col_id, + op, + result_label=label, + ) + block = block.copy_values(result_id, col_id) + block = block.drop_columns([result_id]) + return block + + def apply_window_op( + self, + column: str, + op: agg_ops.WindowOp, + window_spec: core.WindowSpec, + *, + result_label: Label = None, + skip_null_groups: bool = False, + skip_reproject_unsafe: bool = False, + ) -> typing.Tuple[Block, str]: + result_id = guid.generate_guid() + expr = self._expr.project_window_op( + column, + op, + window_spec, + result_id, + skip_null_groups=skip_null_groups, + skip_reproject_unsafe=skip_reproject_unsafe, + ) + block = Block( + expr, + index_columns=self.index_columns, + column_labels=[*self.column_labels, result_label], + index_labels=self._index_labels, + ) + return (block, result_id) + + def copy_values(self, source_column_id: str, destination_column_id: str) -> Block: + expr = self.expr.assign(source_column_id, destination_column_id) + return Block( + expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self._index_labels, + ) + + def create_constant( + self, + scalar_constant: typing.Any, + label: Label = None, + dtype: typing.Optional[bigframes.dtypes.Dtype] = None, + ) -> typing.Tuple[Block, str]: + result_id = guid.generate_guid() + expr = self.expr.assign_constant(result_id, scalar_constant, dtype=dtype) + labels = [*self.column_labels, label] + return ( + Block( + expr, + index_columns=self.index_columns, + column_labels=labels, + index_labels=self.index.names, + ), + result_id, + ) + + def assign_label(self, column_id: str, new_label: Label) -> Block: + col_index = self.value_columns.index(column_id) + new_labels = list(self.column_labels) + new_labels[col_index] = new_label + return self.with_column_labels(new_labels) + + def filter(self, column_name: str): + condition = typing.cast( + ibis_types.BooleanValue, self._expr.get_column(column_name) + ) + filtered_expr = self.expr.filter(condition) + return Block( + filtered_expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + + def aggregate_all_and_pivot( + self, + operation: agg_ops.AggregateOp, + *, + value_col_id: str = "values", + dropna: bool = True, + dtype=pd.Float64Dtype(), + ) -> Block: + aggregations = [(col_id, operation, col_id) for col_id in self.value_columns] + result_expr = self.expr.aggregate( + aggregations, dropna=dropna + ).transpose_single_row( + labels=self.column_labels, + index_col_id="index", + value_col_id=value_col_id, + dtype=dtype, + ) + return Block(result_expr, index_columns=["index"], column_labels=[None]) + + def select_column(self, id: str) -> Block: + return self.select_columns([id]) + + def select_columns(self, ids: typing.Sequence[str]) -> Block: + expr = self._expr.select_columns([*self.index_columns, *ids]) + col_labels = self._get_labels_for_columns(ids) + return Block(expr, self.index_columns, col_labels, self.index.names) + + def drop_columns(self, ids_to_drop: typing.Sequence[str]) -> Block: + """Drops columns by id. Can drop index""" + if set(ids_to_drop) & set(self.index_columns): + raise ValueError( + "Cannot directly drop index column. Use reset_index(drop=True)" + ) + expr = self._expr.drop_columns(ids_to_drop) + remaining_value_col_ids = [ + col_id for col_id in self.value_columns if (col_id not in ids_to_drop) + ] + labels = self._get_labels_for_columns(remaining_value_col_ids) + return Block(expr, self.index_columns, labels, self.index.names) + + def rename(self, *, columns: typing.Mapping[Label, Label]): + # TODO(tbergeron) Support function(Callable) as columns parameter. + col_labels = [ + (columns.get(col_label, col_label)) for col_label in self.column_labels + ] + return self.with_column_labels(col_labels) + + def aggregate( + self, + by_column_ids: typing.Sequence[str], + aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp]], + *, + as_index: bool = True, + dropna: bool = True, + ) -> typing.Tuple[Block, typing.Sequence[str]]: + """ + Apply aggregations to the block. Callers responsible for setting index column(s) after. + Arguments: + by_column_id: column id of the aggregation key, this is preserved through the transform and used as index + aggregations: input_column_id, operation tuples + as_index: if True, grouping keys will be index columns in result, otherwise they will be non-index columns. + dropna: whether null keys should be dropped + """ + agg_specs = [ + (input_id, operation, guid.generate_guid()) + for input_id, operation in aggregations + ] + output_col_ids = [agg_spec[2] for agg_spec in agg_specs] + result_expr = self.expr.aggregate(agg_specs, by_column_ids, dropna=dropna) + + aggregate_labels = self._get_labels_for_columns( + [agg[0] for agg in aggregations] + ) + if as_index: + # TODO: Generalize to multi-index + names: typing.List[Label] = [] + for by_col_id in by_column_ids: + if by_col_id in self.index_columns: + # Groupby level 0 case, keep index name + index_name = self.col_id_to_index_name[by_col_id] + else: + index_name = self.col_id_to_label[by_col_id] + names.append(index_name) + return ( + Block( + result_expr, + index_columns=by_column_ids, + column_labels=aggregate_labels, + index_labels=names, + ), + output_col_ids, + ) + else: + by_column_labels = self._get_labels_for_columns(by_column_ids) + labels = (*by_column_labels, *aggregate_labels) + return Block(result_expr, column_labels=labels), output_col_ids + + def _get_labels_for_columns(self, column_ids: typing.Sequence[str]): + """Get column label for value columns, or index name for index columns""" + lookup = self.col_id_to_label + return [lookup.get(col_id, None) for col_id in column_ids] + + def _normalize_expression( + self, + expr: core.ArrayValue, + index_columns: typing.Sequence[str], + assert_value_size: typing.Optional[int] = None, + ): + """Normalizes expression by moving index columns to left.""" + value_columns = [ + col_id for col_id in expr.column_names.keys() if col_id not in index_columns + ] + if (assert_value_size is not None) and ( + len(value_columns) != assert_value_size + ): + raise ValueError("Unexpected number of value columns.") + return expr.select_columns([*index_columns, *value_columns]) + + def slice( + self: bigframes.core.blocks.Block, + start: typing.Optional[int] = None, + stop: typing.Optional[int] = None, + step: typing.Optional[int] = None, + ) -> bigframes.core.blocks.Block: + sliced_expr = self.expr.slice(start=start, stop=stop, step=step) + # since this is slice, return a copy even if unchanged + block = Block( + sliced_expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self._index_labels, + ) + return block + + def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]: + expr, result_id = self._expr.promote_offsets() + return ( + Block( + expr, + index_columns=self.index_columns, + column_labels=[label, *self.column_labels], + index_labels=self._index_labels, + ), + result_id, + ) + + def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: + axis_number = _get_axis_number(axis) + if axis_number == 0: + expr = self._expr + for index_col in self._index_columns: + expr = expr.project_unary_op(index_col, ops.AsTypeOp("string")) + prefix_op = ops.BinopPartialLeft(ops.add_op, prefix) + expr = expr.project_unary_op(index_col, prefix_op) + return Block( + expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + if axis_number == 1: + expr = self._expr + return Block( + self._expr, + index_columns=self.index_columns, + column_labels=[f"{prefix}{label}" for label in self.column_labels], + index_labels=self.index.names, + ) + + def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block: + axis_number = _get_axis_number(axis) + if axis_number == 0: + expr = self._expr + for index_col in self._index_columns: + expr = expr.project_unary_op(index_col, ops.AsTypeOp("string")) + prefix_op = ops.BinopPartialRight(ops.add_op, suffix) + expr = expr.project_unary_op(index_col, prefix_op) + return Block( + expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + if axis_number == 1: + expr = self._expr + return Block( + self._expr, + index_columns=self.index_columns, + column_labels=[f"{label}{suffix}" for label in self.column_labels], + index_labels=self.index.names, + ) + + def concat( + self, + other: typing.Iterable[Block], + how: typing.Literal["inner", "outer"], + ignore_index=False, + ): + blocks: typing.List[Block] = [self, *other] + if ignore_index: + blocks = [block.reset_index() for block in blocks] + + result_labels = _align_indices(blocks) + + index_nlevels = blocks[0].index.nlevels + + aligned_schema = _align_schema(blocks, how=how) + aligned_blocks = [ + _align_block_to_schema(block, aligned_schema) for block in blocks + ] + result_expr = aligned_blocks[0]._expr.concat( + [block._expr for block in aligned_blocks[1:]] + ) + result_block = Block( + result_expr, + index_columns=list(result_expr.column_names.keys())[:index_nlevels], + column_labels=aligned_blocks[0].column_labels, + index_labels=result_labels, + ) + if ignore_index: + result_block = result_block.reset_index() + return result_block + + +def block_from_local(data, session=None, use_index=True) -> Block: + # TODO(tbergeron): Handle duplicate column labels + pd_data = pd.DataFrame(data) + + column_labels = list(pd_data.columns) + if not all((label is None) or isinstance(label, str) for label in column_labels): + raise NotImplementedError("Only string column labels supported") + + if use_index: + if pd_data.index.nlevels > 1: + raise NotImplementedError("multi-indices not supported.") + index_label = pd_data.index.name + if (index_label is not None) and (not isinstance(index_label, str)): + raise NotImplementedError("Only string index names supported") + + index_id = guid.generate_guid() + pd_data = pd_data.reset_index(names=index_id) + keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) + return Block( + keys_expr, + column_labels=column_labels, + index_columns=[index_id], + index_labels=[index_label], + ) + else: + keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) + # Constructor will create default range index + return Block(keys_expr, column_labels=column_labels) + + +def _align_block_to_schema( + block: Block, schema: dict[Label, bigframes.dtypes.Dtype] +) -> Block: + """For a given schema, remap block to schema by reordering columns and inserting nulls.""" + col_ids: typing.Tuple[str, ...] = () + for label, dtype in schema.items(): + # TODO: Support casting to lcd type - requires mixed type support + matching_ids: typing.Sequence[str] = block.label_to_col_id.get(label, ()) + if len(matching_ids) > 0: + col_id = matching_ids[-1] + col_ids = (*col_ids, col_id) + else: + block, null_column = block.create_constant(None, dtype=dtype) + col_ids = (*col_ids, null_column) + return block.select_columns(col_ids).with_column_labels( + [item for item in schema.keys()] + ) + + +def _align_schema( + blocks: typing.Iterable[Block], how: typing.Literal["inner", "outer"] +) -> typing.Dict[Label, bigframes.dtypes.Dtype]: + schemas = [_get_block_schema(block) for block in blocks] + reduction = _combine_schema_inner if how == "inner" else _combine_schema_outer + return functools.reduce(reduction, schemas) + + +def _align_indices(blocks: typing.Sequence[Block]) -> typing.Sequence[Label]: + """Validates that the blocks have compatible indices and returns the resulting label names.""" + names = blocks[0].index.names + types = blocks[0].index.dtypes + for block in blocks[1:]: + if len(names) != block.index.nlevels: + raise NotImplementedError( + "Cannot combine indices with different number of levels. Use 'ignore_index'=True." + ) + if block.index.dtypes != types: + raise NotImplementedError( + "Cannot combine different index dtypes. Use 'ignore_index'=True." + ) + names = [ + lname if lname == rname else None + for lname, rname in zip(names, block.index.names) + ] + return names + + +def _combine_schema_inner( + left: typing.Dict[Label, bigframes.dtypes.Dtype], + right: typing.Dict[Label, bigframes.dtypes.Dtype], +) -> typing.Dict[Label, bigframes.dtypes.Dtype]: + result = dict() + for label, type in left.items(): + if label in right: + if type != right[label]: + raise ValueError( + f"Cannot concat rows with label {label} due to mismatched types" + ) + result[label] = type + return result + + +def _combine_schema_outer( + left: typing.Dict[Label, bigframes.dtypes.Dtype], + right: typing.Dict[Label, bigframes.dtypes.Dtype], +) -> typing.Dict[Label, bigframes.dtypes.Dtype]: + result = dict() + for label, type in left.items(): + if (label in right) and (type != right[label]): + raise ValueError( + f"Cannot concat rows with label {label} due to mismatched types" + ) + result[label] = type + for label, type in right.items(): + if label not in left: + result[label] = type + return result + + +def _get_block_schema( + block: Block, +) -> typing.Dict[Label, bigframes.dtypes.Dtype]: + """Extracts the schema from the block. Where duplicate labels exist, take the last matching column.""" + result = dict() + for label, dtype in zip(block.column_labels, block.dtypes): + result[label] = typing.cast(bigframes.dtypes.Dtype, dtype) + return result + + +def _get_axis_number(axis: str | int | None) -> typing.Literal[0, 1]: + if axis in {0, "index", "rows", None}: + return 0 + elif axis in {1, "columns"}: + return 1 + else: + raise ValueError(f"Not a valid axis: {axis}") diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py new file mode 100644 index 00000000000..a8655116e93 --- /dev/null +++ b/bigframes/core/groupby/__init__.py @@ -0,0 +1,315 @@ +# 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. + +from __future__ import annotations + +import typing + +import bigframes.core as core +import bigframes.core.blocks as blocks +import bigframes.core.ordering as order +import bigframes.core.window as windows +import bigframes.dataframe as df +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops +import bigframes.series as series +import third_party.bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby + + +class DataFrameGroupBy(vendored_pandas_groupby.DataFrameGroupBy): + __doc__ = vendored_pandas_groupby.GroupBy.__doc__ + + def __init__( + self, + block: blocks.Block, + by_col_ids: typing.Sequence[str], + *, + dropna: bool = True, + as_index: bool = True, + ): + # TODO(tbergeron): Support more group-by expression types + self._block = block + self._col_id_labels = { + value_column: column_label + for value_column, column_label in zip( + block.value_columns, block.column_labels + ) + } + self._by_col_ids = by_col_ids + self._dropna = dropna # Applies to aggregations but not windowing + self._as_index = as_index + + def sum(self, numeric_only: bool = False, *args) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("sum") + return self._aggregate(agg_ops.sum_op, numeric_only=True) + + def mean(self, numeric_only: bool = False, *args) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("mean") + return self._aggregate(agg_ops.mean_op, numeric_only=True) + + def min(self, numeric_only: bool = False, *args) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("min") + return self._aggregate(agg_ops.min_op, numeric_only=True) + + def max(self, numeric_only: bool = False, *args) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("max") + return self._aggregate(agg_ops.max_op, numeric_only=True) + + def std( + self, + *, + numeric_only: bool = False, + ) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("std") + return self._aggregate(agg_ops.std_op, numeric_only=True) + + def var( + self, + *, + numeric_only: bool = False, + ) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("var") + return self._aggregate(agg_ops.var_op, numeric_only=True) + + def all(self) -> df.DataFrame: + return self._aggregate(agg_ops.all_op) + + def any(self) -> df.DataFrame: + return self._aggregate(agg_ops.any_op) + + def count(self) -> df.DataFrame: + return self._aggregate(agg_ops.count_op) + + def cumsum(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("cumsum") + window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) + return self._apply_window_op(agg_ops.sum_op, window, numeric_only=True) + + def cummin(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("cummin") + window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) + return self._apply_window_op(agg_ops.min_op, window, numeric_only=True) + + def cummax(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("cummax") + window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) + return self._apply_window_op(agg_ops.max_op, window, numeric_only=True) + + def cumprod(self, *args, **kwargs) -> df.DataFrame: + window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) + return self._apply_window_op(agg_ops.product_op, window, numeric_only=True) + + def _raise_on_non_numeric(self, op: str): + if not all( + dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES + for dtype in self._block.dtypes + ): + raise NotImplementedError( + f"'{op}' does not support non-numeric columns. Set 'numeric_only'=True to ignore non-numeric columns" + ) + return self + + def _aggregated_columns(self, numeric_only: bool = False): + return [ + col_id + for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) + if col_id not in self._by_col_ids + and ( + (not numeric_only) + or (dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) + ) + ] + + def _aggregate( + self, aggregate_op: agg_ops.AggregateOp, numeric_only: bool = False + ) -> df.DataFrame: + aggregated_col_ids = self._aggregated_columns(numeric_only=numeric_only) + aggregations = [(col_id, aggregate_op) for col_id in aggregated_col_ids] + result_block, _ = self._block.aggregate( + self._by_col_ids, + aggregations, + as_index=self._as_index, + dropna=self._dropna, + ) + return df.DataFrame(result_block) + + def _apply_window_op( + self, + op: agg_ops.WindowOp, + window_spec: bigframes.core.WindowSpec, + numeric_only: bool = False, + ): + columns = self._aggregated_columns(numeric_only=numeric_only) + block = self._block.select_columns([*columns, *window_spec.grouping_keys]) + block = self._block.multi_apply_window_op( + columns, + op, + window_spec=window_spec, + ) + block = block.select_columns(columns) + return df.DataFrame(block) + + +class SeriesGroupBy(vendored_pandas_groupby.SeriesGroupBy): + __doc__ = vendored_pandas_groupby.GroupBy.__doc__ + + def __init__( + self, + block: blocks.Block, + value_column: str, + by_col_ids: typing.Sequence[str], + value_name: typing.Optional[str] = None, + dropna=True, + ): + # TODO(tbergeron): Support more group-by expression types + self._block = block + self._value_column = value_column + self._by_col_ids = by_col_ids + self._value_name = value_name + self._dropna = dropna # Applies to aggregations but not windowing + + @property + def value(self): + return self._block.expr.get_column(self._value_column) + + def all(self) -> series.Series: + return self._aggregate(agg_ops.all_op) + + def any(self) -> series.Series: + return self._aggregate(agg_ops.any_op) + + def count(self) -> series.Series: + return self._aggregate(agg_ops.count_op) + + def sum(self, *args) -> series.Series: + """Sums the numeric values for each group in the series. Ignores null/nan.""" + return self._aggregate(agg_ops.sum_op) + + def mean(self, *args) -> series.Series: + return self._aggregate(agg_ops.mean_op) + + def std(self, *args, **kwargs) -> series.Series: + return self._aggregate(agg_ops.std_op) + + def var(self, *args, **kwargs) -> series.Series: + return self._aggregate(agg_ops.var_op) + + def prod(self, *args) -> series.Series: + return self._aggregate(agg_ops.product_op) + + def cumsum(self, *args, **kwargs) -> series.Series: + return self._apply_window_op( + agg_ops.sum_op, + bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), + ) + + def cumprod(self, *args, **kwargs) -> series.Series: + return self._apply_window_op( + agg_ops.product_op, + bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), + ) + + def cummax(self, *args, **kwargs) -> series.Series: + return self._apply_window_op( + agg_ops.max_op, + bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), + ) + + def cummin(self, *args, **kwargs) -> series.Series: + return self._apply_window_op( + agg_ops.min_op, + bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), + ) + + def cumcount(self, *args, **kwargs) -> series.Series: + return self._apply_window_op( + agg_ops.rank_op, + bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), + discard_name=True, + )._apply_unary_op(ops.partial_right(ops.sub_op, 1)) + + def shift(self, periods=1) -> series.Series: + """Shift index by desired number of periods.""" + window = bigframes.core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.ShiftOp(periods), window) + + def diff(self) -> series.Series: + """Difference between each element and previous element.""" + return self._ungroup() - self.shift(1) + + def rolling(self, window: int, min_periods=None) -> windows.Window: + # To get n size window, need current row and n-1 preceding rows. + window_spec = core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=window - 1, + following=0, + min_periods=min_periods or window, + ) + block = self._block.order_by( + [order.OrderingColumnReference(col) for col in self._by_col_ids], + stable=True, + ) + return windows.Window(block, window_spec, self._value_column) + + def expanding(self, min_periods: int = 1) -> windows.Window: + window_spec = core.WindowSpec( + grouping_keys=self._by_col_ids, following=0, min_periods=min_periods + ) + block = self._block.order_by( + [order.OrderingColumnReference(col) for col in self._by_col_ids], + stable=True, + ) + return windows.Window(block, window_spec, self._value_column) + + def _ungroup(self) -> series.Series: + return series.Series(self._block.select_column(self._value_column)) + + def _aggregate(self, aggregate_op: agg_ops.AggregateOp) -> series.Series: + result_block, _ = self._block.aggregate( + self._by_col_ids, + ((self._value_column, aggregate_op),), + dropna=self._dropna, + ) + + return series.Series(result_block.with_column_labels([self._value_name])) + + def _apply_window_op( + self, + op: agg_ops.WindowOp, + window_spec: bigframes.core.WindowSpec, + discard_name=False, + ): + label = self._value_name if not discard_name else None + block, result_id = self._block.apply_window_op( + self._value_column, + op, + result_label=label, + window_spec=window_spec, + skip_null_groups=self._dropna, + ) + return series.Series(block.select_column(result_id)) diff --git a/bigframes/core/guid.py b/bigframes/core/guid.py new file mode 100644 index 00000000000..4eb6c7a9d62 --- /dev/null +++ b/bigframes/core/guid.py @@ -0,0 +1,21 @@ +# 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. + +_GUID_COUNTER = 0 + + +def generate_guid(prefix="col_"): + global _GUID_COUNTER + _GUID_COUNTER += 1 + return prefix + str(_GUID_COUNTER) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py new file mode 100644 index 00000000000..f65cb02941e --- /dev/null +++ b/bigframes/core/indexers.py @@ -0,0 +1,337 @@ +# 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. + +from __future__ import annotations + +import typing +from typing import Tuple + +import ibis +import pandas as pd + +import bigframes.core as core +import bigframes.core.guid as guid +import bigframes.core.indexes as indexes +import bigframes.core.scalar +import bigframes.dataframe +import bigframes.series + +if typing.TYPE_CHECKING: + LocSingleKey = bigframes.series.Series | indexes.Index | slice + + +class LocSeriesIndexer: + def __init__(self, series: bigframes.series.Series): + self._series = series + + def __getitem__(self, key) -> bigframes.series.Series: + """ + Only indexing by a boolean bigframes.series.Series or list of index entries is currently supported + """ + return typing.cast( + bigframes.series.Series, _loc_getitem_series_or_dataframe(self._series, key) + ) + + def __setitem__(self, key, value) -> None: + # TODO(swast): support MultiIndex + if isinstance(key, slice): + # TODO(swast): Implement loc with slices. + raise NotImplementedError("loc does not yet support slices") + elif isinstance(key, list): + # TODO(tbergeron): Implement loc for index label list. + raise NotImplementedError("loc does not yet support index label lists") + + # Assume the key is for the index label. + block = self._series._block + value_column = self._series._value + index_column = block.expr.get_column(block.index_columns[0]) + new_value = ( + ibis.case() + .when( + index_column == ibis.literal(key, index_column.type()), + ibis.literal(value, value_column.type()), + ) + .else_(value_column) + .end() + .name(value_column.get_name()) + ) + all_columns = [] + for column in block.expr.columns: + if column.get_name() != value_column.get_name(): + all_columns.append(column) + else: + all_columns.append(new_value) + new_expr = block.expr.projection(all_columns) + + # TODO(tbergeron): Use block operators rather than directly building desired ibis expressions. + self._series._set_block( + core.blocks.Block( + new_expr, + self._series._block.index_columns, + self._series._block.column_labels, + self._series._block.index.names, + ) + ) + + +class IlocSeriesIndexer: + def __init__(self, series: bigframes.series.Series): + self._series = series + + def __getitem__( + self, key + ) -> bigframes.core.scalar.Scalar | bigframes.series.Series: + """ + Index series using integer offsets. Currently supports index by key type: + + slice: ex. series.iloc[2:5] returns values at index 2, 3, and 4 as a series + individual offset: ex. series.iloc[0] returns value at index 0 as a scalar + list: ex. series.iloc[1, 1, 2, 0] returns a series with the index 1 item repeated + twice, followed by the index 2 and then and 0 items in that order. + + Other key types are not yet supported. + """ + return _iloc_getitem_series_or_dataframe(self._series, key) + + +class LocDataFrameIndexer: + def __init__(self, dataframe: bigframes.dataframe.DataFrame): + self._dataframe = dataframe + + @typing.overload + def __getitem__(self, key: LocSingleKey) -> bigframes.dataframe.DataFrame: + ... + + # Technically this is wrong since we can have duplicate column labels, but + # this is expected to be rare. + @typing.overload + def __getitem__(self, key: Tuple[LocSingleKey, str]) -> bigframes.series.Series: + ... + + def __getitem__(self, key): + # TODO(swast): If the DataFrame has a MultiIndex, we'll need to + # disambiguate this from a single row selection. + if isinstance(key, tuple) and len(key) == 2: + df = typing.cast( + bigframes.dataframe.DataFrame, + _loc_getitem_series_or_dataframe(self._dataframe, key[0]), + ) + return df[key[1]] + + return typing.cast( + bigframes.dataframe.DataFrame, + _loc_getitem_series_or_dataframe(self._dataframe, key), + ) + + def __setitem__( + self, + key: Tuple[slice, str], + value: bigframes.dataframe.SingleItemValue, + ): + if ( + not isinstance(key, tuple) + or len(key) != 2 + or not isinstance(key[0], slice) + or (key[0].start is not None and key[0].start != 0) + or (key[0].step is not None and key[0].step != 1) + or key[0].stop is not None + ): + raise NotImplementedError( + "Only setting a column by DataFrame.loc[:, 'column'] is supported." + ) + + # TODO(swast): Support setting multiple columns with key[1] as a list + # of labels and value as a DataFrame. + df = self._dataframe.assign(**{key[1]: value}) + self._dataframe._set_block(df._get_block()) + + +class ILocDataFrameIndexer: + def __init__(self, dataframe: bigframes.dataframe.DataFrame): + self._dataframe = dataframe + + def __getitem__(self, key) -> bigframes.dataframe.DataFrame | pd.Series: + """ + Index dataframe using integer offsets. Currently supports index by key type: + + slice: i.e. df.iloc[2:5] returns rows at index 2, 3, and 4 as a dataframe + individual offset: i.e. df.iloc[0] returns row at index 0 as a pandas Series + + Other key types are not yet supported. + """ + return _iloc_getitem_series_or_dataframe(self._dataframe, key) + + +@typing.overload +def _loc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.series.Series, key +) -> bigframes.series.Series: + ... + + +@typing.overload +def _loc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.dataframe.DataFrame, key +) -> bigframes.dataframe.DataFrame: + ... + + +def _loc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, + key: LocSingleKey, +) -> bigframes.dataframe.DataFrame | bigframes.series.Series: + if isinstance(key, bigframes.series.Series) and key.dtype == "boolean": + return series_or_dataframe[key] + elif isinstance(key, bigframes.series.Series): + # TODO(henryjsolberg): support MultiIndex + temp_name = guid.generate_guid(prefix="temp_series_name_") + key = key.rename(temp_name) + keys_df = key.to_frame() + keys_df = keys_df.set_index(temp_name, drop=True) + return _perform_loc_list_join(series_or_dataframe, keys_df) + elif isinstance(key, bigframes.core.indexes.Index): + # TODO(henryjsolberg): support MultiIndex + block = key._data._get_block() + block = block.select_columns(()) + keys_df = bigframes.dataframe.DataFrame(block) + return _perform_loc_list_join(series_or_dataframe, keys_df) + elif pd.api.types.is_list_like(key): + # TODO(henryjsolberg): support MultiIndex + if len(key) == 0: # type: ignore + return typing.cast( + typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], + series_or_dataframe.iloc[0:0], + ) + index_name = series_or_dataframe.index.name + keys_df = bigframes.dataframe.DataFrame( + {index_name: key}, session=series_or_dataframe._get_block().expr._session + ) + keys_df = keys_df.set_index(index_name, drop=True) + return _perform_loc_list_join(series_or_dataframe, keys_df) + elif isinstance(key, slice): + return series_or_dataframe._slice(key.start, key.stop, key.step) + elif callable(key): + raise NotImplementedError("loc does not yet support indexing with a callable") + else: + raise TypeError( + "Invalid argument type. loc currently only supports indexing with a boolean bigframes Series or a list of index entries." + ) + + +@typing.overload +def _perform_loc_list_join( + series_or_dataframe: bigframes.series.Series, + keys_df: bigframes.dataframe.DataFrame, +) -> bigframes.series.Series: + ... + + +@typing.overload +def _perform_loc_list_join( + series_or_dataframe: bigframes.dataframe.DataFrame, + keys_df: bigframes.dataframe.DataFrame, +) -> bigframes.dataframe.DataFrame: + ... + + +def _perform_loc_list_join( + series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, + keys_df: bigframes.dataframe.DataFrame, +) -> bigframes.series.Series | bigframes.dataframe.DataFrame: + # right join based on the old index so that the matching rows from the user's + # original dataframe will be duplicated and reordered appropriately + original_index_names = series_or_dataframe.index.names + if isinstance(series_or_dataframe, bigframes.series.Series): + original_name = series_or_dataframe.name + name = series_or_dataframe.name if series_or_dataframe.name is not None else "0" + result = typing.cast( + bigframes.series.Series, + series_or_dataframe.to_frame().join(keys_df, how="right")[name], + ) + result = result.rename(original_name) + else: + result = series_or_dataframe.join(keys_df, how="right") # type: ignore + result = result.rename_axis(original_index_names) + return result + + +@typing.overload +def _iloc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.series.Series, key +) -> bigframes.series.Series | bigframes.core.scalar.Scalar: + ... + + +@typing.overload +def _iloc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.dataframe.DataFrame, key +) -> bigframes.dataframe.DataFrame | pd.Series: + ... + + +def _iloc_getitem_series_or_dataframe( + series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, key +) -> bigframes.dataframe.DataFrame | bigframes.series.Series | bigframes.core.scalar.Scalar | pd.Series: + if isinstance(key, int): + if key < 0: + raise NotImplementedError( + "iloc does not yet support negative single positional index" + ) + internal_slice_result = series_or_dataframe._slice(key, key + 1, 1) + result_pd_df = internal_slice_result.compute() + if result_pd_df.empty: + raise IndexError("single positional indexer is out-of-bounds") + return result_pd_df.iloc[0] + elif isinstance(key, slice): + return series_or_dataframe._slice(key.start, key.stop, key.step) + elif pd.api.types.is_list_like(key): + # TODO(henryjsolberg): support MultiIndex + + if len(key) == 0: + return typing.cast( + typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], + series_or_dataframe.iloc[0:0], + ) + df = series_or_dataframe + if isinstance(series_or_dataframe, bigframes.series.Series): + original_series_name = series_or_dataframe.name + series_name = ( + original_series_name if original_series_name is not None else "0" + ) + df = series_or_dataframe.to_frame() + original_index_name = df.index.name + temporary_index_name = guid.generate_guid(prefix="temp_iloc_index_") + df = df.rename_axis(temporary_index_name) + + # set to offset index and use regular loc, then restore index + df = df.reset_index(drop=False) + result = df.loc[key] + result = result.set_index(temporary_index_name) + result = result.rename_axis(original_index_name) + + if isinstance(series_or_dataframe, bigframes.series.Series): + result = result[series_name] + result = typing.cast(bigframes.series.Series, result) + result = result.rename(original_series_name) + + return result + + elif isinstance(key, tuple): + raise NotImplementedError( + "iloc does not yet support indexing with a (row, column) tuple" + ) + elif callable(key): + raise NotImplementedError("iloc does not yet support indexing with a callable") + else: + raise TypeError("Invalid argument type.") diff --git a/bigframes/core/indexes/__init__.py b/bigframes/core/indexes/__init__.py new file mode 100644 index 00000000000..d797c57955a --- /dev/null +++ b/bigframes/core/indexes/__init__.py @@ -0,0 +1,23 @@ +# 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. + +from bigframes.core.indexes.index import Index, IndexValue + +INDEX_COLUMN_ID = "bigframes_index_{}" + +__all__ = [ + "Index", + "IndexValue", + "INDEX_COLUMN_ID", +] diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py new file mode 100644 index 00000000000..f7fa7ac1195 --- /dev/null +++ b/bigframes/core/indexes/index.py @@ -0,0 +1,250 @@ +# 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. + +"""An index based on a single column.""" + +from __future__ import annotations + +import typing +from typing import Callable, Tuple + +import numpy as np +import pandas as pd + +import bigframes.core as core +import bigframes.core.blocks as blocks +import bigframes.core.joins as joins +import bigframes.dtypes as bf_dtypes +import third_party.bigframes_vendored.pandas.core.indexes.base as vendored_pandas_index + + +class Index(vendored_pandas_index.Index): + __doc__ = vendored_pandas_index.Index.__doc__ + + def __init__(self, data: blocks.BlockHolder): + self._data = data + + @property + def name(self) -> typing.Optional[str]: + return self.names[0] + + @name.setter + def name(self, value: blocks.Label): + self.names = [value] + + @property + def names(self) -> typing.Sequence[blocks.Label]: + return self._data._get_block()._index_labels + + @names.setter + def names(self, values: typing.Sequence[blocks.Label]): + return self._data._set_block(self._data._get_block().with_index_labels(values)) + + @property + def shape(self) -> typing.Tuple[int]: + return (self._data._get_block().shape[0],) + + @property + def size(self) -> int: + return self.shape[0] + + @property + def empty(self) -> bool: + return self.shape[0] == 0 + + def to_pandas(self) -> pd.Index: + """Get the Index as a pandas Index.""" + return IndexValue(self._data._get_block()).compute() + + def __len__(self): + return self.shape[0] + + compute = to_pandas + + +class IndexValue: + """An immutable index.""" + + def __init__(self, block: blocks.Block): + self._block = block + + @property + def _expr(self) -> core.ArrayValue: + return self._block.expr + + @property + def name(self) -> typing.Optional[str]: + return self._block._index_labels[0] + + @property + def names(self) -> typing.Sequence[typing.Optional[str]]: + return self._block._index_labels + + @property + def nlevels(self) -> int: + return len(self._block._index_columns) + + @property + def dtypes( + self, + ) -> typing.Sequence[typing.Union[bf_dtypes.Dtype, np.dtype[typing.Any]]]: + return self._block.index_dtypes + + def __repr__(self) -> str: + """Converts an Index to a string.""" + # TODO(swast): Add a timeout here? If the query is taking a long time, + # maybe we just print the job metadata that we have so far? + # TODO(swast): Avoid downloading the whole index by using job + # metadata, like we do with DataFrame. + preview = self.compute() + return repr(preview) + + def compute(self) -> pd.Index: + """Executes deferred operations and downloads the results.""" + # Project down to only the index column. So the query can be cached to visualize other data. + index_column = self._block.index_columns[0] + expr = self._expr.projection([self._expr.get_any_column(index_column)]) + results, _ = expr.start_query() + df = expr._session._rows_to_dataframe(results) + df.set_index(index_column) + index = df.index + index.name = self._block._index_labels[0] + return index + + def join( + self, other: IndexValue, *, how="left", sort=False + ) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: + if not isinstance(other, IndexValue): + # TODO(swast): We need to improve this error message to be more + # actionable for the user. For example, it's possible they + # could call set_index and try again to resolve this error. + raise ValueError( + "Can't mixed objects with explicit Index and ImpliedJoiner" + ) + + # TODO(swast): Support cross-joins (requires reindexing). + if how not in {"outer", "left", "right", "inner"}: + raise NotImplementedError( + "Only how='outer','left','right','inner' currently supported" + ) + if self.nlevels == other.nlevels == 1: + return join_mono_indexed(self, other, how=how, sort=sort) + else: + # Always sort mult-index join + return join_multi_indexed(self, other, how=how, sort=sort) + + def resolve_level_name(self: IndexValue, label: blocks.Label) -> str: + matches = self._block.index_name_to_col_id.get(label, []) + if len(matches) > 1: + raise ValueError(f"Ambiguous index level name {label}") + if len(matches) == 0: + raise ValueError(f"Cannot resolve index level name {label}") + return matches[0] + + def is_uniquely_named(self: IndexValue): + return len(set(self.names)) == len(self.names) + + +def join_mono_indexed( + left: IndexValue, right: IndexValue, *, how="left", sort=False +) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: + ( + combined_expr, + joined_index_col_names, + (get_column_left, get_column_right), + ) = joins.join_by_column( + left._block.expr, + left._block.index_columns, + right._block.expr, + right._block.index_columns, + how=how, + sort=sort, + ) + # Drop original indices from each side. and used the coalesced combination generated by the join. + left_indices = [get_column_left(col_id) for col_id in left._block.index_columns] + right_indices = [get_column_right(col_id) for col_id in right._block.index_columns] + combined_expr = combined_expr.drop_columns(left_indices).drop_columns(right_indices) + block = blocks.Block( + combined_expr, + index_columns=[*joined_index_col_names], + column_labels=[*left._block.column_labels, *right._block.column_labels], + index_labels=[left.name] if left.name == right.name else [None], + ) + return ( + typing.cast(IndexValue, block.index), + (get_column_left, get_column_right), + ) + + +def join_multi_indexed( + left: IndexValue, right: IndexValue, *, how="left", sort=False +) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: + if not (left.is_uniquely_named() and right.is_uniquely_named()): + raise ValueError("Joins not supported on indices with non-unique level names") + + common_names = [name for name in left.names if name in right.names] + if len(common_names) == 0: + raise ValueError("Cannot join without a index level in common.") + + left_only_names = [name for name in left.names if name not in right.names] + right_only_names = [name for name in right.names if name not in left.names] + + left_join_ids = [left.resolve_level_name(name) for name in common_names] + right_join_ids = [right.resolve_level_name(name) for name in common_names] + + ( + combined_expr, + joined_index_col_names, + (get_column_left, get_column_right), + ) = joins.join_by_column( + left._block.expr, + left_join_ids, + right._block.expr, + right_join_ids, + how=how, + sort=sort, + ) + # Drop original indices from each side. and used the coalesced combination generated by the join. + combined_expr = combined_expr.drop_columns( + [get_column_left(col) for col in left_join_ids] + ).drop_columns([get_column_right(col) for col in right_join_ids]) + + if left.nlevels == 1: + index_labels = right.names + elif right.nlevels == 1: + index_labels = left.names + else: + index_labels = [*common_names, *left_only_names, *right_only_names] + + def resolve_label_id(label: blocks.Label) -> str: + if label in common_names: + return joined_index_col_names[common_names.index(label)] + if label in left_only_names: + return get_column_left(left.resolve_level_name(label)) + if label in right_only_names: + return get_column_right(right.resolve_level_name(label)) + raise ValueError(f"Unexpected label: {label}") + + index_columns = [resolve_label_id(label) for label in index_labels] + + block = blocks.Block( + combined_expr, + index_columns=index_columns, + column_labels=[*left._block.column_labels, *right._block.column_labels], + index_labels=index_labels, + ) + return ( + typing.cast(IndexValue, block.index), + (get_column_left, get_column_right), + ) diff --git a/bigframes/core/joins/__init__.py b/bigframes/core/joins/__init__.py new file mode 100644 index 00000000000..58a1c2cfd79 --- /dev/null +++ b/bigframes/core/joins/__init__.py @@ -0,0 +1,23 @@ +# 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. + +"""Helpers to join ArrayValue objects.""" + +from bigframes.core.joins.row_identity import join_by_row_identity +from bigframes.core.joins.single_column import join_by_column + +__all__ = ( + "join_by_row_identity", + "join_by_column", +) diff --git a/bigframes/core/joins/row_identity.py b/bigframes/core/joins/row_identity.py new file mode 100644 index 00000000000..b77b5d80adc --- /dev/null +++ b/bigframes/core/joins/row_identity.py @@ -0,0 +1,201 @@ +# 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. + +"""Helpers to join ArrayValue objects.""" + +from __future__ import annotations + +import functools +import typing +from typing import Callable, Tuple + +import ibis +import ibis.expr.types as ibis_types + +import bigframes.core as core + +SUPPORTED_ROW_IDENTITY_HOW = {"outer", "left", "inner"} + + +def join_by_row_identity( + left: core.ArrayValue, right: core.ArrayValue, *, how: str +) -> Tuple[core.ArrayValue, Tuple[Callable[[str], str], Callable[[str], str]],]: + """Compute join when we are joining by row identity not a specific column.""" + if how not in SUPPORTED_ROW_IDENTITY_HOW: + raise NotImplementedError("Only how='outer','left','inner' currently supported") + + if not left.table.equals(right.table): + raise ValueError( + "Cannot combine objects without an explicit join/merge key. " + f"Left based on: {left.table.compile()}, but " + f"right based on: {right.table.compile()}" + ) + + left_predicates = left._predicates + right_predicates = right._predicates + # TODO(tbergeron): Skip generating these for inner part of join + ( + left_relative_predicates, + right_relative_predicates, + ) = _get_relative_predicates(left_predicates, right_predicates) + + combined_predicates = [] + if left_predicates or right_predicates: + joined_predicates = _join_predicates( + left_predicates, right_predicates, join_type=how + ) + combined_predicates = list(joined_predicates) # builder expects mutable list + + left_mask = left_relative_predicates if how in ["right", "outer"] else None + right_mask = right_relative_predicates if how in ["left", "outer"] else None + joined_columns = [ + _mask_value(left.get_column(key), left_mask).name(map_left_id(key)) + for key in left.column_names.keys() + ] + [ + _mask_value(right.get_column(key), right_mask).name(map_right_id(key)) + for key in right.column_names.keys() + ] + + hidden_ordering_columns = [] + new_ordering = core.ExpressionOrdering() + if left._ordering and right._ordering: + # These ordering columns will be present in the ArrayValue, as we + # haven't hidden any value / index column(s). Code that is aware of + # which columns are index columns / value columns columns will need to + # add the previous columns to hidden columns. + new_ordering = left._ordering.with_ordering_columns( + [ + col_ref.with_name(map_left_id(col_ref.column_id)) + for col_ref in left._ordering.ordering_value_columns + ] + + [ + col_ref.with_name(map_right_id(col_ref.column_id)) + for col_ref in right._ordering.ordering_value_columns + ] + ) + + hidden_ordering_columns = [ + left._get_hidden_ordering_column(key.column_id).name( + map_left_id(key.column_id) + ) + for key in left._ordering.ordering_value_columns + if key.column_id in left._hidden_ordering_column_names.keys() + ] + [ + right._get_hidden_ordering_column(key.column_id).name( + map_right_id(key.column_id) + ) + for key in right._ordering.ordering_value_columns + if key.column_id in right._hidden_ordering_column_names.keys() + ] + + left_ordering_id = left._ordering.ordering_id + if left_ordering_id: + new_ordering = new_ordering.with_ordering_id(map_left_id(left_ordering_id)) + if left_ordering_id in left._hidden_ordering_column_names.keys(): + hidden_ordering_columns.append( + left._get_hidden_ordering_column(left_ordering_id).name( + map_left_id(left_ordering_id) + ) + ) + + joined_expr = core.ArrayValue( + left._session, + left.table, + columns=joined_columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=new_ordering, + predicates=combined_predicates, + ) + return joined_expr, ( + lambda key: map_left_id(key), + lambda key: map_right_id(key), + ) + + +def map_left_id(left_side_id): + return f"{left_side_id}_x" + + +def map_right_id(right_side_id): + return f"{right_side_id}_y" + + +def _mask_value( + value: ibis_types.Value, + predicates: typing.Optional[typing.Sequence[ibis_types.BooleanValue]] = None, +): + if predicates: + return ( + ibis.case() + .when(_reduce_predicate_list(predicates), value) + .else_(ibis.null()) + .end() + ) + return value + + +def _join_predicates( + left_predicates: typing.Collection[ibis_types.BooleanValue], + right_predicates: typing.Collection[ibis_types.BooleanValue], + join_type: str = "outer", +) -> typing.Tuple[ibis_types.BooleanValue, ...]: + """Combines predicates lists for each side of a join.""" + if join_type == "outer": + if not left_predicates: + return () + if not right_predicates: + return () + # TODO(tbergeron): Investigate factoring out common predicates + joined_predicates = _reduce_predicate_list(left_predicates).__or__( + _reduce_predicate_list(right_predicates) + ) + return (joined_predicates,) + if join_type == "left": + return tuple(left_predicates) + if join_type == "inner": + _, right_relative_predicates = _get_relative_predicates( + left_predicates, right_predicates + ) + return (*left_predicates, *right_relative_predicates) + else: + raise ValueError("Unsupported join_type: " + join_type) + + +def _get_relative_predicates( + left_predicates: typing.Collection[ibis_types.BooleanValue], + right_predicates: typing.Collection[ibis_types.BooleanValue], +) -> tuple[ + typing.Tuple[ibis_types.BooleanValue, ...], + typing.Tuple[ibis_types.BooleanValue, ...], +]: + """Get predicates that apply to only one side of the join. Not strictly necessary but simplifies resulting query.""" + left_relative_predicates = tuple(left_predicates) or () + right_relative_predicates = tuple(right_predicates) or () + if left_predicates and right_predicates: + # Factor out common predicates needed for left/right column masking + left_relative_predicates = tuple(set(left_predicates) - set(right_predicates)) + right_relative_predicates = tuple(set(right_predicates) - set(left_predicates)) + return (left_relative_predicates, right_relative_predicates) + + +def _reduce_predicate_list( + predicate_list: typing.Collection[ibis_types.BooleanValue], +) -> ibis_types.BooleanValue: + """Converts a list of predicates BooleanValues into a single BooleanValue.""" + if len(predicate_list) == 0: + raise ValueError("Cannot reduce empty list of predicates") + if len(predicate_list) == 1: + (item,) = predicate_list + return item + return functools.reduce(lambda acc, pred: acc.__and__(pred), predicate_list) diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py new file mode 100644 index 00000000000..34873e6f7ee --- /dev/null +++ b/bigframes/core/joins/single_column.py @@ -0,0 +1,281 @@ +# 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. + +"""Helpers to join ArrayValue objects.""" + +from __future__ import annotations + +import typing +from typing import Callable, Literal, Tuple + +import ibis +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types + +import bigframes.core as core +import bigframes.core.guid +import bigframes.core.joins.row_identity +import bigframes.core.ordering + + +def join_by_column( + left: core.ArrayValue, + left_column_ids: typing.Sequence[str], + right: core.ArrayValue, + right_column_ids: typing.Sequence[str], + *, + how: Literal[ + "inner", + "left", + "outer", + "right", + ], + sort: bool = False, + get_both_join_key_cols: bool = False, +) -> Tuple[ + core.ArrayValue, + typing.Sequence[str], + Tuple[Callable[[str], str], Callable[[str], str]], +]: + """Join two expressions by column equality. + + Arguments: + left: Expression for left table to join. + left_column_ids: Column IDs (not label) to join by. + right: Expression for right table to join. + right_column_ids: Column IDs (not label) to join by. + how: The type of join to perform. + get_both_join_key_cols: if set to True, returned column ids will contain + both left and right join key columns. + + Returns: + The joined expression and the objects needed to interpret it. + + * ArrayValue: Joined table with all columns from left and right. + * Sequence[str]: Column IDs of the coalesced join columns. Sometimes either the + left/right table will have missing rows. This column pulls the + non-NULL value from either left/right. + If get_both_join_key_cols is True, will return uncombined left and + right key columns. + * Tuple[Callable, Callable]: For a given column ID from left or right, + respectively, return the new column id from the combined expression. + """ + + if ( + how in bigframes.core.joins.row_identity.SUPPORTED_ROW_IDENTITY_HOW + and left.table.equals(right.table) + # Compare ibis expressions for left/right columns because its possible that + # they both have the same names but were modified in different ways. + and all( + left.get_any_column(lcol).equals(right.get_any_column(rcol)) + for lcol, rcol in zip(left_column_ids, right_column_ids) + ) + ): + combined_expr, ( + get_column_left, + get_column_right, + ) = bigframes.core.joins.row_identity.join_by_row_identity(left, right, how=how) + original_ordering = combined_expr._ordering + else: + # Generate offsets if non-default ordering is applied + # Assumption, both sides are totally ordered, otherwise offsets will be nondeterministic + left_table = left.to_ibis_expr( + ordering_mode="ordered_col", order_col_name=core.ORDER_ID_COLUMN + ) + right_table = right.to_ibis_expr( + ordering_mode="ordered_col", order_col_name=core.ORDER_ID_COLUMN + ) + join_conditions = [ + value_to_join_key(left_table[left_index]) + == value_to_join_key(right_table[right_index]) + for left_index, right_index in zip(left_column_ids, right_column_ids) + ] + + combined_table = ibis.join( + left_table, + right_table, + predicates=join_conditions, + how=how, + lname="{name}_x", + rname="{name}_y", + ) + + def get_column_left(key: str) -> str: + if ( + how == "inner" + and key in left_column_ids + and key in combined_table.columns + ): + # Ibis doesn't rename the column if the values are guaranteed + # to be equal on left and right (because they're part of an + # inner join condition). See: + # https://github.com/ibis-project/ibis/pull/4651 + pass + elif key in right_table.columns: + key = f"{key}_x" + + return key + + def get_column_right(key: str) -> str: + if ( + how == "inner" + and key in right_column_ids + and key in combined_table.columns + ): + # Ibis doesn't rename the column if the values are guaranteed + # to be equal on left and right (because they're part of an + # inner join condition). See: + # https://github.com/ibis-project/ibis/pull/4651 + pass + elif key in left_table.columns: + key = f"{key}_y" + + return key + + left_ordering_encoding_size = ( + left._ordering.ordering_encoding_size + or bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH + ) + right_ordering_encoding_size = ( + right._ordering.ordering_encoding_size + or bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH + ) + + # Preserve original ordering accross joins. + left_order_id = get_column_left(core.ORDER_ID_COLUMN) + right_order_id = get_column_right(core.ORDER_ID_COLUMN) + new_order_id_col = _merge_order_ids( + combined_table[left_order_id], + left_ordering_encoding_size, + combined_table[right_order_id], + right_ordering_encoding_size, + how, + ) + new_order_id = new_order_id_col.get_name() + if new_order_id is None: + raise ValueError("new_order_id unexpectedly has no name") + hidden_columns = (new_order_id_col,) + original_ordering = core.ExpressionOrdering( + ordering_id_column=core.OrderingColumnReference(new_order_id) + if (new_order_id_col is not None) + else None, + ordering_encoding_size=left_ordering_encoding_size + + right_ordering_encoding_size, + ) + combined_expr = core.ArrayValue( + left._session, + combined_table, + hidden_ordering_columns=hidden_columns, + ) + + join_key_cols: list[ibis_types.Value] = [] + for lcol, rcol in zip(left_column_ids, right_column_ids): + if get_both_join_key_cols: + join_key_cols.append( + combined_expr.get_column(get_column_left(lcol)).name( + bigframes.core.guid.generate_guid(prefix="index_") + ) + ) + join_key_cols.append( + combined_expr.get_column(get_column_right(rcol)).name( + bigframes.core.guid.generate_guid(prefix="index_") + ) + ) + else: + if how == "left" or how == "inner": + join_key_cols.append( + combined_expr.get_column(get_column_left(lcol)).name( + bigframes.core.guid.generate_guid(prefix="index_") + ) + ) + elif how == "right": + join_key_cols.append( + combined_expr.get_column(get_column_right(rcol)).name( + bigframes.core.guid.generate_guid(prefix="index_") + ) + ) + elif how == "outer": + # The left index and the right index might contain null values, for + # example due to an outer join with different numbers of rows. Coalesce + # these to take the index value from either column. + # Use a random name in case the left index and the right index have the + # same name. In such a case, _x and _y suffixes will already be used. + join_key_cols.append( + ibis.coalesce( + combined_expr.get_column(get_column_left(lcol)), + combined_expr.get_column(get_column_right(rcol)), + ).name(bigframes.core.guid.generate_guid(prefix="index_")) + ) + else: + raise ValueError(f"Unexpected join type: {how}") + + # We could filter out the original join columns, but predicates/ordering + # might still reference them in implicit joins. + columns = ( + join_key_cols + + [ + combined_expr.get_column(get_column_left(key)) + for key in left.column_names.keys() + ] + + [ + combined_expr.get_column(get_column_right(key)) + for key in right.column_names.keys() + ] + ) + + if sort: + ordering = original_ordering.with_ordering_columns( + [ + core.OrderingColumnReference(join_key_col.get_name()) + for join_key_col in join_key_cols + ] + ) + else: + ordering = original_ordering + + combined_expr_builder = combined_expr.builder() + combined_expr_builder.columns = columns + combined_expr_builder.ordering = ordering + combined_expr = combined_expr_builder.build() + return ( + combined_expr, + [key.get_name() for key in join_key_cols], + (get_column_left, get_column_right), + ) + + +def value_to_join_key(value: ibis_types.Value): + """Converts nullable values to non-null string SQL will not match null keys together - but pandas does.""" + if not value.type().is_string(): + value = value.cast(ibis_dtypes.str) + return value.fillna(ibis_types.literal("$NULL_SENTINEL$")) + + +def _merge_order_ids( + left_id: ibis_types.Value, + left_encoding_size: int, + right_id: ibis_types.Value, + right_encoding_size: int, + how: str, +) -> ibis_types.StringValue: + if how == "right": + return _merge_order_ids( + right_id, right_encoding_size, left_id, left_encoding_size, "left" + ) + return ( + ( + bigframes.core.ordering.stringify_order_id(left_id, left_encoding_size) + + bigframes.core.ordering.stringify_order_id(right_id, right_encoding_size) + ) + ).name(bigframes.core.guid.generate_guid(prefix="bigframes_ordering_id_")) diff --git a/bigframes/core/ordering.py b/bigframes/core/ordering.py new file mode 100644 index 00000000000..6f8d35e52ae --- /dev/null +++ b/bigframes/core/ordering.py @@ -0,0 +1,181 @@ +# 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. + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +import math +import typing +from typing import Optional, Sequence + +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types + +# TODO(tbergeron): Encode more efficiently +ORDERING_ID_STRING_BASE: int = 10 +# Sufficient to store any value up to 2^63 +DEFAULT_ORDERING_ID_LENGTH: int = math.ceil(63 * math.log(2, ORDERING_ID_STRING_BASE)) + + +class OrderingDirection(Enum): + ASC = 1 + DESC = 2 + + def reverse(self): + if self == OrderingDirection.ASC: + return OrderingDirection.DESC + else: + return OrderingDirection.ASC + + @property + def is_ascending(self) -> bool: + return self == OrderingDirection.ASC + + +@dataclass(frozen=True) +class OrderingColumnReference: + """References a column and how to order with respect to values in that column.""" + + column_id: str + direction: OrderingDirection = OrderingDirection.ASC + na_last: bool = True + + def with_name(self, name: str): + return OrderingColumnReference(name, self.direction, self.na_last) + + def with_reverse(self): + return OrderingColumnReference( + self.column_id, self.direction.reverse(), not self.na_last + ) + + +@dataclass(frozen=True) +class ExpressionOrdering: + """Immutable object that holds information about the ordering of rows in a ArrayValue object.""" + + ordering_value_columns: Sequence[OrderingColumnReference] = () + ordering_id_column: Optional[OrderingColumnReference] = None + is_sequential: bool = False + # Encoding size must be tracked in order to know what how to combine ordering ids across tables (eg how much to pad when combining different length). + # Also will be needed to determine when length is too large and need to compact ordering id with a ROW_NUMBER operation. + ordering_encoding_size: int = DEFAULT_ORDERING_ID_LENGTH + + def with_is_sequential(self, is_sequential: bool): + """Create a copy that is marked as non-sequential. + + This is useful when filtering, but not sorting, an expression. + """ + return ExpressionOrdering( + self.ordering_value_columns, + self.ordering_id_column, + is_sequential, + ordering_encoding_size=self.ordering_encoding_size, + ) + + def with_ordering_columns( + self, + ordering_value_columns: Sequence[OrderingColumnReference] = (), + stable: bool = False, + ): + """Creates a new ordering that preserves ordering id, but replaces ordering value column list.""" + if stable: + col_ids_new = [ + ordering_ref.column_id for ordering_ref in ordering_value_columns + ] + # Only reference each column once, so discard old referenc if there is a new reference + old_ordering_keep = [ + ordering_ref + for ordering_ref in self.ordering_value_columns + if ordering_ref.column_id not in col_ids_new + ] + new_ordering = (*ordering_value_columns, *old_ordering_keep) + else: # Not stable, so discard old ordering completely + new_ordering = tuple(ordering_value_columns) + return ExpressionOrdering( + new_ordering, + self.ordering_id_column, + is_sequential=False, + ordering_encoding_size=self.ordering_encoding_size, + ) + + def with_ordering_id(self, ordering_id: str): + """Creates a new ordering that preserves other properties, but with a different ordering id. + + Useful when reprojecting ordering for implicit joins. + """ + return ExpressionOrdering( + self.ordering_value_columns, + OrderingColumnReference(ordering_id), + is_sequential=self.is_sequential, + ordering_encoding_size=self.ordering_encoding_size, + ) + + def with_reverse(self): + """Reverses the ordering.""" + return ExpressionOrdering( + tuple([col.with_reverse() for col in self.ordering_value_columns]), + self.ordering_id_column.with_reverse() + if self.ordering_id_column is not None + else None, + is_sequential=False, + ordering_encoding_size=self.ordering_encoding_size, + ) + + @property + def ordering_id(self) -> Optional[str]: + return self.ordering_id_column.column_id if self.ordering_id_column else None + + @property + def order_id_defined(self) -> bool: + """True if ordering is fully defined in ascending order by its ordering id.""" + return bool( + self.ordering_id_column + and (not self.ordering_value_columns) + and self.ordering_id_column.direction == OrderingDirection.ASC + ) + + @property + def all_ordering_columns(self) -> Sequence[OrderingColumnReference]: + return ( + list(self.ordering_value_columns) + if self.ordering_id_column is None + else [*self.ordering_value_columns, self.ordering_id_column] + ) + + +def stringify_order_id( + order_id: ibis_types.Value, length: int = DEFAULT_ORDERING_ID_LENGTH +) -> ibis_types.StringValue: + """Converts an order id value to string if it is not already a string. MUST produced fixed-length strings.""" + if order_id.type().is_int64(): + # This is very inefficient encoding base-10 string uses only 10 characters per byte(out of 256 bit combinations) + # Furthermore, if know tighter bounds on order id are known, can produce smaller strings. + # 19 characters chosen as it can represent any positive Int64 in base-10 + # For missing values, ":" * 19 is used as it is larger than any other value this function produces, so null values will be last. + string_order_id = ( + typing.cast( + ibis_types.StringValue, + typing.cast(ibis_types.IntegerValue, order_id).cast(ibis_dtypes.string), + ) + .lpad(length, "0") + .fillna(ibis_types.literal(":" * length)) + ) + else: + string_order_id = ( + typing.cast(ibis_types.StringValue, order_id) + .lpad(length, "0") + .fillna(ibis_types.literal(":" * length)) + ) + return typing.cast(ibis_types.StringValue, string_order_id) diff --git a/bigframes/core/reshape/__init__.py b/bigframes/core/reshape/__init__.py new file mode 100644 index 00000000000..1443a8e08fc --- /dev/null +++ b/bigframes/core/reshape/__init__.py @@ -0,0 +1,61 @@ +# 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. +from __future__ import annotations + +import typing +from typing import Iterable, Literal, Union + +import bigframes.dataframe +import bigframes.series + + +@typing.overload +def concat( + objs: Iterable[bigframes.dataframe.DataFrame], *, join, ignore_index +) -> bigframes.dataframe.DataFrame: + ... + + +@typing.overload +def concat( + objs: Iterable[bigframes.series.Series], *, join, ignore_index +) -> bigframes.series.Series: + ... + + +def concat( + objs: Union[ + Iterable[bigframes.dataframe.DataFrame], Iterable[bigframes.series.Series] + ], + *, + join: Literal["inner", "outer"] = "outer", + ignore_index: bool = False, +) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: + contains_dataframes = any( + isinstance(x, bigframes.dataframe.DataFrame) for x in objs + ) + if not contains_dataframes: + # Special case, all series, so align everything into single column even if labels don't match + series = typing.cast(typing.Iterable[bigframes.series.Series], objs) + names = {s.name for s in series} + # For series case, labels are stripped if they don't all match + if len(names) > 1: + blocks = [s._block.with_column_labels([None]) for s in series] + else: + blocks = [s._block for s in series] + block = blocks[0].concat(blocks[1:], how=join, ignore_index=ignore_index) + return bigframes.series.Series(block) + blocks = [obj._block for obj in objs] + block = blocks[0].concat(blocks[1:], how=join, ignore_index=ignore_index) + return bigframes.dataframe.DataFrame(block) diff --git a/bigframes/core/scalar.py b/bigframes/core/scalar.py new file mode 100644 index 00000000000..67ced525b99 --- /dev/null +++ b/bigframes/core/scalar.py @@ -0,0 +1,48 @@ +# 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. + +from __future__ import annotations + +import typing +from typing import Any + +import ibis.expr.types as ibis_types + +if typing.TYPE_CHECKING: + import bigframes.session + + +class DeferredScalar: + """A deferred scalar object.""" + + def __init__(self, value: ibis_types.Scalar, session: bigframes.session.Session): + self._value = value + self._session = session + + def __repr__(self) -> str: + """Converts a Series to a string.""" + # TODO(swast): Add a timeout here? If the query is taking a long time, + # maybe we just print the job metadata that we have so far? + return repr(self.compute()) + + def compute(self) -> Any: + """Executes deferred operations and downloads the resulting scalar.""" + result, _ = self._session._start_query(self._value.compile()) + df = self._session._rows_to_dataframe(result) + return df.iloc[0, 0] + + +# All public APIs return Any at present +# Later implementation may sometimes return a lazy scalar +Scalar = Any diff --git a/bigframes/core/window/__init__.py b/bigframes/core/window/__init__.py new file mode 100644 index 00000000000..8994004e0b8 --- /dev/null +++ b/bigframes/core/window/__init__.py @@ -0,0 +1,85 @@ +# 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. + +from __future__ import annotations + +import typing + +import bigframes.core as core +import bigframes.core.blocks as blocks +import bigframes.operations.aggregations as agg_ops +import third_party.bigframes_vendored.pandas.core.window.rolling as vendored_pandas_rolling + +if typing.TYPE_CHECKING: + from bigframes.series import Series + + +class Window(vendored_pandas_rolling.Window): + __doc__ = vendored_pandas_rolling.Window.__doc__ + + # TODO(tbergeron): Windows with groupings should create multi-indexed results + + def __init__( + self, + block: blocks.Block, + window_spec: core.WindowSpec, + value_column_id: str, + ): + self._block = block + self._window_spec = window_spec + self._value_column_id = value_column_id + + def count(self) -> Series: + return self._apply_aggregate(agg_ops.count_op) + + def sum(self) -> Series: + return self._apply_aggregate(agg_ops.sum_op) + + def mean(self) -> Series: + return self._apply_aggregate(agg_ops.mean_op) + + def var(self) -> Series: + return self._apply_aggregate(agg_ops.var_op) + + def std(self) -> Series: + return self._apply_aggregate(agg_ops.std_op) + + def max(self) -> Series: + return self._apply_aggregate(agg_ops.max_op) + + def min(self) -> Series: + return self._apply_aggregate(agg_ops.min_op) + + def _apply_aggregate( + self, + op: agg_ops.AggregateOp, + ) -> Series: + block = self._block + label = block.col_id_to_label[self._value_column_id] + block, result_id = block.apply_window_op( + self._value_column_id, op, self._window_spec, result_label=label + ) + + if self._window_spec.grouping_keys: + original_index_ids = block.index_columns + block = block.reset_index(drop=False) + index_ids = ( + *[col for col in self._window_spec.grouping_keys], + *original_index_ids, + ) + block = block.set_index(col_ids=index_ids) + + from bigframes.series import Series + + return Series(block.select_column(result_id)) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py new file mode 100644 index 00000000000..0155aafa0b0 --- /dev/null +++ b/bigframes/dataframe.py @@ -0,0 +1,1626 @@ +# 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. + +"""DataFrame is a two dimensional data structure.""" + +from __future__ import annotations + +import random +import re +import textwrap +import typing +from typing import ( + Callable, + Iterable, + List, + Literal, + Mapping, + Optional, + Sequence, + Tuple, + Union, +) + +import google.cloud.bigquery as bigquery +import ibis.expr.datatypes as ibis_dtypes +import numpy +import pandas as pd +import typing_extensions + +import bigframes +import bigframes._config.display_options as display_options +import bigframes.core +import bigframes.core.block_transforms as block_ops +import bigframes.core.blocks as blocks +import bigframes.core.groupby as groupby +import bigframes.core.guid +import bigframes.core.indexers as indexers +import bigframes.core.indexes as indexes +import bigframes.core.joins as joins +import bigframes.core.ordering as order +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops +import bigframes.series +import bigframes.series as bf_series +import third_party.bigframes_vendored.pandas.core.frame as vendored_pandas_frame +import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common +import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing + +if typing.TYPE_CHECKING: + import bigframes.session + + +# BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. +# TODO(tbergeron): Convert to bytes-based limit +MAX_INLINE_DF_SIZE = 5000 + +LevelsType = typing.Union[str, int, typing.Sequence[typing.Union[str, int]]] +SingleItemValue = Union[bigframes.series.Series, int, float, Callable] + + +# Inherits from pandas DataFrame so that we can use the same docstrings. +class DataFrame(vendored_pandas_frame.DataFrame): + __doc__ = vendored_pandas_frame.DataFrame.__doc__ + + def __init__( + self, + data=None, + index: vendored_pandas_typing.Axes | None = None, + columns: vendored_pandas_typing.Axes | None = None, + dtype: typing.Optional[ + bigframes.dtypes.DtypeString | bigframes.dtypes.Dtype + ] = None, + copy: typing.Optional[bool] = None, + *, + session: typing.Optional[bigframes.session.Session] = None, + ): + if copy is not None and not copy: + raise ValueError("DataFrame constructor only supports copy=True") + + # Check to see if constructing from BigQuery-backed objects before + # falling back to pandas constructor + block = None + if isinstance(data, blocks.Block): + block = data + + elif isinstance(data, DataFrame): + block = data._get_block() + + # Dict of Series + elif ( + _is_dict_like(data) + and len(data) >= 1 + and any(isinstance(data[key], bf_series.Series) for key in data.keys()) + ): + if not all(isinstance(data[key], bf_series.Series) for key in data.keys()): + # TODO(tbergeron): Support local list/series data by converting to memtable. + raise NotImplementedError("Cannot mix Series with other types.") + keys = list(data.keys()) + first_label, first_series = keys[0], data[keys[0]] + block = ( + typing.cast(bf_series.Series, first_series) + ._get_block() + .with_column_labels([first_label]) + ) + + for key in keys[1:]: + other = typing.cast(bf_series.Series, data[key]) + other_block = other._block.with_column_labels([key]) + # Pandas will keep original sorting if all indices are aligned. + # We cannot detect this easily however, and so always sort on index + result_index, _ = block.index.join( # type:ignore + other_block.index, how="outer", sort=True + ) + block = result_index._block + + if block: + if index: + raise NotImplementedError( + "DataFrame 'index' constructor parameter not supported when passing BigQuery-backed objects" + ) + if columns: + block = block.select_columns(list(columns)) # type:ignore + if dtype: + block = block.multi_apply_unary_op( + block.value_columns, ops.AsTypeOp(dtype) + ) + self._block = block + + else: + import bigframes.pandas + + pd_dataframe = pd.DataFrame( + data=data, + index=index, # type:ignore + columns=columns, # type:ignore + dtype=dtype, # type:ignore + ) + if pd_dataframe.size < MAX_INLINE_DF_SIZE: + self._block = blocks.block_from_local( + pd_dataframe, session or bigframes.pandas.get_global_session() + ) + elif session: + self._block = session.read_pandas(pd_dataframe)._get_block() + else: + self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block() + self._query_job: Optional[bigquery.QueryJob] = None + + def __dir__(self): + return dir(type(self)) + self._block.column_labels + + def _ipython_key_completions_(self) -> List[str]: + return list([label for label in self._block.column_labels if label]) + + def _find_indices( + self, + columns: Union[blocks.Label, Sequence[blocks.Label]], + tolerance: bool = False, + ) -> Sequence[int]: + """Find corresponding indices in df._block.column_labels for column name(s). + Order is kept the same as input names order. + + Args: + columns: column name(s) + tolerance: True to pass through columns not found. False to raise + ValueError. + """ + col_ids = self._sql_names(columns, tolerance) + return [self._block.value_columns.index(col_id) for col_id in col_ids] + + def _resolve_label_exact(self, label) -> str: + matches = self._block.label_to_col_id.get(label, []) + if len(matches) != 1: + raise ValueError("Index data must be 1-dimensional") + return matches[0] + + def _sql_names( + self, + columns: Union[blocks.Label, Sequence[blocks.Label], pd.Index], + tolerance: bool = False, + ) -> Sequence[str]: + """Retrieve sql name (column name in BQ schema) of column(s).""" + labels = columns if _is_list_like(columns) else [columns] # type:ignore + results: Sequence[str] = [] + for label in labels: + col_ids = self._block.label_to_col_id.get(label, []) + if not tolerance and len(col_ids) == 0: + raise ValueError(f"Column name {label} doesn't exist") + results = (*results, *col_ids) + return results + + @property + def index( + self, + ) -> indexes.Index: + return indexes.Index(self) + + @property + def loc(self) -> indexers.LocDataFrameIndexer: + return indexers.LocDataFrameIndexer(self) + + @property + def iloc(self) -> indexers.ILocDataFrameIndexer: + return indexers.ILocDataFrameIndexer(self) + + @property + def dtypes(self) -> pd.Series: + return pd.Series(data=self._block.dtypes, index=self._block.column_labels) + + @property + def columns(self) -> pd.Index: + return self.dtypes.index + + @property + def shape(self) -> Tuple[int, int]: + return self._block.shape + + @property + def size(self) -> int: + rows, cols = self.shape + return rows * cols + + @property + def ndim(self) -> int: + return 2 + + @property + def empty(self) -> bool: + return self.size == 0 + + @property + def values(self) -> numpy.ndarray: + return self.to_numpy() + + def __len__(self): + rows, _ = self.shape + return rows + + def astype( + self, + dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], + ) -> DataFrame: + return self._apply_to_rows(ops.AsTypeOp(dtype)) + + def to_sql_query( + self, always_include_index: bool + ) -> Tuple[str, List[Tuple[str, bool]]]: + """Compiles this dataframe's expression tree to SQL, optionally + including unnamed index columns + + Args: + always_include_index: whether to include unnamed index columns. + If False, only named indexes are included. + + Returns: a tuple of (sql_string, index_column_list). Each entry in the + index column list is a tuple of (column_name, named). If named is + is false, then the column name exists only in SQL""" + # Has to be unordered as it is impossible to order the sql without + # including metadata columns in selection with ibis. + ibis_expr = self._block.expr.to_ibis_expr(ordering_mode="unordered") + column_labels = self._block.column_labels + + # TODO(swast): Need to have a better way of controlling when to include + # the index or not. + index_has_names = all([name is not None for name in self.index.names]) + if index_has_names: + column_labels = column_labels + list(self.index.names) + elif always_include_index: + # In this mode include the index even if it is a nameless generated + # column like 'bigframes_index_0' + index_labels = [] + unnamed_index_count = 0 + for index_label in self._block.index_labels: + if index_label is None: + index_labels.append( + indexes.INDEX_COLUMN_ID.format(unnamed_index_count), + ) + unnamed_index_count += 1 + else: + index_labels.append(index_label) + + column_labels = column_labels + typing.cast( + List[Optional[str]], index_labels + ) + + column_labels_deduped = typing.cast( + List[str], + vendored_pandas_io_common.dedup_names( + column_labels, is_potential_multiindex=False + ), + ) + column_ids = self._block.value_columns + substitutions = {} + for column_id, column_label in zip(column_ids, column_labels_deduped): + # TODO(swast): Do we need to further escape this, or can we rely on + # the BigQuery unicode column name feature? + substitutions[column_id] = column_label + + index_cols: List[Tuple[str, bool]] = [] + first_index_offset = len(self._block.column_labels) + if index_has_names or always_include_index: + for i, index_col in enumerate(self._block.index_columns): + offset = first_index_offset + i + substitutions[index_col] = column_labels_deduped[offset] + index_cols = [ + (label, index_has_names) + for label in column_labels_deduped[first_index_offset:] + ] + else: + ibis_expr = ibis_expr.drop(*self._block.index_columns) + + ibis_expr = ibis_expr.relabel(substitutions) + return typing.cast(str, ibis_expr.compile()), index_cols + + @property + def sql(self) -> str: + """Compiles this dataframe's expression tree to SQL""" + sql, _ = self.to_sql_query(always_include_index=False) + return sql + + @property + def query_job(self) -> Optional[bigquery.QueryJob]: + """BigQuery job metadata for the most recent query.""" + return self._query_job + + @typing.overload + def __getitem__(self, key: bigframes.series.Series) -> DataFrame: + ... + + @typing.overload + def __getitem__(self, key: Sequence[blocks.Label]) -> DataFrame: # type:ignore + ... + + @typing.overload + def __getitem__(self, key: pd.Index) -> DataFrame: # type:ignore + ... + + @typing.overload + def __getitem__(self, key: blocks.Label) -> bigframes.series.Series: # type:ignore + ... + + def __getitem__( + self, + key: Union[ + blocks.Label, + Sequence[blocks.Label], + # Index of column labels can be treated the same as a sequence of column labels. + pd.Index, + bigframes.series.Series, + ], + ) -> Union[bigframes.series.Series, "DataFrame"]: + """Gets the specified column(s) from the DataFrame.""" + # NOTE: This implements the operations described in + # https://pandas.pydata.org/docs/getting_started/intro_tutorials/03_subset_data.html + + if isinstance(key, bigframes.series.Series): + return self._getitem_bool_series(key) + + sql_names = self._sql_names(key) + # Only input is a str and only find one column, returns a Series + if isinstance(key, str) and len(sql_names) == 1: + return bigframes.series.Series(self._block.select_column(sql_names[0])) + + # Select a subset of columns or re-order columns. + # In Ibis after you apply a projection, any column objects from the + # table before the projection can't be combined with column objects + # from the table after the projection. This is because the table after + # a projection is considered a totally separate table expression. + # + # This is unexpected behavior for a pandas user, who expects their old + # Series objects to still work with the new / mutated DataFrame. We + # avoid applying a projection in Ibis until it's absolutely necessary + # to provide pandas-like semantics. + # TODO(swast): Do we need to apply implicit join when doing a + # projection? + + # Select a number of columns as DF. + key = key if _is_list_like(key) else [key] # type:ignore + + selected_ids: Tuple[str, ...] = () + for label in key: + col_ids = self._block.label_to_col_id.get(label, []) + selected_ids = (*selected_ids, *col_ids) + + return DataFrame(self._block.select_columns(selected_ids)) + + # Bool Series selects rows + def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame: + if not key._to_ibis_expr().type() == ibis_dtypes.bool: + raise ValueError("Only boolean series currently supported for indexing.") + # TODO: enforce stricter alignment + combined_index, ( + get_column_left, + get_column_right, + ) = self._block.index.join(key._block.index, how="left") + block = combined_index._block + filter_col_id = get_column_right(key._value_column) + block = block.filter(filter_col_id) + block = block.drop_columns([filter_col_id]) + return DataFrame(block) + + def __getattr__(self, key: str): + if key in self._block.column_labels: + return self.__getitem__(key) + elif hasattr(pd.DataFrame, key): + raise NotImplementedError( + textwrap.dedent( + f""" + BigQuery DataFrames has not yet implemented an equivalent to + 'pandas.DataFrame.{key}'. Please check + https://github.com/googleapis/python-bigquery-dataframes/issues for + existing feature requests, or file your own. + Please include information about your use case, as well as + relevant code snippets. + """ + ) + ) + else: + raise AttributeError(key) + + def __repr__(self) -> str: + """Converts a DataFrame to a string. Calls compute. + + Only represents the first ``bigframes.options.display.max_rows``. + """ + opts = bigframes.options.display + max_results = opts.max_rows + # TODO(swast): pass max_columns and get the true column count back. Maybe + # get 1 more column than we have requested so that pandas can add the + # ... for us? + pandas_df, row_count = self._retrieve_repr_request_results(max_results) + column_count = len(pandas_df.columns) + + with display_options.pandas_repr(opts): + repr_string = repr(pandas_df) + + # Modify the end of the string to reflect count. + lines = repr_string.split("\n") + pattern = re.compile("\\[[0-9]+ rows x [0-9]+ columns\\]") + if pattern.match(lines[-1]): + lines = lines[:-2] + + if row_count > len(lines) - 1: + lines.append("...") + + lines.append("") + lines.append(f"[{row_count} rows x {column_count} columns]") + return "\n".join(lines) + + def _repr_html_(self) -> str: + """ + Returns an html string primarily for use by notebooks for displaying + a representation of the DataFrame. Displays 20 rows by default since + many notebooks are not configured for large tables. + """ + opts = bigframes.options.display + max_results = bigframes.options.display.max_rows + # TODO(swast): pass max_columns and get the true column count back. Maybe + # get 1 more column than we have requested so that pandas can add the + # ... for us? + pandas_df, row_count = self._retrieve_repr_request_results(max_results) + column_count = len(pandas_df.columns) + + with display_options.pandas_repr(opts): + # _repr_html_ stub is missing so mypy thinks it's a Series. Ignore mypy. + html_string = pandas_df._repr_html_() # type:ignore + + html_string += f"[{row_count} rows x {column_count} columns in total]" + return html_string + + def _retrieve_repr_request_results( + self, max_results: int + ) -> Tuple[pd.DataFrame, int]: + """ + Retrieves a pandas dataframe containing only max_results many rows for use + with printing methods. + + Returns a tuple of the dataframe and the overall number of rows of the query. + """ + # TODO(swast): Select a subset of columns if max_columns is less than the + # number of columns in the schema. + count = self.shape[0] + if count > max_results: + head_df = self.head(n=max_results) + computed_df, query_job = head_df._block.compute(max_results=max_results) + else: + head_df = self + computed_df, query_job = head_df._block.compute() + + formatted_df = computed_df.set_axis(self._block.column_labels, axis=1) + # don't update details when the cache is hit + if self.query_job is None or not query_job.cache_hit: + self._query_job = query_job + # we reset the axis and substitute the bf index name for the default + formatted_df.index.name = self.index.name + return formatted_df, count + + def __setitem__(self, key: str, value: SingleItemValue): + """Modify or insert a column into the DataFrame. + + Note: This does **not** modify the original table the DataFrame was + derived from. + """ + df = self._assign_single_item(key, value) + self._set_block(df._get_block()) + + def _apply_binop( + self, + other: float | int | bigframes.series.Series, + op, + axis: str | int = "columns", + ): + if isinstance(other, (float, int)): + return self._apply_scalar_binop(other, op) + elif isinstance(other, bigframes.series.Series): + return self._apply_series_binop(other, op, axis=axis) + raise NotImplementedError( + f"binary operation is not implemented on the second operand of type {type(other).__name__}." + ) + + def _apply_scalar_binop(self, other: float | int, op: ops.BinaryOp) -> DataFrame: + block = self._block + partial_op = ops.BinopPartialRight(op, other) + for column_id, label in zip( + self._block.value_columns, self._block.column_labels + ): + block, _ = block.apply_unary_op(column_id, partial_op, result_label=label) + block = block.drop_columns([column_id]) + return DataFrame(block) + + def _apply_series_binop( + self, + other: bigframes.series.Series, + op: ops.BinaryOp, + axis: str | int = "columns", + ) -> DataFrame: + if axis not in ("columns", "index", 0, 1): + raise ValueError(f"Invalid input: axis {axis}.") + + if axis in ("columns", 1): + raise NotImplementedError("Row Series operations haven't been supported.") + + joined_index, (get_column_left, get_column_right) = self._block.index.join( + other._block.index, how="outer" + ) + + series_column_id = other._value.get_name() + series_col = get_column_right(series_column_id) + block = joined_index._block + for column_id, label in zip( + self._block.value_columns, self._block.column_labels + ): + block, _ = block.apply_binary_op( + get_column_left(column_id), + series_col, + op, + result_label=label, + ) + block = block.drop_columns([get_column_left(column_id)]) + + block = block.drop_columns([series_col]) + block = block.with_index_labels(self.index.names) + return DataFrame(block) + + def le(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.le_op, axis=axis) + + def lt(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.lt_op, axis=axis) + + def ge(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.ge_op, axis=axis) + + def gt(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.gt_op, axis=axis) + + __lt__ = lt + + __le__ = le + + __gt__ = gt + + __ge__ = ge + + def add( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + # TODO(swast): Support fill_value parameter. + # TODO(swast): Support level parameter with MultiIndex. + return self._apply_binop(other, ops.add_op, axis=axis) + + __radd__ = __add__ = radd = add + + def sub( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.sub_op, axis=axis) + + __sub__ = subtract = sub + + def rsub( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.reverse(ops.sub_op), axis=axis) + + __rsub__ = rsub + + def mul( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.mul_op, axis=axis) + + __rmul__ = __mul__ = rmul = multiply = mul + + def truediv( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.div_op, axis=axis) + + div = divide = __truediv__ = truediv + + def rtruediv( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.reverse(ops.div_op), axis=axis) + + __rtruediv__ = rdiv = rtruediv + + def floordiv( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.floordiv_op, axis=axis) + + __floordiv__ = floordiv + + def rfloordiv( + self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.reverse(ops.floordiv_op), axis=axis) + + __rfloordiv__ = rfloordiv + + def mod(self, other: int | bigframes.series.Series, axis: str | int = "columns") -> DataFrame: # type: ignore + return self._apply_binop(other, ops.mod_op, axis=axis) + + def rmod(self, other: int | bigframes.series.Series, axis: str | int = "columns") -> DataFrame: # type: ignore + return self._apply_binop(other, ops.reverse(ops.mod_op), axis=axis) + + __mod__ = mod + + __rmod__ = rmod + + def compute(self) -> pd.DataFrame: + """Executes deferred operations and downloads the results.""" + # TODO(orrbradford): Optimize this in future. Potentially some cases where we can return the stored query job + df, query_job = self._block.compute() + self._query_job = query_job + return df.set_axis(self._block.column_labels, axis=1) + + def copy(self) -> DataFrame: + return DataFrame(self._block) + + def head(self, n: int = 5) -> DataFrame: + return typing.cast(DataFrame, self.iloc[:n]) + + def tail(self, n: int = 5) -> DataFrame: + return typing.cast(DataFrame, self.iloc[-n:]) + + def drop(self, *, columns: Union[str, Iterable[str]]) -> DataFrame: + if not _is_list_like(columns): + columns = [columns] # type:ignore + columns = list(columns) + + block = self._block.drop_columns(self._sql_names(columns)) + return DataFrame(block) + + def droplevel(self, level: LevelsType): + resolved_level_ids = self._resolve_levels(level) + return DataFrame(self._block.drop_levels(resolved_level_ids)) + + def reorder_levels(self, order: LevelsType): + resolved_level_ids = self._resolve_levels(order) + return DataFrame(self._block.reorder_levels(resolved_level_ids)) + + def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: + if _is_list_like(level): + levels = list(level) + else: + levels = [level] + resolved_level_ids = [] + for level_ref in levels: + if isinstance(level_ref, int): + resolved_level_ids.append(self._block.index_columns[level_ref]) + elif isinstance(level_ref, str): + matching_ids = self._block.index_name_to_col_id.get(level_ref, []) + if len(matching_ids) != 1: + raise ValueError("level name cannot be found or is ambiguous") + resolved_level_ids.append(matching_ids[0]) + else: + raise ValueError(f"Unexpected level: {level_ref}") + return resolved_level_ids + + def rename(self, *, columns: Mapping[blocks.Label, blocks.Label]) -> DataFrame: + block = self._block.rename(columns=columns) + return DataFrame(block) + + def rename_axis( + self, + mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]], + **kwargs, + ) -> DataFrame: + if len(kwargs) != 0: + raise NotImplementedError( + "rename_axis does not currently support any keyword arguments." + ) + # limited implementation: the new index name is simply the 'mapper' parameter + if _is_list_like(mapper): + labels = mapper + else: + labels = [mapper] + return DataFrame(self._block.with_index_labels(labels)) + + def assign(self, **kwargs) -> DataFrame: + # TODO(garrettwu) Support list-like values. Requires ordering. + # TODO(garrettwu) Support callable values. + + cur = self + for k, v in kwargs.items(): + cur = cur._assign_single_item(k, v) + + return cur + + def _assign_single_item( + self, + k: str, + v: SingleItemValue, + ) -> DataFrame: + if isinstance(v, bigframes.series.Series): + return self._assign_series_join_on_index(k, v) + elif callable(v): + copy = self.copy() + copy[k] = v(copy) + return copy + else: + return self._assign_scalar(k, v) + + def _assign_scalar(self, label: str, value: Union[int, float]) -> DataFrame: + # TODO(swast): Make sure that k is the ID / SQL name, not a label, + # which could be invalid SQL. + col_ids = self._sql_names(label, tolerance=True) + + block, constant_col_id = self._block.create_constant(value, label) + for col_id in col_ids: + block = block.copy_values(constant_col_id, col_id) + + if len(col_ids) > 0: + block = block.drop_columns([constant_col_id]) + + return DataFrame(block) + + def _assign_series_join_on_index( + self, label: str, series: bigframes.series.Series + ) -> DataFrame: + joined_index, (get_column_left, get_column_right) = self._block.index.join( + series._block.index, how="left" + ) + + column_ids = [ + get_column_left(col_id) for col_id in self._sql_names(label, tolerance=True) + ] + block = joined_index._block + source_column = get_column_right(series._value_column) + + # Replace each column matching the label + for column_id in column_ids: + block = block.copy_values(source_column, column_id).assign_label( + column_id, label + ) + + if not column_ids: + # Append case, so new column needs appropriate label + block = block.assign_label(source_column, label) + else: + # Update case, remove after copying into columns + block = block.drop_columns([source_column]) + + return DataFrame(block.with_index_labels([self.index.name])) + + def reset_index(self, *, drop: bool = False) -> DataFrame: + block = self._block.reset_index(drop) + return DataFrame(block) + + def set_index( + self, + keys: typing.Union[blocks.Label, typing.Sequence[blocks.Label]], + append: bool = False, + drop: bool = True, + ) -> DataFrame: + if not _is_list_like(keys): + keys = typing.cast(typing.Sequence[blocks.Label], (keys,)) + else: + keys = typing.cast(typing.Sequence[blocks.Label], tuple(keys)) + col_ids = [self._resolve_label_exact(key) for key in keys] + return DataFrame(self._block.set_index(col_ids, append=append, drop=drop)) + + def sort_index(self) -> DataFrame: + index_columns = self._block.index_columns + ordering = [order.OrderingColumnReference(column) for column in index_columns] + return DataFrame(self._block.order_by(ordering)) + + def sort_values( + self, + by: str | typing.Sequence[str], + *, + ascending: bool | typing.Sequence[bool] = True, + na_position: typing.Literal["first", "last"] = "last", + ) -> DataFrame: + if na_position not in {"first", "last"}: + raise ValueError("Param na_position must be one of 'first' or 'last'") + + sort_labels = (by,) if isinstance(by, str) else tuple(by) + sort_column_ids = self._sql_names(sort_labels) + + len_by = len(sort_labels) + if not isinstance(ascending, bool): + if len(ascending) != len_by: + raise ValueError("Length of 'ascending' must equal length of 'by'") + sort_directions = ascending + else: + sort_directions = (ascending,) * len_by + + ordering = [] + for i in range(len(sort_labels)): + column_id = sort_column_ids[i] + direction = ( + order.OrderingDirection.ASC + if sort_directions[i] + else order.OrderingDirection.DESC + ) + na_last = na_position == "last" + ordering.append( + order.OrderingColumnReference( + column_id, direction=direction, na_last=na_last + ) + ) + + return DataFrame(self._block.order_by(ordering)) + + def value_counts( + self, + subset: typing.Union[blocks.Label, typing.Sequence[blocks.Label]] = None, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + dropna: bool = True, + ): + # 'sort'=False allows arbitrary sorting, so we will sort anyways and ignore the param + columns = self._sql_names(subset) if subset else self._block.value_columns + block = block_ops.value_counts( + self._block, + columns, + normalize=normalize, + sort=sort, + ascending=ascending, + dropna=dropna, + ) + return bigframes.series.Series(block) + + def add_prefix(self, prefix: str, axis: int | str | None = None) -> DataFrame: + return DataFrame(self._get_block().add_prefix(prefix, axis)) + + def add_suffix(self, suffix: str, axis: int | str | None = None) -> DataFrame: + return DataFrame(self._get_block().add_suffix(suffix, axis)) + + def dropna(self) -> DataFrame: + block = self._block + for column in self._block.value_columns: + block, result_id = block.apply_unary_op(column, ops.notnull_op) + block = block.filter(result_id) + block = block.drop_columns([result_id]) + + return DataFrame(block) + + def any( + self, + *, + bool_only: bool = False, + ) -> bigframes.series.Series: + if not bool_only: + frame = self._raise_on_non_boolean("any") + else: + frame = self._drop_non_bool() + block = frame._block.aggregate_all_and_pivot( + agg_ops.any_op, dtype=pd.BooleanDtype() + ) + return bigframes.series.Series(block.select_column("values")) + + def all(self, *, bool_only: bool = False) -> bigframes.series.Series: + if not bool_only: + frame = self._raise_on_non_boolean("all") + else: + frame = self._drop_non_bool() + block = frame._block.aggregate_all_and_pivot( + agg_ops.all_op, dtype=pd.BooleanDtype() + ) + return bigframes.series.Series(block.select_column("values")) + + def sum(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("sum") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.sum_op) + return bigframes.series.Series(block.select_column("values")) + + def mean(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("mean") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.mean_op) + return bigframes.series.Series(block.select_column("values")) + + def std(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("std") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.std_op) + return bigframes.series.Series(block.select_column("values")) + + def var(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("var") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.var_op) + return bigframes.series.Series(block.select_column("values")) + + def min(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("min") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.min_op) + return bigframes.series.Series(block.select_column("values")) + + def max(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("max") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.max_op) + return bigframes.series.Series(block.select_column("values")) + + def prod(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self._raise_on_non_numeric("prod") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.product_op) + return bigframes.series.Series(block.select_column("values")) + + product = prod + + def count(self, *, numeric_only: bool = False) -> bigframes.series.Series: + if not numeric_only: + frame = self + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.count_op) + return bigframes.series.Series(block.select_column("values")) + + def nunique(self) -> bigframes.series.Series: + block = self._block.aggregate_all_and_pivot(agg_ops.nunique_op) + return bigframes.series.Series(block.select_column("values")) + + def _drop_non_numeric(self) -> DataFrame: + non_numeric_cols = [ + col_id + for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) + if dtype not in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES + ] + return DataFrame(self._block.drop_columns(non_numeric_cols)) + + def _drop_non_bool(self) -> DataFrame: + non_bool_cols = [ + col_id + for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) + if dtype not in bigframes.dtypes.BOOL_BIGFRAMES_TYPES + ] + return DataFrame(self._block.drop_columns(non_bool_cols)) + + def _raise_on_non_numeric(self, op: str): + if not all( + dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES + for dtype in self._block.dtypes + ): + raise NotImplementedError( + f"'{op}' does not support non-numeric columns. Set 'numeric_only'=True to ignore non-numeric columns" + ) + return self + + def _raise_on_non_boolean(self, op: str): + if not all( + dtype in bigframes.dtypes.BOOL_BIGFRAMES_TYPES + for dtype in self._block.dtypes + ): + raise NotImplementedError( + f"'{op}' does not support non-bool columns. Set 'bool_only'=True to ignore non-bool columns" + ) + return self + + def merge( + self, + right: DataFrame, + how: Literal[ + "inner", + "left", + "outer", + "right", + ] = "inner", + # TODO(garrettwu): Currently can take inner, outer, left and right. To support + # cross joins + # TODO(garrettwu): Support "on" list of columns and None. Currently a single + # column must be provided + on: Optional[str] = None, + *, + left_on: Optional[str] = None, + right_on: Optional[str] = None, + sort: bool = False, + suffixes: tuple[str, str] = ("_x", "_y"), + ) -> DataFrame: + if on is None: + if left_on is None or right_on is None: + raise ValueError("Must specify either on or left_on + right_on.") + else: + if left_on is not None or right_on is not None: + raise ValueError("Can not pass both on and left_on + right_on params.") + left_on, right_on = on, on + + left = self + left_on_sql = self._sql_names(left_on) + # 0 elements already throws an exception + if len(left_on_sql) > 1: + raise ValueError(f"The column label {left_on} is not unique.") + left_on_sql = left_on_sql[0] + + right_on_sql = right._sql_names(right_on) + if len(right_on_sql) > 1: + raise ValueError(f"The column label {right_on} is not unique.") + right_on_sql = right_on_sql[0] + + ( + joined_expr, + join_key_ids, + (get_column_left, get_column_right), + ) = joins.join_by_column( + left._block.expr, + [left_on_sql], + right._block.expr, + [right_on_sql], + how=how, + sort=sort, + # In merging on the same column, it only returns 1 key column from coalesced both. + # While if 2 different columns, both will be presented in the result. + get_both_join_key_cols=(left_on != right_on), + ) + # TODO(swast): Add suffixes to the column labels instead of reusing the + # column IDs as the new labels. + # Drop the index column(s) to be consistent with pandas. + left_columns = [ + join_key_ids[0] if (col_id == left_on_sql) else get_column_left(col_id) + for col_id in left._block.value_columns + ] + + right_columns = [] + for col_id in right._block.value_columns: + if col_id == right_on_sql: + # When left_on == right_on + if len(join_key_ids) > 1: + right_columns.append(join_key_ids[1]) + else: + right_columns.append(get_column_right(col_id)) + + expr = joined_expr.select_columns([*left_columns, *right_columns]) + labels = self._get_merged_col_labels( + right, left_on=left_on, right_on=right_on, suffixes=suffixes + ) + + # Constructs default index + block = blocks.Block(expr, column_labels=labels) + return DataFrame(block) + + def _get_merged_col_labels( + self, + right: DataFrame, + left_on: str, + right_on: str, + suffixes: tuple[str, str] = ("_x", "_y"), + ) -> List[blocks.Label]: + on_col_equal = left_on == right_on + + left_col_labels: list[blocks.Label] = [] + for col_label in self._block.column_labels: + if col_label in right._block.column_labels: + if on_col_equal and col_label == left_on: + # Merging on the same column only returns 1 key column from coalesce both. + # Take the left key column. + left_col_labels.append(col_label) + else: + left_col_labels.append(str(col_label) + suffixes[0]) + else: + left_col_labels.append(col_label) + + right_col_labels: list[blocks.Label] = [] + for col_label in right._block.column_labels: + if col_label in self._block.column_labels: + if on_col_equal and col_label == left_on: + # Merging on the same column only returns 1 key column from coalesce both. + # Pass the right key column. + pass + else: + right_col_labels.append(str(col_label) + suffixes[1]) + else: + right_col_labels.append(col_label) + + return left_col_labels + right_col_labels + + def join(self, other: DataFrame, *, how: str = "left") -> DataFrame: + if not self.columns.intersection(other.columns).empty: + raise NotImplementedError("Deduping column names is not implemented") + + left = self + right = other + combined_index, (get_column_left, get_column_right) = left._block.index.join( + right._block.index, how=how + ) + return DataFrame(combined_index._block) + + def groupby( + self, + by: typing.Union[ + blocks.Label, + bigframes.series.Series, + typing.Sequence[typing.Union[blocks.Label, bigframes.series.Series]], + ] = None, + *, + level: typing.Optional[LevelsType] = None, + as_index: bool = True, + dropna: bool = True, + ) -> groupby.DataFrameGroupBy: + if (by is not None) and (level is not None): + raise ValueError("Do not specify both 'by' and 'level'") + if by is not None: + return self._groupby_series(by, as_index=as_index, dropna=dropna) + if level is not None: + return self._groupby_level(level, as_index=as_index, dropna=dropna) + else: + raise TypeError("You have to supply one of 'by' and 'level'") + + def _groupby_level( + self, + level: LevelsType, + as_index: bool = True, + dropna: bool = True, + ): + return groupby.DataFrameGroupBy( + self._block, + self._resolve_levels(level), + as_index=as_index, + dropna=dropna, + ) + + def _groupby_series( + self, + by: typing.Union[ + blocks.Label, + bigframes.series.Series, + typing.Sequence[typing.Union[blocks.Label, bigframes.series.Series]], + ], + as_index: bool = True, + dropna: bool = True, + ): + if not isinstance(by, bigframes.series.Series) and _is_list_like(by): + by = list(by) + else: + by = [typing.cast(typing.Union[blocks.Label, bigframes.series.Series], by)] + + block = self._block + col_ids: typing.Sequence[str] = [] + for key in by: + if isinstance(key, bigframes.series.Series): + combined_index, ( + get_column_left, + get_column_right, + ) = block.index.join( + key._block.index, how="inner" if dropna else "left" + ) + col_ids = [ + *[get_column_left(value) for value in col_ids], + get_column_right(key._value_column), + ] + block = combined_index._block + else: + # Interpret as index level or column name + col_matches = block.label_to_col_id.get(key, []) + level_matches = block.index_name_to_col_id.get(key, []) + matches = [*col_matches, *level_matches] + if len(matches) != 1: + raise ValueError( + f"GroupBy key {key} does not map to unambiguous column or index level" + ) + col_ids = [*col_ids, matches[0]] + + return groupby.DataFrameGroupBy( + block, + col_ids, + as_index=as_index, + dropna=dropna, + ) + + def abs(self) -> DataFrame: + return self._apply_to_rows(ops.abs_op) + + def isna(self) -> DataFrame: + return self._apply_to_rows(ops.isnull_op) + + isnull = isna + + def notna(self) -> DataFrame: + return self._apply_to_rows(ops.notnull_op) + + notnull = notna + + def cumsum(self): + is_numeric_types = [ + (dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) + for _, dtype in self.dtypes.items() + ] + if not all(is_numeric_types): + raise ValueError("All values must be numeric to apply cumsum.") + return self._apply_window_op( + agg_ops.sum_op, + bigframes.core.WindowSpec(following=0), + ) + + def cumprod(self) -> DataFrame: + is_numeric_types = [ + (dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) + for _, dtype in self.dtypes.items() + ] + if not all(is_numeric_types): + raise ValueError("All values must be numeric to apply cumsum.") + return self._apply_window_op( + agg_ops.product_op, + bigframes.core.WindowSpec(following=0), + ) + + def cummin(self) -> DataFrame: + return self._apply_window_op( + agg_ops.min_op, + bigframes.core.WindowSpec(following=0), + ) + + def cummax(self) -> DataFrame: + return self._apply_window_op( + agg_ops.max_op, + bigframes.core.WindowSpec(following=0), + ) + + def shift(self, periods: int = 1) -> DataFrame: + window = bigframes.core.WindowSpec( + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.ShiftOp(periods), window) + + def _apply_window_op( + self, + op: agg_ops.WindowOp, + window_spec: bigframes.core.WindowSpec, + ): + block = self._block.multi_apply_window_op( + self._block.value_columns, + op, + window_spec=window_spec, + ) + return DataFrame(block) + + def sample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + *, + random_state: Optional[int] = None, + ) -> DataFrame: + if n is not None and frac is not None: + raise ValueError("Only one of 'n' or 'frac' parameter can be specified.") + + ns = (n,) if n is not None else () + fracs = (frac,) if frac is not None else () + + return self._split(ns=ns, fracs=fracs, random_state=random_state)[0] + + def _split( + self, + ns: Iterable[int] = (), + fracs: Iterable[float] = (), + *, + random_state: Optional[int] = None, + ) -> List[DataFrame]: + """Internal function to support splitting DF to multiple parts along index axis. + + At most one of ns and fracs can be passed in. If neither, default to ns = (1,). + Return a list of sampled DataFrames. + """ + if ns and fracs: + raise ValueError("Only one of 'ns' or 'fracs' parameter must be specified.") + + block = self._block + if not ns and not fracs: + ns = (1,) + + if ns: + sample_sizes = ns + else: + total_rows = block.shape[0] + # Round to nearest integer. "round half to even" rule applies. + # At least to be 1. + sample_sizes = [round(frac * total_rows) or 1 for frac in fracs] + + # Set random_state if it is not provided + if random_state is None: + random_state = random.randint(-(2**30), 2**30) + + # Create a new column with random_state value. + block, random_state_col = block.create_constant(random_state) + + # Create an ordering col and a new sum col which is ordering+random_state. + block, ordering_col = block.promote_offsets() + block, sum_col = block.apply_binary_op( + ordering_col, random_state_col, ops.add_op + ) + + # Apply hash method to sum col and order by it. + block, string_sum_col = block.apply_unary_op( + sum_col, ops.AsTypeOp("string[pyarrow]") + ) + block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op) + block = block.order_by([order.OrderingColumnReference(hash_string_sum_col)]) + + drop_cols = [ + random_state_col, + ordering_col, + sum_col, + string_sum_col, + hash_string_sum_col, + ] + block = block.drop_columns(drop_cols) + df = DataFrame(block) + + intervals = [] + cur = 0 + for sample_size in sample_sizes: + intervals.append((cur, cur + sample_size)) + cur += sample_size + + # DF.iloc[slice] always returns DF. + return [ + typing.cast(DataFrame, df.iloc[lower:upper]) for lower, upper in intervals + ] + + def to_pandas(self) -> pd.DataFrame: + """Writes DataFrame to Pandas DataFrame.""" + # TODO(chelsealin): Support block parameters. + # TODO(chelsealin): Add to_pandas_batches() API. + return self.compute() + + def to_csv(self, path_or_buf: str, *, index: bool = True) -> None: + # TODO(swast): Can we support partition columns argument? + # TODO(chelsealin): Support local file paths. + # TODO(swast): Some warning that wildcard is recommended for large + # query results? See: + # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size + if not path_or_buf.startswith("gs://"): + raise NotImplementedError( + "Only Google Cloud Storage (gs://...) paths are supported." + ) + + source_table = self._execute_query(index=index) + job_config = bigquery.ExtractJobConfig( + destination_format=bigquery.DestinationFormat.CSV + ) + self._block.expr._session._extract_table( + source_table, + destination_uris=[path_or_buf], + job_config=job_config, + ) + + def to_json( + self, + path_or_buf: str, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + *, + lines: bool = False, + index: bool = True, + ) -> None: + # TODO(swast): Can we support partition columns argument? + # TODO(chelsealin): Support local file paths. + # TODO(swast): Some warning that wildcard is recommended for large + # query results? See: + # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size + if not path_or_buf.startswith("gs://"): + raise NotImplementedError( + "Only Google Cloud Storage (gs://...) paths are supported." + ) + + if lines is True and orient != "records": + raise ValueError( + "'lines' keyword is only valid when 'orient' is 'records'." + ) + + # TODO(ashleyxu) Support lines=False for small tables with arrays and TO_JSON_STRING. + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#to_json_string + if lines is False: + raise NotImplementedError( + "Only newline delimited JSON format is supported." + ) + + source_table = self._execute_query(index=index) + job_config = bigquery.ExtractJobConfig( + destination_format=bigquery.DestinationFormat.NEWLINE_DELIMITED_JSON + ) + self._block.expr._session._extract_table( + source_table, + destination_uris=[path_or_buf], + job_config=job_config, + ) + + def to_gbq( + self, + destination_table: str, + *, + if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", + index: bool = True, + ) -> None: + if "." not in destination_table: + raise ValueError( + "Invalid Table Name. Should be of the form 'datasetId.tableId' or " + "'projectId.datasetId.tableId'" + ) + + dispositions = { + "fail": bigquery.WriteDisposition.WRITE_EMPTY, + "replace": bigquery.WriteDisposition.WRITE_TRUNCATE, + "append": bigquery.WriteDisposition.WRITE_APPEND, + } + if if_exists not in dispositions: + raise ValueError("'{0}' is not valid for if_exists".format(if_exists)) + + job_config = bigquery.QueryJobConfig( + write_disposition=dispositions[if_exists], + destination=bigquery.table.TableReference.from_string( + destination_table, + default_project=self._block.expr._session.bqclient.project, + ), + ) + + self._execute_query(index=index, job_config=job_config) + + def to_numpy( + self, dtype=None, copy=False, na_value=None, **kwargs + ) -> numpy.ndarray: + return self.compute().to_numpy(dtype, copy, na_value, **kwargs) + + __array__ = to_numpy + + def to_parquet(self, path: str, *, index: bool = True) -> None: + # TODO(swast): Can we support partition columns argument? + # TODO(chelsealin): Support local file paths. + # TODO(swast): Some warning that wildcard is recommended for large + # query results? See: + # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size + if not path.startswith("gs://"): + raise NotImplementedError( + "Only Google Cloud Storage (gs://...) paths are supported." + ) + + source_table = self._execute_query(index=index) + job_config = bigquery.ExtractJobConfig( + destination_format=bigquery.DestinationFormat.PARQUET + ) + + self._block.expr._session._extract_table( + source_table, + destination_uris=[path], + job_config=job_config, + ) + + def _apply_to_rows(self, operation: ops.UnaryOp): + block = self._block.multi_apply_unary_op(self._block.value_columns, operation) + return DataFrame(block) + + def _execute_query( + self, index: bool, job_config: Optional[bigquery.job.QueryJobConfig] = None + ): + """Executes a query job presenting this dataframe and returns the destination + table.""" + expr = self._block.expr + session = expr._session + columns = list(self._block.value_columns) + column_labels = list(self._block.column_labels) + # This code drops unnamed indexes to keep consistent with the behavior of + # most pandas write APIs. The exception is `pandas.to_csv`, which keeps + # unnamed indexes as `Unnamed: 0`. + # TODO(chelsealin): check if works for multiple indexes. + if index and self.index.name is not None: + columns.extend(self._block.index_columns) + column_labels.extend(self.index.names) + # TODO(chelsealin): normalize the file formats if we needs, such as arbitrary + # unicode for column labels. + value_columns = (expr.get_column(column_name) for column_name in columns) + expr = expr.projection(value_columns) + + # Make columns in SQL reflect _labels_ not _ids_. Note: This may use + # the arbitrary unicode column labels feature in BigQuery, which is + # currently (June 2023) in preview. + # TODO(swast): Handle duplicate and NULL labels. + ibis_expr = expr.to_ibis_expr() + renamed_columns = [ + ibis_expr[col_id].name(col_label) + for col_id, col_label in zip(columns, column_labels) + ] + ibis_expr = ibis_expr.select(*renamed_columns) + sql = session.ibis_client.compile(ibis_expr) # type: ignore + _, query_job = session._start_query( + sql=sql, job_config=job_config # type: ignore + ) + return query_job.destination + + def map(self, func, na_action: Optional[str] = None) -> DataFrame: + if not callable(func): + raise TypeError("the first argument must be callable") + + if na_action not in {None, "ignore"}: + raise ValueError(f"na_action={na_action} not supported") + + # TODO(shobs): Support **kwargs + + return self._apply_to_rows( + ops.RemoteFunctionOp(func, apply_on_null=(na_action is None)) + ) + + def drop_duplicates(self, subset=None, *, keep: str = "first") -> DataFrame: + if subset is None: + column_ids = self._block.value_columns + else: + column_ids = [ + id for label in subset for id in self._block.label_to_col_id[label] + ] + block = block_ops.drop_duplicates(self._block, column_ids, keep) + return DataFrame(block) + + def duplicated(self, subset=None, keep: str = "first") -> bigframes.series.Series: + if subset is None: + column_ids = self._block.value_columns + else: + column_ids = [ + id for label in subset for id in self._block.label_to_col_id[label] + ] + block, indicator = block_ops.indicate_duplicates(self._block, column_ids, keep) + return bigframes.series.Series( + block.select_column( + indicator, + ) + ) + + def rank( + self, + axis=0, + method: str = "average", + numeric_only=False, + na_option: str = "keep", + ascending=True, + ) -> DataFrame: + df = self._drop_non_numeric() if numeric_only else self + return DataFrame(block_ops.rank(df._block, method, na_option, ascending)) + + applymap = map + + def _slice( + self, + start: typing.Optional[int] = None, + stop: typing.Optional[int] = None, + step: typing.Optional[int] = None, + ) -> DataFrame: + block = self._block.slice(start=start, stop=stop, step=step) + return DataFrame(block) + + def _set_block(self, block: blocks.Block): + self._block = block + + def _get_block(self) -> blocks.Block: + return self._block + + +def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: + return pd.api.types.is_list_like(obj) + + +def _is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]: + return pd.api.types.is_dict_like(obj) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py new file mode 100644 index 00000000000..c257d6943eb --- /dev/null +++ b/bigframes/dtypes.py @@ -0,0 +1,304 @@ +# 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. + +"""Mappings for Pandas dtypes supported by BigQuery DataFrames package""" + +import typing +from typing import Any, Dict, Iterable, Literal, Tuple, Union + +import geopandas as gpd # type: ignore +import ibis +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types +import numpy as np +import pandas as pd +import pyarrow as pa + +# Type hints for Pandas dtypes supported by BigQuery DataFrame +Dtype = Union[ + pd.BooleanDtype, + pd.Float64Dtype, + pd.Int64Dtype, + pd.StringDtype, + pd.ArrowDtype, +] + +# Corresponds to the pandas concept of numeric type (such as when 'numeric_only' is specified in an operation) +NUMERIC_BIGFRAMES_TYPES = [pd.BooleanDtype(), pd.Float64Dtype(), pd.Int64Dtype()] + +# Type hints for dtype strings supported by BigQuery DataFrame +DtypeString = Literal[ + "boolean", + "Float64", + "Int64", + "string", + "string[pyarrow]", + "timestamp[us, tz=UTC][pyarrow]", + "timestamp[us][pyarrow]", + "date32[day][pyarrow]", + "time64[us][pyarrow]", +] + +# Type hints for Ibis data types supported by BigQuery DataFrame +IbisDtype = Union[ + ibis_dtypes.Boolean, + ibis_dtypes.Float64, + ibis_dtypes.Int64, + ibis_dtypes.String, + ibis_dtypes.Date, + ibis_dtypes.Time, + ibis_dtypes.Timestamp, +] + +BOOL_BIGFRAMES_TYPES = [pd.BooleanDtype()] + +# Several operations are restricted to these types. +NUMERIC_BIGFRAMES_TYPES = [pd.BooleanDtype(), pd.Float64Dtype(), pd.Int64Dtype()] + +# Type hints for Ibis data types that can be read to Python objects by BigQuery DataFrame +ReadOnlyIbisDtype = Union[ + ibis_dtypes.Binary, + ibis_dtypes.JSON, + ibis_dtypes.Decimal, + ibis_dtypes.GeoSpatial, + ibis_dtypes.Array, + ibis_dtypes.Struct, +] + +BIDIRECTIONAL_MAPPINGS: Iterable[Tuple[IbisDtype, Dtype]] = ( + (ibis_dtypes.boolean, pd.BooleanDtype()), + (ibis_dtypes.float64, pd.Float64Dtype()), + (ibis_dtypes.int64, pd.Int64Dtype()), + (ibis_dtypes.string, pd.StringDtype(storage="pyarrow")), + (ibis_dtypes.date, pd.ArrowDtype(pa.date32())), + (ibis_dtypes.time, pd.ArrowDtype(pa.time64("us"))), + (ibis_dtypes.Timestamp(timezone=None), pd.ArrowDtype(pa.timestamp("us"))), + ( + ibis_dtypes.Timestamp(timezone="UTC"), + pd.ArrowDtype(pa.timestamp("us", tz="UTC")), + ), +) + +BIGFRAMES_TO_IBIS: Dict[Dtype, IbisDtype] = { + pandas: ibis for ibis, pandas in BIDIRECTIONAL_MAPPINGS +} + +IBIS_TO_BIGFRAMES: Dict[ + Union[IbisDtype, ReadOnlyIbisDtype], Union[Dtype, np.dtype[Any]] +] = {ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS} +# Allow REQUIRED fields to map correctly. +IBIS_TO_BIGFRAMES.update( + {ibis.copy(nullable=False): pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS} +) +IBIS_TO_BIGFRAMES.update( + { + ibis_dtypes.binary: np.dtype("O"), + ibis_dtypes.json: np.dtype("O"), + ibis_dtypes.Decimal(precision=38, scale=9, nullable=True): np.dtype("O"), + ibis_dtypes.Decimal(precision=76, scale=38, nullable=True): np.dtype("O"), + ibis_dtypes.GeoSpatial( + geotype="geography", srid=4326, nullable=True + ): gpd.array.GeometryDtype(), + # TODO: Interval + } +) + +BIGFRAMES_STRING_TO_BIGFRAMES: Dict[DtypeString, Dtype] = { + typing.cast(DtypeString, dtype.name): dtype for dtype in BIGFRAMES_TO_IBIS.keys() +} + +# special case - string[pyarrow] doesn't include the storage in its name, and both +# "string" and "string[pyarrow] are accepted" +BIGFRAMES_STRING_TO_BIGFRAMES["string[pyarrow]"] = pd.StringDtype(storage="pyarrow") + + +def ibis_dtype_to_bigframes_dtype( + ibis_dtype: Union[IbisDtype, ReadOnlyIbisDtype] +) -> Union[Dtype, np.dtype[Any]]: + """Converts an Ibis dtype to a BigQuery DataFrames dtype + + Args: + ibis_dtype: The ibis dtype used to represent this type, which + should in turn correspond to an underlying BigQuery type + + Returns: + The supported BigQuery DataFrames dtype, which may be provided by + pandas, numpy, or db_types + + Raises: + ValueError: if passed an unexpected type + """ + # Special cases: Ibis supports variations on these types, but currently + # our IO returns them as objects. Eventually, we should support them as + # ArrowDType (and update the IO accordingly) + if isinstance(ibis_dtype, ibis_dtypes.Array) or isinstance( + ibis_dtype, ibis_dtypes.Struct + ): + return np.dtype("O") + + if ibis_dtype in IBIS_TO_BIGFRAMES: + return IBIS_TO_BIGFRAMES[ibis_dtype] + else: + raise ValueError(f"Unexpected Ibis data type {type(ibis_dtype)}") + + +def ibis_value_to_canonical_type(value: ibis_types.Value) -> ibis_types.Value: + """Converts an Ibis expression to canonical type. + + This is useful in cases where multiple types correspond to the same BigFrames dtype. + """ + ibis_type = value.type() + # Allow REQUIRED fields to be joined with NULLABLE fields. + nullable_type = ibis_type.copy(nullable=True) + return value.cast(nullable_type).name(value.get_name()) + + +def ibis_table_to_canonical_types(table: ibis_types.Table) -> ibis_types.Table: + """Converts an Ibis table expression to canonical types. + + This is useful in cases where multiple types correspond to the same BigFrames dtype. + """ + casted_columns = [] + for column_name in table.columns: + column = typing.cast(ibis_types.Value, table[column_name]) + casted_columns.append(ibis_value_to_canonical_type(column)) + return table.select(*casted_columns) + + +def bigframes_dtype_to_ibis_dtype( + bigframes_dtype: Union[DtypeString, Dtype] +) -> IbisDtype: + """Converts a BigQuery DataFrames supported dtype to an Ibis dtype. + + Args: + bigframes_dtype: A dtype supported by BigQuery DataFrame + + Returns: + The corresponding Ibis type + + Raises: + ValueError: + If passed a dtype not supported by BigQuery DataFrames. + """ + type_string = str(bigframes_dtype) + if type_string in BIGFRAMES_STRING_TO_BIGFRAMES: + bigframes_dtype = BIGFRAMES_STRING_TO_BIGFRAMES[ + typing.cast(DtypeString, type_string) + ] + else: + raise ValueError(f"Unexpected data type {bigframes_dtype}") + + return BIGFRAMES_TO_IBIS[bigframes_dtype] + + +def literal_to_ibis_scalar( + literal, force_dtype: typing.Optional[Dtype] = None, validate: bool = True +): + """Accept any literal and, if possible, return an Ibis Scalar + expression with a BigQuery DataFrames compatible data type + + Args: + literal: any value accepted by Ibis + force_dtype: force the value to a specific dtype + validate: + If true, will raise ValueError if type cannot be stored in a + BigQuery DataFrames object. If used as a subexpression, this should + be disabled. + + Returns: + An ibis Scalar supported by BigQuery DataFrame + + Raises: + ValueError: if passed literal cannot be coerced to a + BigQuery DataFrames compatible scalar + """ + ibis_dtype = BIGFRAMES_TO_IBIS[force_dtype] if force_dtype else None + + if pd.api.types.is_list_like(literal): + if validate: + raise ValueError("List types can't be stored in BigQuery DataFrames") + # "correct" way would be to use ibis.array, but this produces invalid BQ SQL syntax + return tuple(literal) + if not pd.api.types.is_list_like(literal) and pd.isna(literal): + if ibis_dtype: + return ibis.null().cast(ibis_dtype) + else: + return ibis.null() + + scalar_expr = ibis.literal(literal) + if ibis_dtype: + scalar_expr = ibis.literal(literal, ibis_dtype) + elif scalar_expr.type().is_floating(): + scalar_expr = ibis.literal(literal, ibis_dtypes.float64) + elif scalar_expr.type().is_integer(): + scalar_expr = ibis.literal(literal, ibis_dtypes.int64) + + # TODO(bmil): support other literals that can be coerced to compatible types + if validate and (scalar_expr.type() not in BIGFRAMES_TO_IBIS.values()): + raise ValueError(f"Literal did not coerce to a supported data type: {literal}") + + return scalar_expr + + +def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.Value: + """Perform compatible type casts of ibis values + + Args: + value: Ibis value, which could be a literal, scalar, or column + + to_type: The Ibis type to cast to + + Returns: + A new Ibis value of type to_type + + Raises: + TypeError: if the type cast cannot be executed""" + if value.type() == to_type: + return value + # casts that just work + # TODO(bmil): add to this as more casts are verified + good_casts = { + ibis_dtypes.bool: (ibis_dtypes.int64,), + ibis_dtypes.int64: ( + ibis_dtypes.bool, + ibis_dtypes.float64, + ibis_dtypes.string, + ), + ibis_dtypes.float64: (ibis_dtypes.string,), + ibis_dtypes.string: (), + ibis_dtypes.date: (), + ibis_dtypes.time: (), + ibis_dtypes.timestamp: (ibis_dtypes.Timestamp(timezone="UTC"),), + ibis_dtypes.Timestamp(timezone="UTC"): (ibis_dtypes.timestamp,), + } + + value = ibis_value_to_canonical_type(value) + if value.type() in good_casts: + if to_type in good_casts[value.type()]: + return value.cast(to_type) + else: + # this should never happen + raise TypeError(f"Unexpected value type {value.type()}") + + # casts that need some encouragement + + # BigQuery casts bools to lower case strings. Capitalize the result to match Pandas + # TODO(bmil): remove this workaround after fixing Ibis + if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.string: + return typing.cast(ibis_types.StringValue, value.cast(to_type)).capitalize() + + if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.float64: + return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64) + + raise TypeError(f"Unsupported cast {value.type()} to {to_type}") diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py new file mode 100644 index 00000000000..fa0ed0c64d1 --- /dev/null +++ b/bigframes/formatting_helpers.py @@ -0,0 +1,120 @@ +# 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. + +"""Shared helper functions for formatting jobs related info.""" + +import datetime +from typing import Optional + +import google.cloud.bigquery as bigquery +import humanize +import IPython +import IPython.display as display +import ipywidgets as widgets + +query_job_prop_pairs = { + "Job Id": "job_id", + "Destination Table": "destination", + "Slot Time": "slot_millis", + "Bytes Processed": "estimated_bytes_processed", + "Cache hit": "cache_hit", +} + + +def repr_query_job(query_job: Optional[bigquery.QueryJob]): + if query_job is None: + return widgets.HTML("No job information available") + table_html = "" + table_html += "" + for key, value in query_job_prop_pairs.items(): + job_val = getattr(query_job, value) + if job_val is not None: + if key == "Job Id": # add link to job + table_html += f"""""" + elif key == "Slot Time": + table_html += ( + f"""""" + ) + elif key == "Bytes Processed": + table_html += f"""""" + else: + table_html += f"""""" + table_html += "
{key}{job_val}
{key}{get_formatted_time(job_val)}
{key}{get_formatted_bytes(job_val)}
{key}{job_val}
" + return widgets.HTML(table_html) + + +def wait_for_job( + query_job: bigquery.QueryJob, + max_results: Optional[int] = None, + progress_bar: Optional[str] = None, +) -> bigquery.table.RowIterator: + """Return query results. Displays a progress bar while the query is running + Args: + query_job: + The job representing the execution of the query on the server. + max_results: + The maximum number of rows the row iterator should return. + Returns: + A row iterator over the query results. + """ + loading_bar = widgets.HTML(get_query_job_loading_html(query_job)) + if progress_bar == "auto": + progress_bar = "notebook" if in_ipython() else "terminal" + + if progress_bar == "notebook": + display.display(loading_bar) + query_result = query_job.result(max_results=max_results) + query_job.reload() + loading_bar.close() + elif progress_bar == "terminal": + print(get_query_job_loading_string(query_job)) + query_result = query_job.result(max_results=max_results) + query_job.reload() + return query_result + + +def get_job_url(query_job: bigquery.QueryJob): + if ( + query_job.project is None + or query_job.location is None + or query_job.job_id is None + ): + return None + return f"""https://console.cloud.google.com/bigquery?project={query_job.project}&j=bq:{query_job.location}:{query_job.job_id}&page=queryresults""" + + +def get_query_job_loading_html(query_job: bigquery.QueryJob): + return f"""Job {query_job.job_id} is {query_job.state}. Open Job""" + + +def get_query_job_loading_string(query_job: bigquery.QueryJob): + return ( + f"""Job {query_job.job_id} is {query_job.state}. \n{get_job_url(query_job)}""" + ) + + +def get_formatted_time(val): + try: + return humanize.naturaldelta(datetime.timedelta(milliseconds=float(val))) + except Exception: + return val + + +def get_formatted_bytes(val): + return humanize.naturalsize(val) + + +def in_ipython(): + """Return True iff we're in a colab-like IPython.""" + return hasattr(IPython.get_ipython(), "kernel") diff --git a/bigframes/ml/__init__.py b/bigframes/ml/__init__.py new file mode 100644 index 00000000000..55c8709d8d8 --- /dev/null +++ b/bigframes/ml/__init__.py @@ -0,0 +1,29 @@ +# 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. + +"""BigQuery DataFrames ML provides a SKLearn-like API on the BigQuery engine.""" + +__all__ = [ + "cluster", + "compose", + "decomposition", + "linear_model", + "metrics", + "model_selection", + "pipeline", + "preprocessing", + "llm", + "forecasting", + "imported", +] diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py new file mode 100644 index 00000000000..e0a1027dfea --- /dev/null +++ b/bigframes/ml/base.py @@ -0,0 +1,135 @@ +# 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. + +""" +Wraps primitives for machine learning with BQML + +This library is an evolving attempt to +- implement BigQuery DataFrames API for BQML +- follow as close as possible the API design of SKLearn + https://arxiv.org/pdf/1309.0238.pdf +""" + +import abc +from typing import Optional, TypeVar + +from bigframes.ml.core import BqmlModel +import third_party.bigframes_vendored.sklearn.base + + +class BaseEstimator(third_party.bigframes_vendored.sklearn.base.BaseEstimator, abc.ABC): + """ + A BigQuery DataFrames machine learning component following the SKLearn API + design Ref: https://bit.ly/3NyhKjN + + The estimator is the fundamental abstraction for all learning components. This includes learning + algorithms, and also some preprocessing routines. + + This base class provides shared methods for inspecting parameters, and for building a consistent + string representation of the component. By convention, the __init__ of all descendents will be + assumed to be the list of hyperparameters. + + All descendents of this class should implement: + def __init__(self, hyperparameter_1=default_1, hyperparameter_2=default_2, hyperparameter3, ...): + '''Set hyperparameters''' + self.hyperparameter_1 = hyperparameter_1 + self.hyperparameter_2 = hyperparameter_2 + self.hyperparameter3 = hyperparameter3 + ... + Note: the object variable names must be exactly the same with parameter names. In order to utilize __repr__. + + fit(X, y) method is optional. + + The types of decendents of this class should be: + + 1) Predictors + These extend the interface with a .predict(self, x_test) method which predicts the target values + according to the parameters that were calculated in .fit() + + def predict(self, x_test: Union[DataFrame, Series]) -> Union[DataFrame, Series]: + '''Predict the target values according to the parameters that were calculated in .fit''' + ... + + 2) Transformers + These extend the interface with .transform(self, x) and .fit_transform(x_train) methods, which + apply data processing steps such as scaling that must be fitted to training data + + def transform(self, x: Union[DataFrame, Series]) -> Union[DataFrame, Series]: + '''Transform the data according to the parameters that were calculated in .fit()''' + ... + + def fit_transform(self, x_train: Union[DataFrame, Series], y_train: Union[DataFrame, Series]): + '''Perform both fit() and transform()''' + ... + """ + + def __repr__(self): + """Print the estimator's constructor with all non-default parameter values""" + + # Estimator pretty printer adapted from Sklearn's, which is in turn an adaption of + # the inbuilt pretty-printer in CPython + import third_party.bigframes_vendored.cpython._pprint as adapted_pprint + + prettyprinter = adapted_pprint._EstimatorPrettyPrinter( + compact=True, indent=1, indent_at_name=True, n_max_elements_to_show=30 + ) + + return prettyprinter.pformat(self) + + +class Predictor(BaseEstimator): + """A BigQuery DataFrames ML Model base class that can be used to predict outputs.""" + + def __init__(self): + self._bqml_model: Optional[BqmlModel] = None + + @abc.abstractmethod + def predict(self, X): + pass + + _T = TypeVar("_T", bound="Predictor") + + def register(self: _T, vertex_ai_model_id: Optional[str] = None) -> _T: + """Register the model to Vertex AI. + Args: + vertex_ai_model_id: optional string id as model id in Vertex. If not set, will by default to 'bigframes_{bq_model_id}'. + + Returns: + BigQuery DataFrames Model after register. + """ + if not self._bqml_model: + raise RuntimeError("A model must be trained before register.") + + self._bqml_model.register(vertex_ai_model_id) + return self + + +class TrainablePredictor(Predictor): + """A BigQuery DataFrame ML Model base class that can be used to fit and predict outputs. + + Also the predictor can be attached to a pipeline with transformers.""" + + @abc.abstractmethod + def fit(self, X, y, transforms): + pass + + # TODO(b/289280565): enable signatures after updating KMeans and PCA + # @abc.abstractmethod + # def score(self, X, y): + # pass + + # TODO(b/291812029): move to Predictor after implement in LLM and imported models + @abc.abstractmethod + def to_gbq(self, model_name, replace): + pass diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py new file mode 100644 index 00000000000..2dd95bfc1c9 --- /dev/null +++ b/bigframes/ml/cluster.py @@ -0,0 +1,97 @@ +# 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. + +"""Clustering models. This module is styled after Scikit-Learn's cluster module: +https://scikit-learn.org/stable/modules/clustering.html""" + +from __future__ import annotations + +from typing import cast, Dict, List, Optional, TYPE_CHECKING + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core +import third_party.bigframes_vendored.sklearn.cluster._kmeans + + +class KMeans( + third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans, + bigframes.ml.base.TrainablePredictor, +): + + __doc__ = third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans.__doc__ + + def __init__(self, n_clusters=8): + self.n_clusters = n_clusters + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> KMeans: + assert model.model_type == "KMEANS" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + if "numClusters" in last_fitting: + kwargs["n_clusters"] = int(last_fitting["numClusters"]) + + new_kmeans = KMeans(**kwargs) + new_kmeans._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_kmeans + + @property + def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: + """The model options as they will be set for BQML""" + return {"model_type": "KMEANS", "num_clusters": self.n_clusters} + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y=None, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + train_X=X, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + return cast( + bigframes.dataframe.DataFrame, self._bqml_model.predict(X)[["CENTROID_ID"]] + ) + + def to_gbq(self, model_name: str, replace: bool = False) -> KMeans: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py new file mode 100644 index 00000000000..09645d4cf8e --- /dev/null +++ b/bigframes/ml/compose.py @@ -0,0 +1,123 @@ +# 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. + +"""Build composite transformers on heterogenous data. This module is styled +after Scikit-Learn's compose module: +https://scikit-learn.org/stable/modules/classes.html#module-sklearn.compose""" + +from __future__ import annotations + +import typing +from typing import List, Optional, Tuple, TYPE_CHECKING, Union + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.compose +import bigframes.ml.core +import bigframes.ml.preprocessing +import third_party.bigframes_vendored.sklearn.compose._column_transformer + +CompilablePreprocessorType = Union[ + bigframes.ml.preprocessing.OneHotEncoder, + bigframes.ml.preprocessing.StandardScaler, +] + + +class ColumnTransformer( + third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, + bigframes.ml.base.BaseEstimator, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer.__doc__ + ) + + def __init__( + self, + transformers: List[ + Tuple[ + str, + CompilablePreprocessorType, + Union[str, List[str]], + ] + ], + ): + self.transformers = transformers + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @property + def transformers_( + self, + ) -> List[Tuple[str, CompilablePreprocessorType, str,]]: + """The collection of transformers as tuples of (name, transformer, column)""" + result: List[ + Tuple[ + str, + CompilablePreprocessorType, + str, + ] + ] = [] + for entry in self.transformers: + name, transformer, column_or_columns = entry + if isinstance(column_or_columns, str): + result.append((name, transformer, column_or_columns)) + else: + for column in column_or_columns: + result.append((name, transformer, column)) + + return result + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + transformer._compile_to_sql([column])[0] + for column in columns + for _, transformer, target_column in self.transformers_ + if column == target_column + ] + + def fit( + self, + X: bigframes.dataframe.DataFrame, + ): + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + + def transform( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + df = self._bqml_model.transform(X) + return typing.cast( + bigframes.dataframe.DataFrame, + df[self._output_names], + ) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py new file mode 100644 index 00000000000..2663a8e17e5 --- /dev/null +++ b/bigframes/ml/core.py @@ -0,0 +1,334 @@ +# 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. + +"""Core operations for BQML based models""" + +from __future__ import annotations + +from typing import Callable, cast, Iterable, Mapping, Optional, Union +import uuid + +from google.cloud import bigquery + +import bigframes.dataframe +import bigframes.ml.sql +import bigframes.session + + +class BqmlModel: + """Represents an existing BQML model in BigQuery. + + Wraps the BQML API and SQL interface to expose the functionality needed for + BigQuery DataFrames ML. + """ + + def __init__(self, session: bigframes.session.Session, model: bigquery.Model): + self._session = session + self._model = model + + @property + def session(self) -> bigframes.Session: + """Get the BigQuery DataFrames session that this BQML model wrapper is tied to""" + return self._session + + @property + def model_name(self) -> str: + """Get the fully qualified name of the model, i.e. project_id.dataset_id.model_id""" + return f"{self._model.project}.{self._model.dataset_id}.{self._model.model_id}" + + @property + def model(self) -> bigquery.Model: + """Get the BQML model associated with this wrapper""" + return self._model + + @staticmethod + def _apply_sql( + session: bigframes.Session, + input_data: bigframes.dataframe.DataFrame, + func: Callable[[str], str], + ) -> bigframes.dataframe.DataFrame: + """Helper to wrap a dataframe in a SQL query, keeping the index intact. + + Args: + session: the active bigframes.Session + + input_data: the dataframe to be wrapped + + func: a function that will accept a SQL string and produce a new SQL + string from which to construct the output dataframe. It must + include the index columns of the input SQL. + """ + source_sql, tagged_index_cols = input_data.to_sql_query( + always_include_index=True + ) + + if len(tagged_index_cols) != 1: + raise NotImplementedError("Only exactly one index column is supported") + + index_col_name, is_named_index = tagged_index_cols[0] + sql = func(source_sql) + df = session.read_gbq(sql, index_col=[index_col_name]) + if not is_named_index: + df.index.name = None + + return df + + def predict( + self, input_data: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + # TODO: validate input data schema + return self._apply_sql( + self._session, + input_data, + lambda source_sql: bigframes.ml.sql.ml_predict( + model_name=self.model_name, source_sql=source_sql + ), + ) + + def transform( + self, input_data: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + # TODO: validate input data schema + return self._apply_sql( + self._session, + input_data, + lambda source_sql: bigframes.ml.sql.ml_transform( + model_name=self.model_name, source_sql=source_sql + ), + ) + + def generate_text( + self, + input_data: bigframes.dataframe.DataFrame, + options: Mapping[str, int | float], + ) -> bigframes.dataframe.DataFrame: + # TODO: validate input data schema + return self._apply_sql( + self._session, + input_data, + lambda source_sql: bigframes.ml.sql.ml_generate_text( + model_name=self.model_name, + source_sql=source_sql, + struct_options=bigframes.ml.sql.struct_options(**options), + ), + ) + + def embed_text( + self, + input_data: bigframes.dataframe.DataFrame, + options: Mapping[str, int | float], + ) -> bigframes.dataframe.DataFrame: + # TODO: validate input data schema + return self._apply_sql( + self._session, + input_data, + lambda source_sql: bigframes.ml.sql.ml_embed_text( + model_name=self.model_name, + source_sql=source_sql, + struct_options=bigframes.ml.sql.struct_options(**options), + ), + ) + + def forecast(self) -> bigframes.dataframe.DataFrame: + sql = bigframes.ml.sql.ml_forecast(self.model_name) + return self._session.read_gbq(sql) + + def evaluate(self, input_data: Union[bigframes.dataframe.DataFrame, None] = None): + # TODO: validate input data schema + # Note: don't need index as evaluate returns a new table + source_sql, _ = ( + input_data.to_sql_query(always_include_index=False) + if (input_data is not None) + else (None, None) + ) + sql = bigframes.ml.sql.ml_evaluate(self.model_name, source_sql) + + return self._session.read_gbq(sql) + + def copy(self, new_model_name, replace=False) -> BqmlModel: + job_config = bigquery.job.CopyJobConfig() + if replace: + job_config.write_disposition = "WRITE_TRUNCATE" + + self._session.bqclient.copy_table( + self.model_name, new_model_name, job_config=job_config + ).result() + + new_model = self._session.bqclient.get_model(new_model_name) + return BqmlModel(self._session, new_model) + + def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: + if vertex_ai_model_id is None: + # vertex id needs to start with letters. https://cloud.google.com/vertex-ai/docs/general/resource-naming + vertex_ai_model_id = "bigframes_" + cast(str, self._model.model_id) + + options_sql = bigframes.ml.sql.options( + **{"vertex_ai_model_id": vertex_ai_model_id} + ) + sql = bigframes.ml.sql.alter_model(self.model_name, options_sql=options_sql) + # Register the model and wait it to finish + self._session.bqclient.query(sql).result() + + self._model = self._session.bqclient.get_model(self.model_name) + return self + + +def create_bqml_model( + train_X: bigframes.dataframe.DataFrame, + train_y: Optional[bigframes.dataframe.DataFrame] = None, + transforms: Optional[Iterable[str]] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, +) -> BqmlModel: + """Create a session-temporary BQML model with the CREATE MODEL statement + + Args: + train_X: features columns for training + train_y: labels columns for training, if applicable + transforms: an optional list of SQL expressions that implement preprocessing + on top of the input data. Generates a BQML TRANSFORM clause + options: a dict of options to configure the model. Generates a BQML OPTIONS + clause + + Returns: a BqmlModel, wrapping a trained model in BigQuery + """ + options = dict(options) + if train_y is None: + input_data = train_X + else: + # TODO: handle case where train_y columns are renamed in the join + input_data = train_X.join(train_y, how="outer") + options.update({"INPUT_LABEL_COLS": train_y.columns.tolist()}) + + # pickpocket session object from the dataframe + session = train_X._get_block().expr._session + + # TODO(garrettwu): add wrapper to select the feature columns + # for now, drop index to avoid including the index in feature columns + input_data = input_data.reset_index(drop=True) + + model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" + source_sql = input_data.sql + options_sql = bigframes.ml.sql.options(**options) + transform_sql = ( + bigframes.ml.sql.transform(*transforms) if transforms is not None else None + ) + sql = bigframes.ml.sql.create_model( + model_name=model_name, + source_sql=source_sql, + transform_sql=transform_sql, + options_sql=options_sql, + ) + + # fit the model, synchronously + session.bqclient.query(sql).result() + + model = session.bqclient.get_model(model_name) + return BqmlModel(session, model) + + +def create_bqml_time_series_model( + train_X: bigframes.dataframe.DataFrame, + train_y: bigframes.dataframe.DataFrame, + transforms: Optional[Iterable[str]] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, +) -> BqmlModel: + + assert ( + train_X.columns.size == 1 + ), "Time series timestamp input must only contain 1 column." + assert ( + train_y.columns.size == 1 + ), "Time stamp data input must only contain 1 column." + + options = dict(options) + input_data = train_X.join(train_y, how="outer") + options.update({"TIME_SERIES_TIMESTAMP_COL": train_X.columns.tolist()[0]}) + options.update({"TIME_SERIES_DATA_COL": train_y.columns.tolist()[0]}) + # pickpocket session object from the dataframe + session = train_X._get_block().expr._session + + model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" + source_sql = input_data.sql + options_sql = bigframes.ml.sql.options(**options) + + transform_sql = ( + bigframes.ml.sql.transform(*transforms) if transforms is not None else None + ) + sql = bigframes.ml.sql.create_model( + model_name=model_name, + source_sql=source_sql, + transform_sql=transform_sql, + options_sql=options_sql, + ) + + # fit the model, synchronously + session.bqclient.query(sql).result() + + model = session.bqclient.get_model(model_name) + return BqmlModel(session, model) + + +def create_bqml_remote_model( + session: bigframes.Session, + connection_name: str, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, +) -> BqmlModel: + """Create a session-temporary BQML remote model with the CREATE MODEL statement + + Args: + connection_name: a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection + options: a dict of options to configure the model. Generates a BQML OPTIONS + clause + + Returns: a BqmlModel, wrapping a trained model in BigQuery + """ + model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" + options_sql = bigframes.ml.sql.options(**options) + sql = bigframes.ml.sql.create_remote_model( + model_name=model_name, + connection_name=connection_name, + options_sql=options_sql, + ) + + # create the model, synchronously + session.bqclient.query(sql).result() + + model = session.bqclient.get_model(model_name) + return BqmlModel(session, model) + + +def create_bqml_imported_model( + session: bigframes.Session, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, +) -> BqmlModel: + """Create a session-temporary BQML imported model with the CREATE MODEL statement + + Args: + options: a dict of options to configure the model. Generates a BQML OPTIONS + clause + + Returns: a BqmlModel, wrapping a trained model in BigQuery + """ + model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" + options_sql = bigframes.ml.sql.options(**options) + sql = bigframes.ml.sql.create_imported_model( + model_name=model_name, + options_sql=options_sql, + ) + + # create the model, synchronously + session.bqclient.query(sql).result() + + model = session.bqclient.get_model(model_name) + return BqmlModel(session, model) diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py new file mode 100644 index 00000000000..08551d1cb8c --- /dev/null +++ b/bigframes/ml/decomposition.py @@ -0,0 +1,104 @@ +# 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. + +"""Matrix Decomposition models. This module is styled after Scikit-Learn's decomposition module: +https://scikit-learn.org/stable/modules/decomposition.html""" + +from __future__ import annotations + +from typing import cast, List, Optional, TYPE_CHECKING + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core +import third_party.bigframes_vendored.sklearn.decomposition._pca + + +class PCA( + third_party.bigframes_vendored.sklearn.decomposition._pca.PCA, + bigframes.ml.base.TrainablePredictor, +): + __doc__ = third_party.bigframes_vendored.sklearn.decomposition._pca.PCA.__doc__ + + def __init__(self, n_components=3): + self.n_components = n_components + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> PCA: + assert model.model_type == "PCA" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + if "numPrincipalComponents" in last_fitting: + kwargs["n_components"] = int(last_fitting["numPrincipalComponents"]) + + new_pca = PCA(**kwargs) + new_pca._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_pca + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y=None, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + train_X=X, + transforms=transforms, + options={ + "model_type": "PCA", + "num_principal_components": self.n_components, + }, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + """Predict the closest cluster for each sample in X. + + Args: + X: a BigQuery DataFrame to predict. + y: ignored for API consistency. + + Returns: predicted BigQuery DataFrames.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + return cast( + bigframes.dataframe.DataFrame, + self._bqml_model.predict(X)[ + ["principal_component_" + str(i + 1) for i in range(self.n_components)] + ], + ) + + def to_gbq(self, model_name: str, replace: bool = False) -> PCA: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py new file mode 100644 index 00000000000..117759ca1c6 --- /dev/null +++ b/bigframes/ml/ensemble.py @@ -0,0 +1,672 @@ +# 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. + +"""Linear models. This module is styled after Scikit-Learn's linear_model module: +https://scikit-learn.org/stable/modules/linear_model.html""" + +from __future__ import annotations + +from typing import cast, Dict, List, Literal, Optional, TYPE_CHECKING + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core +import third_party.bigframes_vendored.sklearn.ensemble._forest +import third_party.bigframes_vendored.xgboost.sklearn + +_BQML_PARAMS_MAPPING = { + "booster": "boosterType", + "tree_method": "treeMethod", + "early_stop": "earlyStop", + "colsample_bytree": "colsampleBylevel", + "colsample_bylevel": "colsampleBytree", + "colsample_bynode": "colsampleBynode", + "gamma": "minSplitLoss", + "subsample": "subsample", + "reg_alpha": "l1Regularization", + "reg_lambda": "l2Regularization", + "learning_rate": "learnRate", + "min_rel_progress": "minRelativeProgress", + "num_parallel_tree": "numParallelTree", + "min_tree_child_weight": "minTreeChildWeight", + "max_depth": "maxTreeDepth", + "max_iterations": "maxIterations", +} + + +class XGBRegressor( + third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor, + bigframes.ml.base.TrainablePredictor, +): + __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor.__doc__ + + def __init__( + self, + num_parallel_tree: int = 1, + booster: Literal["gbtree", "dart"] = "gbtree", + dart_normalized_type: Literal["TREE", "FOREST"] = "TREE", + tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", + min_tree_child_weight: int = 1, + colsample_bytree=1.0, + colsample_bylevel=1.0, + colsample_bynode=1.0, + gamma=0.0, + max_depth: int = 6, + subsample=1.0, + reg_alpha=0.0, + reg_lambda=1.0, + early_stop=True, + learning_rate=0.3, + max_iterations: int = 20, + min_rel_progress=0.01, + enable_global_explain=False, + xgboost_version: Literal["0.9", "1.1"] = "0.9", + ): + self.num_parallel_tree = num_parallel_tree + self.booster = booster + self.dart_normalized_type = dart_normalized_type + self.tree_method = tree_method + self.min_tree_child_weight = min_tree_child_weight + self.colsample_bytree = colsample_bytree + self.colsample_bylevel = colsample_bylevel + self.colsample_bynode = colsample_bynode + self.gamma = gamma + self.max_depth = max_depth + self.subsample = subsample + self.reg_alpha = reg_alpha + self.reg_lambda = reg_lambda + self.early_stop = early_stop + self.learning_rate = learning_rate + self.max_iterations = max_iterations + self.min_rel_progress = min_rel_progress + self.enable_global_explain = enable_global_explain + self.xgboost_version = xgboost_version + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBRegressor: + assert model.model_type == "BOOSTED_TREE_REGRESSOR" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + + dummy_regressor = XGBRegressor() + for bf_param, bf_value in dummy_regressor.__dict__.items(): + bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) + if bqml_param is not None: + kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) + + new_xgb_regressor = XGBRegressor(**kwargs) + new_xgb_regressor._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_xgb_regressor + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "BOOSTED_TREE_REGRESSOR", + "data_split_method": "NO_SPLIT", + "num_parallel_tree": self.num_parallel_tree, + "booster_type": self.booster, + "tree_method": self.tree_method, + "min_tree_child_weight": self.min_tree_child_weight, + "colsample_bytree": self.colsample_bytree, + "colsample_bylevel": self.colsample_bylevel, + "colsample_bynode": self.colsample_bynode, + "min_split_loss": self.gamma, + "max_tree_depth": self.max_depth, + "subsample": self.subsample, + "l1_reg": self.reg_alpha, + "l2_reg": self.reg_lambda, + "early_stop": self.early_stop, + "learn_rate": self.learning_rate, + "max_iterations": self.max_iterations, + "min_rel_progress": self.min_rel_progress, + "enable_global_explain": self.enable_global_explain, + "xgboost_version": self.xgboost_version, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ): + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) + + +class XGBClassifier( + third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier, + bigframes.ml.base.TrainablePredictor, +): + + __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier.__doc__ + + def __init__( + self, + num_parallel_tree: int = 1, + booster: Literal["gbtree", "dart"] = "gbtree", + dart_normalized_type: Literal["TREE", "FOREST"] = "TREE", + tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", + min_tree_child_weight: int = 1, + colsample_bytree=1.0, + colsample_bylevel=1.0, + colsample_bynode=1.0, + gamma=0.0, + max_depth: int = 6, + subsample=1.0, + reg_alpha=0.0, + reg_lambda=1.0, + early_stop=True, + learning_rate=0.3, + max_iterations: int = 20, + min_rel_progress=0.01, + enable_global_explain=False, + xgboost_version: Literal["0.9", "1.1"] = "0.9", + ): + self.num_parallel_tree = num_parallel_tree + self.booster = booster + self.dart_normalized_type = dart_normalized_type + self.tree_method = tree_method + self.min_tree_child_weight = min_tree_child_weight + self.colsample_bytree = colsample_bytree + self.colsample_bylevel = colsample_bylevel + self.colsample_bynode = colsample_bynode + self.gamma = gamma + self.max_depth = max_depth + self.subsample = subsample + self.reg_alpha = reg_alpha + self.reg_lambda = reg_lambda + self.early_stop = early_stop + self.learning_rate = learning_rate + self.max_iterations = max_iterations + self.min_rel_progress = min_rel_progress + self.enable_global_explain = enable_global_explain + self.xgboost_version = xgboost_version + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBClassifier: + assert model.model_type == "BOOSTED_TREE_CLASSIFIER" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + + dummy_classifier = XGBClassifier() + for bf_param, bf_value in dummy_classifier.__dict__.items(): + bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) + if bqml_param is not None: + kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) + + new_xgb_classifier = XGBClassifier(**kwargs) + new_xgb_classifier._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_xgb_classifier + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "BOOSTED_TREE_CLASSIFIER", + "data_split_method": "NO_SPLIT", + "num_parallel_tree": self.num_parallel_tree, + "booster_type": self.booster, + "tree_method": self.tree_method, + "min_tree_child_weight": self.min_tree_child_weight, + "colsample_bytree": self.colsample_bytree, + "colsample_bylevel": self.colsample_bylevel, + "colsample_bynode": self.colsample_bynode, + "min_split_loss": self.gamma, + "max_tree_depth": self.max_depth, + "subsample": self.subsample, + "l1_reg": self.reg_alpha, + "l2_reg": self.reg_lambda, + "early_stop": self.early_stop, + "learn_rate": self.learning_rate, + "max_iterations": self.max_iterations, + "min_rel_progress": self.min_rel_progress, + "enable_global_explain": self.enable_global_explain, + "xgboost_version": self.xgboost_version, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ): + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) + + +class RandomForestRegressor( + third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor, + bigframes.ml.base.TrainablePredictor, +): + + __doc__ = ( + third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor.__doc__ + ) + + def __init__( + self, + num_parallel_tree: int = 100, + tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", + min_tree_child_weight: int = 1, + colsample_bytree=1.0, + colsample_bylevel=1.0, + colsample_bynode=0.8, + gamma=0.00, + max_depth: int = 15, + subsample=0.8, + reg_alpha=0.0, + reg_lambda=1.0, + early_stop=True, + min_rel_progress=0.01, + enable_global_explain=False, + xgboost_version: Literal["0.9", "1.1"] = "0.9", + ): + self.num_parallel_tree = num_parallel_tree + self.tree_method = tree_method + self.min_tree_child_weight = min_tree_child_weight + self.colsample_bytree = colsample_bytree + self.colsample_bylevel = colsample_bylevel + self.colsample_bynode = colsample_bynode + self.gamma = gamma + self.max_depth = max_depth + self.subsample = subsample + self.reg_alpha = reg_alpha + self.reg_lambda = reg_lambda + self.early_stop = early_stop + self.min_rel_progress = min_rel_progress + self.enable_global_explain = enable_global_explain + self.xgboost_version = xgboost_version + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq( + session: bigframes.Session, model: bigquery.Model + ) -> RandomForestRegressor: + assert model.model_type == "RANDOM_FOREST_REGRESSOR" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + + dummy_model = RandomForestRegressor() + for bf_param, bf_value in dummy_model.__dict__.items(): + bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) + if bqml_param is not None: + kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) + + new_random_forest_regressor = RandomForestRegressor(**kwargs) + new_random_forest_regressor._bqml_model = bigframes.ml.core.BqmlModel( + session, model + ) + return new_random_forest_regressor + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "RANDOM_FOREST_REGRESSOR", + "num_parallel_tree": self.num_parallel_tree, + "tree_method": self.tree_method, + "min_tree_child_weight": self.min_tree_child_weight, + "colsample_bytree": self.colsample_bytree, + "colsample_bylevel": self.colsample_bylevel, + "colsample_bynode": self.colsample_bynode, + "min_split_loss": self.gamma, + "max_tree_depth": self.max_depth, + "subsample": self.subsample, + "l1_reg": self.reg_alpha, + "l2_reg": self.reg_lambda, + "early_stop": self.early_stop, + "min_rel_progress": self.min_rel_progress, + "data_split_method": "NO_SPLIT", + "enable_global_explain": self.enable_global_explain, + "xgboost_version": self.xgboost_version, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ): + """Calculate evaluation metrics of the model. + + Args: + X: a BigFrames DataFrame as evaluation data. + y: a BigFrames DataFrame as evaluation labels. + + Returns: a BigFrames DataFrame as evaluation result.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + if (X is None) != (y is None): + raise ValueError( + "Either both or neither of test_X and test_y must be specified" + ) + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegressor: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) + + +class RandomForestClassifier( + third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier, + bigframes.ml.base.TrainablePredictor, +): + + __doc__ = ( + third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier.__doc__ + ) + + def __init__( + self, + num_parallel_tree: int = 100, + tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", + min_tree_child_weight: int = 1, + colsample_bytree: float = 1.0, + colsample_bylevel: float = 1.0, + colsample_bynode: float = 0.8, + gamma: float = 0.00, + max_depth: int = 15, + subsample: float = 0.8, + reg_alpha: float = 0.0, + reg_lambda: float = 1.0, + early_stop=True, + min_rel_progress: float = 0.01, + enable_global_explain=False, + xgboost_version: Literal["0.9", "1.1"] = "0.9", + ): + self.num_parallel_tree = num_parallel_tree + self.tree_method = tree_method + self.min_tree_child_weight = min_tree_child_weight + self.colsample_bytree = colsample_bytree + self.colsample_bylevel = colsample_bylevel + self.colsample_bynode = colsample_bynode + self.gamma = gamma + self.max_depth = max_depth + self.subsample = subsample + self.reg_alpha = reg_alpha + self.reg_lambda = reg_lambda + self.early_stop = early_stop + self.min_rel_progress = min_rel_progress + self.enable_global_explain = enable_global_explain + self.xgboost_version = xgboost_version + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq( + session: bigframes.Session, model: bigquery.Model + ) -> RandomForestClassifier: + assert model.model_type == "RANDOM_FOREST_CLASSIFIER" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + + dummy_model = RandomForestClassifier() + for bf_param, bf_value in dummy_model.__dict__.items(): + bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) + if bqml_param is not None: + kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) + + new_random_forest_classifier = RandomForestClassifier(**kwargs) + new_random_forest_classifier._bqml_model = bigframes.ml.core.BqmlModel( + session, model + ) + return new_random_forest_classifier + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "RANDOM_FOREST_CLASSIFIER", + "num_parallel_tree": self.num_parallel_tree, + "tree_method": self.tree_method, + "min_tree_child_weight": self.min_tree_child_weight, + "colsample_bytree": self.colsample_bytree, + "colsample_bylevel": self.colsample_bylevel, + "colsample_bynode": self.colsample_bynode, + "min_split_loss": self.gamma, + "max_tree_depth": self.max_depth, + "subsample": self.subsample, + "l1_reg": self.reg_alpha, + "l2_reg": self.reg_lambda, + "early_stop": self.early_stop, + "min_rel_progress": self.min_rel_progress, + "data_split_method": "NO_SPLIT", + "enable_global_explain": self.enable_global_explain, + "xgboost_version": self.xgboost_version, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ): + """Calculate evaluation metrics of the model. + + Args: + X: a BigFrames DataFrame as evaluation data. + y: a BigFrames DataFrame as evaluation labels. + + Returns: a BigFrames DataFrame as evaluation result.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + if (X is None) != (y is None): + raise ValueError( + "Either both or neither of test_X and test_y must be specified" + ) + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestClassifier: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py new file mode 100644 index 00000000000..24afaad7f26 --- /dev/null +++ b/bigframes/ml/forecasting.py @@ -0,0 +1,125 @@ +# 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. + +from __future__ import annotations + +from typing import cast, Dict, List, Optional, TYPE_CHECKING + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core + +_PREDICT_OUTPUT_COLUMNS = ["forecast_timestamp", "forecast_value"] + + +class ARIMAPlus(bigframes.ml.base.TrainablePredictor): + """Time Series ARIMA Plus model.""" + + def __init__(self): + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> ARIMAPlus: + assert model.model_type == "ARIMA_PLUS" + + kwargs: Dict[str, str | int | bool | float | List[str]] = {} + + new_arima_plus = ARIMAPlus(**kwargs) + new_arima_plus._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_arima_plus + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML.""" + return {"model_type": "ARIMA_PLUS"} + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + """Fit the model to training data + + Args: + X: A dataframe of training timestamp. + + y: Target values for training.""" + self._bqml_model = bigframes.ml.core.create_bqml_time_series_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict(self, X=None) -> bigframes.dataframe.DataFrame: + """Predict the closest cluster for each sample in X. + + Args: + X: ignored, to be compatible with other APIs. + Returns: + The predicted BigQuery DataFrames. Which contains 2 columns + "forecast_timestamp" and "forecast_value". + """ + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + return cast( + bigframes.dataframe.DataFrame, + self._bqml_model.forecast()[_PREDICT_OUTPUT_COLUMNS], + ) + + # Unlike regression models, time series forcasting can only evaluate with unseen data. X and y must be providee. + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ) -> bigframes.dataframe.DataFrame: + """Calculate evaluation metrics of the model. + + Args: + X: + A BigQuery DataFrames only contains 1 column as + evaluation timestamp. The timestamp must be within the horizon + of the model, which by default is 1000 data points. + y: + A BigQuery DataFrames only contains 1 column as + evaluation numeric values. + + Returns: + A BigQuery DataFrames as evaluation result. + """ + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + input_data = X.join(y, how="outer") + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> ARIMAPlus: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py new file mode 100644 index 00000000000..581ee2b1e24 --- /dev/null +++ b/bigframes/ml/imported.py @@ -0,0 +1,101 @@ +# 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. + +from __future__ import annotations + +from typing import cast, TYPE_CHECKING + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core + + +class TensorFlowModel(bigframes.ml.base.Predictor): + """Imported TensorFlow model. + + Args: + session: BQ session to create the model + model_path: GCS path that holds the model files.""" + + def __init__(self, session: bigframes.Session, model_path: str): + self.session = session + self.model_path = model_path + self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + + def _create_bqml_model(self): + options = {"model_type": "TENSORFLOW", "model_path": self.model_path} + return bigframes.ml.core.create_bqml_imported_model( + session=self.session, options=options + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + """Predict the result from input DataFrame. + + Args: + X: Input DataFrame, schema is defined by the model. + + Returns: Output DataFrame, schema is defined by the model.""" + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + +class OnnxModel(bigframes.ml.base.BaseEstimator): + """Imported Open Neural Network Exchange (ONNX) model. + + Args: + session: BQ session to create the model + model_path: GCS path that holds the model files.""" + + def __init__(self, session: bigframes.Session, model_path: str): + self.session = session + self.model_path = model_path + self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + + def _create_bqml_model(self): + options = {"model_type": "ONNX", "model_path": self.model_path} + return bigframes.ml.core.create_bqml_imported_model( + session=self.session, options=options + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + """Predict the result from input DataFrame. + + Args: + X: Input DataFrame, schema is defined by the model. + + Returns: Output DataFrame, schema is defined by the model.""" + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py new file mode 100644 index 00000000000..2b2ea5c2afc --- /dev/null +++ b/bigframes/ml/linear_model.py @@ -0,0 +1,247 @@ +# 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. + +"""Linear models. This module is styled after Scikit-Learn's linear_model module: +https://scikit-learn.org/stable/modules/linear_model.html""" + +from __future__ import annotations + +from typing import cast, Dict, List, Optional, TYPE_CHECKING + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.base +import bigframes.ml.core +import third_party.bigframes_vendored.sklearn.linear_model._base +import third_party.bigframes_vendored.sklearn.linear_model._logistic + + +class LinearRegression( + third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression, + bigframes.ml.base.TrainablePredictor, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression.__doc__ + ) + + def __init__( + self, + fit_intercept=True, + ): + self.fit_intercept = fit_intercept + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq(session: bigframes.Session, model: bigquery.Model) -> LinearRegression: + assert model.model_type == "LINEAR_REGRESSION" + + # TODO(bmil): construct a standard way to extract these properties + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + if "fitIntercept" in last_fitting: + kwargs["fit_intercept"] = last_fitting["fitIntercept"] + + new_linear_regression = LinearRegression(**kwargs) + new_linear_regression._bqml_model = bigframes.ml.core.BqmlModel(session, model) + return new_linear_regression + + @property + def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "LINEAR_REG", + "data_split_method": "NO_SPLIT", + "fit_intercept": self.fit_intercept, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) + + +class LogisticRegression( + third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression, + bigframes.ml.base.TrainablePredictor, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression.__doc__ + ) + + # TODO(ashleyxu) support class_weights in the constructor. + def __init__( + self, + fit_intercept: bool = True, + auto_class_weights: bool = False, + ): + self.fit_intercept = fit_intercept + self.auto_class_weights = auto_class_weights + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + @staticmethod + def _from_bq( + session: bigframes.Session, model: bigquery.Model + ) -> LogisticRegression: + assert model.model_type == "LOGISTIC_REGRESSION" + + kwargs = {} + + # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun + last_fitting = model.training_runs[-1]["trainingOptions"] + if "fitIntercept" in last_fitting: + kwargs["fit_intercept"] = last_fitting["fitIntercept"] + # TODO(ashleyxu): b/285162045 support auto_class_weights once the API is + # fixed and enable the tests. + if "autoClassWeights" in last_fitting: + kwargs["auto_class_weights"] = last_fitting["autoClassWeights"] + # TODO(ashleyxu) support class_weights in the constructor. + # if "labelClassWeights" in last_fitting: + # kwargs["class_weights"] = last_fitting["labelClassWeights"] + + new_logistic_regression = LogisticRegression(**kwargs) + new_logistic_regression._bqml_model = bigframes.ml.core.BqmlModel( + session, model + ) + return new_logistic_regression + + @property + def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: + """The model options as they will be set for BQML""" + return { + "model_type": "LOGISTIC_REG", + "data_split_method": "NO_SPLIT", + "fit_intercept": self.fit_intercept, + "auto_class_weights": self.auto_class_weights, + # TODO(ashleyxu): support class_weights (struct array) + # "class_weights": self.class_weights, + } + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + transforms: Optional[List[str]] = None, + ): + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + y, + transforms=transforms, + options=self._bqml_options, + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before predict") + + df = self._bqml_model.predict(X) + return cast( + bigframes.dataframe.DataFrame, + df[ + [ + cast(str, field.name) + for field in self._bqml_model.model.label_columns + ] + ], + ) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + input_data = ( + X.join(y, how="outer") if (X is not None) and (y is not None) else None + ) + return self._bqml_model.evaluate(input_data) + + def to_gbq(self, model_name: str, replace: bool = False) -> LogisticRegression: + """Save the model to Google Cloud BigQuey. + + Args: + model_name: the name of the model. + replace: whether to replace if the model already exists. Default to False. + + Returns: saved model.""" + if not self._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + # TODO(ashleyxu): b/285162045 support auto_class_weights once the API is + # fixed and enable the tests. + if self.auto_class_weights is True: + raise NotImplementedError("auto_class_weight is not supported yet.") + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py new file mode 100644 index 00000000000..ef4f28f1f7d --- /dev/null +++ b/bigframes/ml/llm.py @@ -0,0 +1,168 @@ +# 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. + +from __future__ import annotations + +from typing import cast + +import bigframes +from bigframes.core import blocks +import bigframes.ml.base +import bigframes.ml.core + +_REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" +_TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" + +_REMOTE_EMBEDDING_GENERATOR_MODEL_CODE = "CLOUD_AI_TEXT_EMBEDDING_MODEL_V1" +_EMBED_TEXT_RESULT_COLUMN = "ml_embed_text_embedding" + + +class PaLM2TextGenerator(bigframes.ml.base.Predictor): + """PaLM2 text generator LLM model. + + Args: + session: BQ session to create the model + connection_name: connection to connect with remote service. str of the format ..""" + + def __init__(self, session: bigframes.Session, connection_name: str): + self.session = session + self.connection_name = connection_name + self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + + def _create_bqml_model(self): + options = { + "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, + } + + return bigframes.ml.core.create_bqml_remote_model( + session=self.session, connection_name=self.connection_name, options=options + ) + + def predict( + self, + X: bigframes.dataframe.DataFrame, + temperature: float = 0.0, + max_output_tokens: int = 128, + top_k: int = 40, + top_p: float = 0.95, + ) -> bigframes.dataframe.DataFrame: + """Predict the result from input DataFrame. + + Args: + X: Input DataFrame, which needs to contain a column with name "prompt". Only the column will be used as input. Prompts can include preamble, questions, suggestions, instructions, or examples. + + temperature: The temperature is used for sampling during the response generation, which occurs when topP and topK are applied. + Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that expect a true or correct response, + while higher temperatures can lead to more diverse or unexpected results. A temperature of 0 is deterministic: + the highest probability token is always selected. For most use cases, try starting with a temperature of 0.2. + Default 0. + + max_output_tokens: Maximum number of tokens that can be generated in the response. Specify a lower value for shorter responses and a higher value for longer responses. + A token may be smaller than a word. A token is approximately four characters. 100 tokens correspond to roughly 60-80 words. + Default 128. + + top_k: Top-k changes how the model selects tokens for output. A top-k of 1 means the selected token is the most probable among all tokens + in the model’s vocabulary (also called greedy decoding), while a top-k of 3 means that the next token is selected from among the 3 most probable tokens (using temperature). + For each token selection step, the top K tokens with the highest probabilities are sampled. Then tokens are further filtered based on topP with the final token selected using temperature sampling. + Specify a lower value for less random responses and a higher value for more random responses. + Default 40. + + top_p: Top-p changes how the model selects tokens for output. Tokens are selected from most K (see topK parameter) probable to least until the sum of their probabilities equals the top-p value. + For example, if tokens A, B, and C have a probability of 0.3, 0.2, and 0.1 and the top-p value is 0.5, then the model will select either A or B as the next token (using temperature) + and not consider C at all. + Specify a lower value for less random responses and a higher value for more random responses. + Default 0.95. + + + Returns: Output DataFrame with only 1 column as the output text results.""" + + # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models + if temperature < 0.0 or temperature > 1.0: + raise ValueError(f"temperature must be [0.0, 1.0], but is {temperature}.") + if max_output_tokens not in range(1, 1025): + raise ValueError( + f"max_output_token must be [1, 1024], but is {max_output_tokens}." + ) + if top_k not in range(1, 41): + raise ValueError(f"top_k must be [1, 40], but is {top_k}.") + if top_p < 0.0 or top_p > 1.0: + raise ValueError(f"top_p must be [0.0, 1.0], but is {top_p}.") + if len(X.columns) != 1: + raise ValueError("Only support one column as input.") + + # BQML identified the column by name + col_label = cast(blocks.Label, X.columns[0]) + X = X.rename(columns={col_label: "prompt"}) + + options = { + "temperature": temperature, + "max_output_tokens": max_output_tokens, + "top_k": top_k, + "top_p": top_p, + "flatten_json_output": True, + } + df = self._bqml_model.generate_text(X, options) + return cast( + bigframes.dataframe.DataFrame, + df[[_TEXT_GENERATE_RESULT_COLUMN]], + ) + + +class PaLM2EmbeddingGenerator(bigframes.ml.base.Predictor): + """PaLM2 embedding generator LLM model. + + Args: + session: BQ session to create the model + connection_name: connection to connect with remote service. str of the format ..""" + + def __init__(self, session: bigframes.Session, connection_name: str): + self.session = session + self.connection_name = connection_name + self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + + def _create_bqml_model(self): + options = { + "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, + } + + return bigframes.ml.core.create_bqml_remote_model( + session=self.session, connection_name=self.connection_name, options=options + ) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + """Predict the result from input DataFrame. + + Args: + X: Input DataFrame, which needs to contain a column with name "content". Only the column will be used as input. Content can include preamble, questions, suggestions, instructions, or examples. + + Returns: Output DataFrame with only 1 column as the output embedding results.""" + + # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models + if len(X.columns) != 1: + raise ValueError("Only support one column as input.") + + # BQML identified the column by name + col_label = cast(blocks.Label, X.columns[0]) + X = X.rename(columns={col_label: "content"}) + + options = { + "flatten_json_output": True, + } + df = self._bqml_model.embed_text(X, options) + return cast( + bigframes.dataframe.DataFrame, + df[[_EMBED_TEXT_RESULT_COLUMN]], + ) diff --git a/bigframes/ml/loader.py b/bigframes/ml/loader.py new file mode 100644 index 00000000000..cc43b47698e --- /dev/null +++ b/bigframes/ml/loader.py @@ -0,0 +1,74 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Union + +from google.cloud import bigquery + +if TYPE_CHECKING: + import bigframes + +import bigframes.ml.cluster +import bigframes.ml.decomposition +import bigframes.ml.ensemble +import bigframes.ml.forecasting +import bigframes.ml.linear_model + + +def from_bq( + session: bigframes.Session, model: bigquery.Model +) -> Union[ + bigframes.ml.decomposition.PCA, + bigframes.ml.cluster.KMeans, + bigframes.ml.linear_model.LinearRegression, + bigframes.ml.linear_model.LogisticRegression, + bigframes.ml.ensemble.XGBRegressor, + bigframes.ml.ensemble.XGBClassifier, + bigframes.ml.forecasting.ARIMAPlus, + bigframes.ml.ensemble.RandomForestRegressor, + bigframes.ml.ensemble.RandomForestClassifier, +]: + """Load a BQML model to BigQuery DataFrames ML. + + Args: + session: a BigQuery DataFrames session. + model: a BigQuery model. + + Returns: + A BigQuery DataFrames ML model object. + """ + if model.model_type == "LINEAR_REGRESSION": + return bigframes.ml.linear_model.LinearRegression._from_bq(session, model) + elif model.model_type == "KMEANS": + return bigframes.ml.cluster.KMeans._from_bq(session, model) + elif model.model_type == "PCA": + return bigframes.ml.decomposition.PCA._from_bq(session, model) + elif model.model_type == "LOGISTIC_REGRESSION": + return bigframes.ml.linear_model.LogisticRegression._from_bq(session, model) + elif model.model_type == "BOOSTED_TREE_REGRESSOR": + return bigframes.ml.ensemble.XGBRegressor._from_bq(session, model) + elif model.model_type == "BOOSTED_TREE_CLASSIFIER": + return bigframes.ml.ensemble.XGBClassifier._from_bq(session, model) + elif model.model_type == "ARIMA_PLUS": + return bigframes.ml.forecasting.ARIMAPlus._from_bq(session, model) + elif model.model_type == "RANDOM_FOREST_REGRESSOR": + return bigframes.ml.ensemble.RandomForestRegressor._from_bq(session, model) + elif model.model_type == "RANDOM_FOREST_CLASSIFIER": + return bigframes.ml.ensemble.RandomForestClassifier._from_bq(session, model) + else: + raise NotImplementedError( + f"Model type {model.model_type} is not yet supported by BigQuery DataFrames." + ) diff --git a/bigframes/ml/metrics.py b/bigframes/ml/metrics.py new file mode 100644 index 00000000000..861e1f02d10 --- /dev/null +++ b/bigframes/ml/metrics.py @@ -0,0 +1,371 @@ +# 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. + +"""Metrics functions for evaluating models. This module is styled after +Scikit-Learn's metrics module: https://scikit-learn.org/stable/modules/metrics.html""" + +import inspect +import typing +from typing import Tuple + +import numpy as np +import pandas as pd +import sklearn.metrics as sklearn_metrics # type: ignore + +import bigframes.core.blocks as blocks +import bigframes.pandas as bpd +import third_party.bigframes_vendored.sklearn.metrics._classification as vendored_mertics_classification +import third_party.bigframes_vendored.sklearn.metrics._ranking as vendored_mertics_ranking +import third_party.bigframes_vendored.sklearn.metrics._regression as vendored_metrics_regression + + +def r2_score( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, + force_finite=True, +) -> float: + # TODO(bmil): support multioutput + if len(y_true.columns) > 1 or len(y_pred.columns) > 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + y_true_series = typing.cast( + bpd.Series, y_true[typing.cast(str, y_true.columns.tolist()[0])] + ) + y_pred_series = typing.cast( + bpd.Series, y_pred[typing.cast(str, y_pred.columns.tolist()[0])] + ) + + # total sum of squares + # (dataframe, scalar) binops + # TODO(bmil): remove multiply by self when bigframes supports pow() + # TODO(tbergeron): These stats are eagerly evaluated. Move to lazy representation once scalar subqueries supported. + delta_from_mean = y_true_series - y_true_series.mean() + ss_total = (delta_from_mean * delta_from_mean).sum() + + # residual sum of squares + # (scalar, scalar) binops + # TODO(bmil): remove multiply by self when bigframes supports pow() + delta_from_pred = y_true_series - y_pred_series + ss_res = (delta_from_pred * delta_from_pred).sum() + + if force_finite and ss_total == 0: + return 0.0 if ss_res > 0 else 1.0 + + return 1 - (ss_res / ss_total) + + +r2_score.__doc__ = inspect.getdoc(vendored_metrics_regression.r2_score) + + +def accuracy_score( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, + normalize=True, +) -> float: + # TODO(ashleyxu): support sample_weight as the parameter + if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + y_true_series = typing.cast( + bpd.Series, y_true[typing.cast(str, y_true.columns.tolist()[0])] + ) + y_pred_series = typing.cast( + bpd.Series, y_pred[typing.cast(str, y_pred.columns.tolist()[0])] + ) + + # Compute accuracy for each possible representation + # TODO(ashleyxu): add multilabel classification support where y_type + # starts with "multilabel" + score = (y_true_series == y_pred_series).astype(pd.Int64Dtype()) + + if normalize: + return score.mean() + else: + return score.sum() + + +accuracy_score.__doc__ = inspect.getdoc(vendored_mertics_classification.accuracy_score) + + +def roc_curve( + y_true: bpd.DataFrame, + y_score: bpd.DataFrame, + drop_intermediate: bool = True, +) -> Tuple[bpd.Series, bpd.Series, bpd.Series]: + # TODO(bmil): Add multi-class support + # TODO(bmil): Add multi-label support + if len(y_true.columns) > 1 or len(y_score.columns) > 1: + raise NotImplementedError("Only binary classification is supported") + + # TODO(bmil): Implement drop_intermediate + if drop_intermediate: + raise NotImplementedError("drop_intermediate is not yet implemented") + + # TODO(bmil): remove this once bigframes supports the necessary operations + session = y_true._block.expr._session + pd_y_true = y_true.to_pandas() + pd_y_score = y_score.to_pandas() + + # We operate on rows, so, remove the index if there is one + # TODO(bmil): check that the indexes are equivalent before removing + pd_y_true = pd_y_true.reset_index(drop=True) + pd_y_score = pd_y_score.reset_index(drop=True) + + pd_df = pd.DataFrame( + { + "y_true": pd_y_true[pd_y_true.columns[0]], + "y_score": pd_y_score[pd_y_score.columns[0]], + } + ) + + total_positives = pd_df.y_true.sum() + total_negatives = len(pd_df) - total_positives + + pd_df = pd_df.sort_values(by="y_score", ascending=False) + pd_df["cum_tp"] = pd_df.y_true.cumsum() + pd_df["cum_fp"] = (~pd_df.y_true.astype(bool)).cumsum() + + # produce just one data point per y_score + pd_df = pd_df.groupby("y_score", as_index=False).last() + pd_df = pd_df.sort_values(by="y_score", ascending=False) + + pd_df["tpr"] = pd_df.cum_tp / total_positives + pd_df["fpr"] = pd_df.cum_fp / total_negatives + pd_df["thresholds"] = pd_df.y_score + + # sklearn includes an extra datapoint for the origin with threshold np.inf + pd_origin = pd.DataFrame({"tpr": [0.0], "fpr": [0.0], "thresholds": np.inf}) + pd_df = pd.concat([pd_origin, pd_df]) + + df = session.read_pandas(pd_df) + return df.fpr, df.tpr, df.thresholds + + +roc_curve.__doc__ = inspect.getdoc(vendored_mertics_ranking.roc_curve) + + +def roc_auc_score(y_true: bpd.DataFrame, y_score: bpd.DataFrame) -> float: + # TODO(bmil): Add multi-class support + # TODO(bmil): Add multi-label support + if len(y_true.columns) > 1 or len(y_score.columns) > 1: + raise NotImplementedError("Only binary classification is supported") + + fpr, tpr, _ = roc_curve(y_true, y_score, drop_intermediate=False) + + # TODO(bmil): remove this once bigframes supports the necessary operations + pd_fpr = fpr.compute() + pd_tpr = tpr.compute() + + # Use the trapezoid rule to compute the area under the ROC curve + width_diff = pd_fpr.diff().iloc[1:].reset_index(drop=True) + height_avg = (pd_tpr.iloc[:-1] + pd_tpr.iloc[1:].reset_index(drop=True)) / 2 + return (width_diff * height_avg).sum() + + +roc_auc_score.__doc__ = inspect.getdoc(vendored_mertics_ranking.roc_auc_score) + + +def auc( + x: bpd.DataFrame, + y: bpd.DataFrame, +) -> float: + if len(x.columns) != 1 or len(y.columns) != 1: + raise ValueError("Only 1-D data structure is supported") + + # TODO(b/286410053) Support ML exceptions and error handling. + auc = sklearn_metrics.auc(x.to_pandas(), y.to_pandas()) + return auc + + +auc.__doc__ = inspect.getdoc(vendored_mertics_ranking.auc) + + +def confusion_matrix( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, +) -> pd.DataFrame: + # TODO(ashleyxu): support labels and sample_weight parameters + # TODO(ashleyxu): support bpd.Series as input type + if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + y_true_column = typing.cast(blocks.Label, y_true.columns[0]) + y_pred_series = typing.cast( + bpd.Series, + y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], + ) + confusion_df = y_true.assign(y_pred=y_pred_series) + confusion_df = confusion_df.assign(dummy=0) + groupby_count = ( + confusion_df.groupby(by=[y_true_column, "y_pred"], as_index=False) + .count() + .to_pandas() + ) + + unique_values = sorted(set(groupby_count.y_true).union(set(groupby_count.y_pred))) + + confusion_matrix = pd.DataFrame( + 0, index=pd.Index(unique_values), columns=pd.Index(unique_values), dtype=int + ) + + # Loop through the result by rows and columns + for _, row in groupby_count.iterrows(): + y_true = row["y_true"] + y_pred = row["y_pred"] + count = row["dummy"] + confusion_matrix[y_pred][y_true] = count + + return confusion_matrix + + +confusion_matrix.__doc__ = inspect.getdoc( + vendored_mertics_classification.confusion_matrix +) + + +def recall_score( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, + average: str = "binary", +) -> pd.Series: + # TODO(ashleyxu): support more average type, default to "binary" + # TODO(ashleyxu): support bpd.Series as input type + if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + if average is not None: + raise NotImplementedError("Only average=None is supported") + + y_true_series = typing.cast( + bpd.Series, + y_true[typing.cast(blocks.Label, y_true.columns.tolist()[0])], + ) + y_pred_series = typing.cast( + bpd.Series, + y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], + ) + + is_accurate = y_true_series == y_pred_series + unique_labels = ( + bpd.concat([y_true_series, y_pred_series], join="outer") + .drop_duplicates() + .sort_values() + ) + index = unique_labels.to_list() + + recall = ( + is_accurate.groupby(y_true_series).sum() + / is_accurate.groupby(y_true_series).count() + ).compute() + + recall_score = pd.Series(0, index=index) + for i in recall_score.index: + recall_score.loc[i] = recall.loc[i] + + return recall_score + + +recall_score.__doc__ = inspect.getdoc(vendored_mertics_classification.recall_score) + + +def precision_score( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, + average: str = "binary", +) -> pd.Series: + # TODO(ashleyxu): support more average type, default to "binary" + # TODO(ashleyxu): support bpd.Series as input type + if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + if average is not None: + raise NotImplementedError("Only average=None is supported") + + y_true_series = typing.cast( + bpd.Series, + y_true[typing.cast(blocks.Label, y_true.columns.tolist()[0])], + ) + y_pred_series = typing.cast( + bpd.Series, + y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], + ) + + is_accurate = y_true_series == y_pred_series + unique_labels = ( + bpd.concat([y_true_series, y_pred_series], join="outer") + .drop_duplicates() + .sort_values() + ) + index = unique_labels.to_list() + + precision = ( + is_accurate.groupby(y_pred_series).sum() + / is_accurate.groupby(y_pred_series).count() + ).compute() + + precision_score = pd.Series(0, index=index) + for i in precision.index: + precision_score.loc[i] = precision.loc[i] + + return precision_score + + +precision_score.__doc__ = inspect.getdoc( + vendored_mertics_classification.precision_score +) + + +def f1_score( + y_true: bpd.DataFrame, + y_pred: bpd.DataFrame, + average: str = "binary", +) -> pd.Series: + # TODO(ashleyxu): support more average type, default to "binary" + # TODO(ashleyxu): support bpd.Series as input type + if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + raise NotImplementedError( + "Only one labels column, one predictions column is supported" + ) + + if average is not None: + raise NotImplementedError("Only average=None is supported") + + recall = recall_score(y_true, y_pred, average=None) + precision = precision_score(y_true, y_pred, average=None) + + f1_score = pd.Series(0, index=recall.index) + for index in recall.index: + if precision[index] + recall[index] != 0: + f1_score[index] = ( + 2 + * (precision[index] * recall[index]) + / (precision[index] + recall[index]) + ) + else: + f1_score[index] = 0 + + return f1_score + + +f1_score.__doc__ = inspect.getdoc(vendored_mertics_classification.f1_score) diff --git a/bigframes/ml/model_selection.py b/bigframes/ml/model_selection.py new file mode 100644 index 00000000000..73a2fd76a9a --- /dev/null +++ b/bigframes/ml/model_selection.py @@ -0,0 +1,86 @@ +# 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. + +"""Functions for test/train split and model tuning. This module is styled after +Scikit-Learn's model_selection module: +https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection""" + + +from typing import List, Union + +import bigframes +import bigframes.dataframe + + +def train_test_split( + *dataframes: bigframes.dataframe.DataFrame, + test_size: Union[float, None] = None, + train_size: Union[float, None] = None, + random_state: Union[int, None] = None, +) -> List[bigframes.dataframe.DataFrame]: + """Splits dataframes into random train and test subsets + + Args: + *dataframes: + A sequence of BigQuery DataFrames that can be joined on + their indexes + test_size: + The proportion of the dataset to include in the test split. If + None, this will default to the complement of train_size. If both + are none, it will be set to 0.25. + train_size: + The proportion of the dataset to include in the train split. If + None, this will default to the complement of test_size. + random_state: + A seed to use for randomly choosing the rows of the split. If not + set, a random split will be generated each time. + + Returns: + A list of BigQuery DataFrames. + """ + + # TODO(garrettwu): Scikit-Learn throws an error when the dataframes don't have the same + # number of rows. We probably want to do something similar. Now the implementation is based + # on index. We'll move to based on ordering first. + + if test_size is None: + if train_size is None: + test_size = 0.25 + else: + test_size = 1.0 - train_size + if train_size is None: + train_size = 1.0 - test_size + + if train_size <= 0.0 or train_size >= 1.0: + raise ValueError(f"train_size must be within (0.0, 1.0). But is {train_size}.") + + if test_size <= 0.0 or test_size >= 1.0: + raise ValueError(f"test_size must be within (0.0, 1.0). But is {test_size}.") + + if train_size + test_size > 1.0: + raise ValueError( + f"The sum of train_size and test_size exceeds 1.0. train_size: {train_size}. test_size: {test_size}" + ) + + results = dataframes[0]._split( + fracs=(train_size, test_size), random_state=random_state + ) + train_index = results[0].index + test_index = results[1].index + + results += [ + df.loc[index] for df in dataframes[1:] for index in (train_index, test_index) + ] + + return results diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py new file mode 100644 index 00000000000..022ba0148d7 --- /dev/null +++ b/bigframes/ml/pipeline.py @@ -0,0 +1,102 @@ +# 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. + +"""For composing estimators together. This module is styled after Scikit-Learn's +pipeline module: https://scikit-learn.org/stable/modules/pipeline.html""" + + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import bigframes +from bigframes.ml import base, cluster, compose, decomposition, preprocessing +import third_party.bigframes_vendored.sklearn.pipeline + + +class Pipeline( + third_party.bigframes_vendored.sklearn.pipeline.Pipeline, + base.BaseEstimator, +): + __doc__ = third_party.bigframes_vendored.sklearn.pipeline.Pipeline.__doc__ + + def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): + self.steps = steps + + if len(steps) != 2: + raise NotImplementedError( + "Currently only two step (transform, estimator) pipelines are supported" + ) + + transform, estimator = steps[0][1], steps[1][1] + if isinstance( + transform, + ( + compose.ColumnTransformer, + preprocessing.StandardScaler, + preprocessing.OneHotEncoder, + ), + ): + self._transform = transform + else: + raise NotImplementedError( + f"Transform {transform} is not yet supported by Pipeline" + ) + + if not isinstance( + estimator, + base.TrainablePredictor, + ): + raise NotImplementedError( + f"Estimator {estimator} is not supported by Pipeline" + ) + + self._transform = transform + self._estimator = estimator + + def fit( + self, + X: bigframes.dataframe.DataFrame, + y: Optional[bigframes.dataframe.DataFrame] = None, + ): + compiled_transforms = self._transform._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + if y is not None: + # If labels columns are present, they should pass through un-transformed + transform_sqls.extend(y.columns.tolist()) + + self._estimator.fit(X=X, y=y, transforms=transform_sqls) + + def predict( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + return self._estimator.predict(X) + + def score( + self, + X: bigframes.dataframe.DataFrame, + y: bigframes.dataframe.DataFrame, + ): + if isinstance(self._estimator, (cluster.KMeans, decomposition.PCA)): + raise NotImplementedError("KMeans/PCA haven't supported score method.") + + # TODO(b/289280565): remove type ignore after updating KMeans and PCA + return self._estimator.score(X=X, y=y) # type: ignore + + def to_gbq(self, model_name: str, replace: bool = False): + self._estimator.to_gbq(model_name, replace) + + # TODO: should instead load from GBQ, but loading pipelines is not implemented yet + return self diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py new file mode 100644 index 00000000000..59d2cc2ae9e --- /dev/null +++ b/bigframes/ml/preprocessing.py @@ -0,0 +1,140 @@ +# 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. + +"""Transformers that prepare data for other estimators. This module is styled after +Scikit-Learn's preprocessing module: https://scikit-learn.org/stable/modules/preprocessing.html""" + + +import typing +from typing import List, Optional, Tuple + +import bigframes +import bigframes.ml +import bigframes.ml.sql +import third_party.bigframes_vendored.sklearn.preprocessing._data +import third_party.bigframes_vendored.sklearn.preprocessing._encoder + + +class StandardScaler( + third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, + bigframes.ml.base.BaseEstimator, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler.__doc__ + ) + + def __init__(self): + self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + ( + bigframes.ml.sql.ml_standard_scaler(column, f"scaled_{column}"), + f"scaled_{column}", + ) + for column in columns + ] + + def fit( + self, + X: bigframes.dataframe.DataFrame, + ): + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + + def transform( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + df = self._bqml_model.transform(X) + return typing.cast( + bigframes.dataframe.DataFrame, + df[self._output_names], + ) + + +class OneHotEncoder( + third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, + bigframes.ml.base.BaseEstimator, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder.__doc__ + ) + + # All estimators must implement __init__ to document their parameters, even + # if they don't have any + def __init__(self): + pass + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + ( + bigframes.ml.sql.ml_one_hot_encoder(column, f"onehotencoded_{column}"), + f"onehotencoded_{column}", + ) + for column in columns + ] + + def fit( + self, + X: bigframes.dataframe.DataFrame, + ): + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = bigframes.ml.core.create_bqml_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + + def transform( + self, X: bigframes.dataframe.DataFrame + ) -> bigframes.dataframe.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + df = self._bqml_model.transform(X) + return typing.cast( + bigframes.dataframe.DataFrame, + df[self._output_names], + ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py new file mode 100644 index 00000000000..7017b2a4cd1 --- /dev/null +++ b/bigframes/ml/sql.py @@ -0,0 +1,173 @@ +# 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. + +""" +Generates SQL queries needed for BigQuery DataFrames ML +""" + +from typing import Iterable, Optional, Union + + +def _encode_value(v: Union[str, int, float, Iterable[str]]) -> str: + """Encode a parameter value for SQL""" + if isinstance(v, str): + return f'"{v}"' + elif isinstance(v, int) or isinstance(v, float): + return f"{v}" + elif isinstance(v, Iterable): + inner = ", ".join([_encode_value(x) for x in v]) + return f"[{inner}]" + else: + raise ValueError("Unexpected value type") + + +def _build_param_Iterable(**kwargs: Union[str, int, float, Iterable[str]]) -> str: + """Encode a dict of values into a formatted Iterable of KVPs for SQL""" + indent_str = " " + param_strs = [f"{k}={_encode_value(v)}" for k, v in kwargs.items()] + return "\n" + indent_str + f",\n{indent_str}".join(param_strs) + + +def options(**kwargs: Union[str, int, float, Iterable[str]]) -> str: + """Encode the OPTIONS clause for BQML""" + return f"OPTIONS({_build_param_Iterable(**kwargs)})" + + +def _build_struct_param_Iterable(**kwargs: Union[int, float]) -> str: + """Encode a dict of values into a formatted STRUCT items for SQL""" + indent_str = " " + param_strs = [f"{v} AS {k}" for k, v in kwargs.items()] + return "\n" + indent_str + f",\n{indent_str}".join(param_strs) + + +def struct_options(**kwargs: Union[int, float]) -> str: + """Encode a BQ STRUCT as options.""" + return f"STRUCT({_build_struct_param_Iterable(**kwargs)})" + + +def _build_expr_Iterable(*expr_sqls: str) -> str: + """Encode a Iterable of SQL expressions into a formatted Iterable for SQL""" + indent_str = " " + return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls) + + +def transform(*expr_sqls: str) -> str: + """Encode the TRANSFORM clause for BQML""" + return f"TRANSFORM({_build_expr_Iterable(*expr_sqls)})" + + +def connection(conn_name: str) -> str: + """Encode the REMOTE WITH CONNECTION clause for BQML. conn_name is of the format ...""" + return f"REMOTE WITH CONNECTION `{conn_name}`" + + +def ml_standard_scaler(numeric_expr_sql: str, name: str) -> str: + """Encode ML.STANDARD_SCALER for BQML""" + return f"""ML.STANDARD_SCALER({numeric_expr_sql}) OVER() AS {name}""" + + +def ml_one_hot_encoder(numeric_expr_sql: str, name: str) -> str: + """Encode ML.ONE_HOT_ENCODER for BQML""" + return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}) OVER() AS {name}""" + + +def create_model( + model_name: str, + source_sql: str, + transform_sql: Optional[str] = None, + options_sql: Optional[str] = None, +) -> str: + """Encode the CREATE MODEL statement for BQML""" + # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed + parts = [f"CREATE MODEL `{model_name}`"] + if transform_sql: + parts.append(transform_sql) + if options_sql: + parts.append(options_sql) + parts.append(f"AS {source_sql}") + return "\n".join(parts) + + +def create_remote_model( + model_name: str, + connection_name: str, + options_sql: Optional[str] = None, +) -> str: + """Encode the CREATE MODEL statement for BQML""" + # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed + parts = [f"CREATE MODEL `{model_name}`"] + parts.append(connection(connection_name)) + if options_sql: + parts.append(options_sql) + return "\n".join(parts) + + +def create_imported_model( + model_name: str, + options_sql: Optional[str] = None, +) -> str: + """Encode the CREATE MODEL statement for BQML""" + # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed + parts = [f"CREATE MODEL `{model_name}`"] + if options_sql: + parts.append(options_sql) + return "\n".join(parts) + + +def alter_model( + model_name: str, + options_sql: str, +) -> str: + """Encode the ALTER MODEL statement for BQML""" + parts = [f"ALTER MODEL `{model_name}`"] + parts.append(f"SET {options_sql}") + return "\n".join(parts) + + +def ml_evaluate(model_name: str, source_sql: Union[str, None] = None) -> str: + """Encode ML.EVALUATE for BQML""" + if source_sql is None: + return f"""SELECT * FROM ML.EVALUATE(MODEL `{model_name}`)""" + else: + return f"""SELECT * FROM ML.EVALUATE(MODEL `{model_name}`, + ({source_sql}))""" + + +def ml_predict(model_name: str, source_sql: str) -> str: + """Encode ML.PREDICT for BQML""" + return f"""SELECT * FROM ML.PREDICT(MODEL `{model_name}`, + ({source_sql}))""" + + +def ml_transform(model_name: str, source_sql: str) -> str: + """Encode ML.TRANSFORM for BQML""" + return f"""SELECT * FROM ML.TRANSFORM(MODEL `{model_name}`, + ({source_sql}))""" + + +def ml_generate_text(model_name: str, source_sql: str, struct_options: str) -> str: + """Encode ML.GENERATE_TEXT for BQML""" + return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL `{model_name}`, + ({source_sql}), {struct_options})""" + + +def ml_embed_text(model_name: str, source_sql: str, struct_options: str) -> str: + """Encode ML.EMBED_TEXT for BQML""" + return f"""SELECT * FROM ML.EMBED_TEXT(MODEL `{model_name}`, + ({source_sql}), {struct_options})""" + + +def ml_forecast(model_name: str) -> str: + """Encode ML.FORECAST for BQML""" + return f"""SELECT * FROM ML.FORECAST(MODEL `{model_name}`)""" diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py new file mode 100644 index 00000000000..781cc493397 --- /dev/null +++ b/bigframes/operations/__init__.py @@ -0,0 +1,546 @@ +# 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. + +from __future__ import annotations + +import functools +import typing + +import ibis +import ibis.common.exceptions +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.operations.generic +import ibis.expr.types as ibis_types +import numpy as np + +import bigframes.dtypes +import bigframes.dtypes as dtypes + +_ZERO = typing.cast(ibis_types.NumericValue, ibis_types.literal(0)) +_INF = typing.cast(ibis_types.NumericValue, ibis_types.literal(np.inf)) + +BinaryOp = typing.Callable[[ibis_types.Value, ibis_types.Value], ibis_types.Value] +TernaryOp = typing.Callable[ + [ibis_types.Value, ibis_types.Value, ibis_types.Value], ibis_types.Value +] + + +### Unary Ops +class UnaryOp: + def _as_ibis(self, x): + raise NotImplementedError("Base class UnaryOp has no implementation.") + + @property + def is_windowed(self): + return False + + +class AbsOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).abs() + + +class InvertOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).negate() + + +class IsNullOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return x.isnull() + + +class LenOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).length() + + +class NotNullOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return x.notnull() + + +class ReverseOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).reverse() + + +class LowerOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).lower() + + +class UpperOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).upper() + + +class StripOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).strip() + + +class IsNumericOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + # catches all members of the Unicode number class, which matches pandas isnumeric + # see https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#regexp_contains + return typing.cast(ibis_types.StringValue, x).re_search(r"^(\pN*)$") + + +class RstripOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).rstrip() + + +class LstripOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).lstrip() + + +class CapitalizeOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).capitalize() + + +class HashOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.IntegerValue, x).hash() + + +class DayOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).day() + + +class DateOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).date() + + +class DayofweekOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).day_of_week.index() + + +class HourOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).hour() + + +class MinuteOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).minute() + + +class MonthOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).month() + + +class QuarterOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).quarter() + + +class SecondOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).second() + + +class TimeOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).time() + + +class YearOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.TimestampValue, x).year() + + +# Parameterized ops +class AsTypeOp(UnaryOp): + def __init__(self, to_type: dtypes.DtypeString | dtypes.Dtype): + self.to_type = bigframes.dtypes.bigframes_dtype_to_ibis_dtype(to_type) + + def _as_ibis(self, x: ibis_types.Value): + if isinstance(x, ibis_types.NullScalar): + return ibis_types.null().cast(self.to_type) + + return bigframes.dtypes.cast_ibis_value(x, self.to_type) + + +class FindOp(UnaryOp): + def __init__(self, sub, start, end): + self._sub = sub + self._start = start + self._end = end + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).find( + self._sub, self._start, self._end + ) + + +class SliceOp(UnaryOp): + def __init__(self, start, stop): + self._start = start + self._stop = stop + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x)[self._start : self._stop] + + +class BinopPartialRight(UnaryOp): + def __init__(self, binop: BinaryOp, right_scalar: typing.Any): + self._binop = binop + self._right = dtypes.literal_to_ibis_scalar(right_scalar, validate=False) + + def _as_ibis(self, x): + return self._binop(x, self._right) + + +class BinopPartialLeft(UnaryOp): + def __init__(self, binop: BinaryOp, left_scalar: typing.Any): + self._binop = binop + self._left = dtypes.literal_to_ibis_scalar(left_scalar, validate=False) + + def _as_ibis(self, x): + return self._binop(self._left, x) + + +class RepeatOp(UnaryOp): + def __init__(self, repeats): + self._repeats = repeats + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).repeat(self._repeats) + + +class RemoteFunctionOp(UnaryOp): + def __init__(self, func: typing.Callable, apply_on_null=True): + if not hasattr(func, "bigframes_remote_function"): + raise TypeError( + "only a bigframes remote function is supported as a callable" + ) + + self._func = func + self._apply_on_null = apply_on_null + + def _as_ibis(self, x: ibis_types.Value): + x_transformed = self._func(x) + if not self._apply_on_null: + x_transformed = where_op(x, x.isnull(), x_transformed) + return x_transformed + + +abs_op = AbsOp() +invert_op = InvertOp() +isnull_op = IsNullOp() +len_op = LenOp() +notnull_op = NotNullOp() +reverse_op = ReverseOp() +lower_op = LowerOp() +upper_op = UpperOp() +strip_op = StripOp() +isnumeric_op = IsNumericOp() +rstrip_op = RstripOp() +lstrip_op = LstripOp() +hash_op = HashOp() +day_op = DayOp() +dayofweek_op = DayofweekOp() +date_op = DateOp() +hour_op = HourOp() +minute_op = MinuteOp() +month_op = MonthOp() +quarter_op = QuarterOp() +second_op = SecondOp() +time_op = TimeOp() +year_op = YearOp() +capitalize_op = CapitalizeOp() + + +### Binary Ops +def short_circuit_nulls(type_override: typing.Optional[ibis_dtypes.DataType] = None): + """Wraps a binary operator to generate nulls of the expected type if either input is a null scalar.""" + + def short_circuit_nulls_inner(binop): + @functools.wraps(binop) + def wrapped_binop(x: ibis_types.Value, y: ibis_types.Value): + if isinstance(x, ibis_types.NullScalar): + return ibis_types.null().cast(type_override or y.type()) + elif isinstance(y, ibis_types.NullScalar): + return ibis_types.null().cast(type_override or x.type()) + else: + return binop(x, y) + + return wrapped_binop + + return short_circuit_nulls_inner + + +def concat_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + x_string = typing.cast(ibis_types.StringValue, x) + y_string = typing.cast(ibis_types.StringValue, y) + return x_string.concat(y_string) + + +def eq_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x == y + + +def ne_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x != y + + +def and_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.BooleanValue, x) & typing.cast( + ibis_types.BooleanValue, y + ) + + +def or_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.BooleanValue, x) | typing.cast( + ibis_types.BooleanValue, y + ) + + +def isin_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x.isin(y) + + +@short_circuit_nulls() +def add_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + if isinstance(x, ibis_types.NullScalar) or isinstance(x, ibis_types.NullScalar): + return + return typing.cast(ibis_types.NumericValue, x) + typing.cast( + ibis_types.NumericValue, y + ) + + +@short_circuit_nulls() +def sub_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.NumericValue, x) - typing.cast( + ibis_types.NumericValue, y + ) + + +@short_circuit_nulls() +def mul_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.NumericValue, x) * typing.cast( + ibis_types.NumericValue, y + ) + + +@short_circuit_nulls(ibis_dtypes.float) +def div_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.NumericValue, x) / typing.cast( + ibis_types.NumericValue, y + ) + + +@short_circuit_nulls(ibis_dtypes.bool) +def lt_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x < y + + +@short_circuit_nulls(ibis_dtypes.bool) +def le_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x <= y + + +@short_circuit_nulls(ibis_dtypes.bool) +def gt_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x > y + + +@short_circuit_nulls(ibis_dtypes.bool) +def ge_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x >= y + + +@short_circuit_nulls(ibis_dtypes.int) +def floordiv_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + x_numeric = typing.cast(ibis_types.NumericValue, x) + y_numeric = typing.cast(ibis_types.NumericValue, y) + floordiv_expr = x_numeric // y_numeric + + # DIV(N, 0) will error in bigquery, but needs to return 0 for int, and inf for float in BQ so we short-circuit in this case. + # Multiplying left by zero propogates nulls. + zero_result = _INF if (x.type().is_floating() or y.type().is_floating()) else _ZERO + return ( + ibis.case() + .when(y_numeric == _ZERO, zero_result * x_numeric) + .else_(floordiv_expr) + .end() + ) + + +@short_circuit_nulls() +def mod_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + # TODO(tbergeron): fully support floats, including when mixed with integer + # Pandas has inconsitency about whether N mod 0. Most conventions have this be NAN. + # For some dtypes, the result is 0 instead. This implementation results in NA always. + x_numeric = typing.cast(ibis_types.NumericValue, x) + y_numeric = typing.cast(ibis_types.NumericValue, y) + # Hacky short-circuit to avoid passing zero-literal to sql backend, evaluate locally instead to null. + op = y.op() + if isinstance(op, ibis.expr.operations.generic.Literal) and op.value == 0: + return ibis_types.null().cast(x.type()) + + bq_mod = x_numeric % y_numeric # Bigquery will maintain x sign here + # In BigQuery returned value has the same sign as X. In pandas, the sign of y is used, so we need to flip the result if sign(x) != sign(y) + return ( + ibis.case() + .when( + y_numeric == _ZERO, _ZERO * x_numeric + ) # Dummy op to propogate nulls and type from x arg + .when( + (y_numeric < _ZERO) & (bq_mod > _ZERO), (y_numeric + bq_mod) + ) # Convert positive result to negative + .when( + (y_numeric > _ZERO) & (bq_mod < _ZERO), (y_numeric + bq_mod) + ) # Convert negative result to positive + .else_(bq_mod) + .end() + ) + + +def fillna_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return x.fillna(typing.cast(ibis_types.Scalar, y)) + + +def clip_lower( + value: ibis_types.Value, + lower: ibis_types.Value, +): + return ibis.case().when(lower.isnull() | (value < lower), lower).else_(value).end() + + +def clip_upper( + value: ibis_types.Value, + upper: ibis_types.Value, +): + return ibis.case().when(upper.isnull() | (value > upper), upper).else_(value).end() + + +def reverse(op: BinaryOp) -> BinaryOp: + return lambda x, y: op(y, x) + + +def partial_left(op: BinaryOp, scalar: typing.Any) -> UnaryOp: + return BinopPartialLeft(op, scalar) + + +def partial_right(op: BinaryOp, scalar: typing.Any) -> UnaryOp: + return BinopPartialRight(op, scalar) + + +# Ternary ops +def where_op( + original: ibis_types.Value, + condition: ibis_types.Value, + replacement: ibis_types.Value, +) -> ibis_types.Value: + """Returns x if y is true, otherwise returns z.""" + return ibis.case().when(condition, original).else_(replacement).end() + + +def clip_op( + original: ibis_types.Value, + lower: ibis_types.Value, + upper: ibis_types.Value, +) -> ibis_types.Value: + """Clips value to lower and upper bounds.""" + if isinstance(lower, ibis_types.NullScalar) and ( + not isinstance(upper, ibis_types.NullScalar) + ): + return ( + ibis.case() + .when(upper.isnull() | (original > upper), upper) + .else_(original) + .end() + ) + elif (not isinstance(lower, ibis_types.NullScalar)) and isinstance( + upper, ibis_types.NullScalar + ): + return ( + ibis.case() + .when(lower.isnull() | (original < lower), lower) + .else_(original) + .end() + ) + elif isinstance(lower, ibis_types.NullScalar) and ( + isinstance(upper, ibis_types.NullScalar) + ): + return original + else: + # Note: Pandas has unchanged behavior when upper bound and lower bound are flipped. This implementation requires that lower_bound < upper_bound + return ( + ibis.case() + .when(lower.isnull() | (original < lower), lower) + .when(upper.isnull() | (original > upper), upper) + .else_(original) + .end() + ) diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py new file mode 100644 index 00000000000..028ed4b6065 --- /dev/null +++ b/bigframes/operations/aggregations.py @@ -0,0 +1,288 @@ +# 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. + +from __future__ import annotations + +import typing + +import ibis +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types + + +class WindowOp: + def _as_ibis(self, value: ibis_types.Column, window=None): + raise NotImplementedError("Base class WindowOp has no implementaiton.") + + @property + def skips_nulls(self): + """Whether the window op skips null rows.""" + return True + + @property + def handles_ties(self): + """Whether the operator can handle ties without nondeterministic output. (eg. rank operator can handle ties but not the count operator)""" + return False + + +class AggregateOp(WindowOp): + def _as_ibis(self, value: ibis_types.Column, window=None): + raise NotImplementedError("Base class AggregateOp has no implementaiton.") + + +def numeric_op(operation): + def constrained_op(op, column: ibis_types.Column, window=None): + if column.type().is_boolean(): + column = typing.cast( + ibis_types.NumericColumn, column.cast(ibis_dtypes.int64) + ) + if column.type().is_numeric(): + return operation(op, column, window) + else: + raise ValueError( + f"Numeric operation cannot be applied to type {column.type()}" + ) + + return constrained_op + + +class SumOp(AggregateOp): + @numeric_op + def _as_ibis( + self, column: ibis_types.NumericColumn, window=None + ) -> ibis_types.NumericValue: + # Will be null if all inputs are null. Pandas defaults to zero sum though. + bq_sum = _apply_window_if_present(column.sum(), window) + return ( + ibis.case().when(bq_sum.isnull(), ibis_types.literal(0)).else_(bq_sum).end() + ) + + +class MeanOp(AggregateOp): + @numeric_op + def _as_ibis( + self, column: ibis_types.NumericColumn, window=None + ) -> ibis_types.NumericValue: + return _apply_window_if_present(column.mean(), window) + + +class ProductOp(AggregateOp): + @numeric_op + def _as_ibis( + self, column: ibis_types.NumericColumn, window=None + ) -> ibis_types.NumericValue: + # Need to short-circuit as log with zeroes is illegal sql + is_zero = typing.cast(ibis_types.BooleanColumn, (column == 0)) + + # There is no product sql aggregate function, so must implement as a sum of logs, and then + # apply power after. Note, log and power base must be equal! This impl uses base 2. + logs = typing.cast( + ibis_types.NumericColumn, + ibis.case().when(is_zero, 0).else_(column.abs().log2()).end(), + ) + logs_sum = _apply_window_if_present(logs.sum(), window) + magnitude = typing.cast(ibis_types.NumericValue, ibis_types.literal(2)).pow( + logs_sum + ) + + # Can't determine sign from logs, so have to determine parity of count of negative inputs + is_negative = typing.cast( + ibis_types.NumericColumn, + ibis.case().when(column.sign() == -1, 1).else_(0).end(), + ) + negative_count = _apply_window_if_present(is_negative.sum(), window) + negative_count_parity = negative_count % typing.cast( + ibis_types.NumericValue, ibis.literal(2) + ) # 1 if result should be negative, otherwise 0 + + any_zeroes = _apply_window_if_present(is_zero.any(), window) + float_result = ( + ibis.case() + .when(any_zeroes, ibis_types.literal(0)) + .else_(magnitude * pow(-1, negative_count_parity)) + .end() + ) + return float_result.cast(column.type()) + + +class MaxOp(AggregateOp): + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present(column.max(), window) + + +class MinOp(AggregateOp): + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present(column.min(), window) + + +class StdOp(AggregateOp): + @numeric_op + def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present( + typing.cast(ibis_types.NumericColumn, x).std(), window + ) + + +class VarOp(AggregateOp): + @numeric_op + def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present( + typing.cast(ibis_types.NumericColumn, x).var(), window + ) + + +class CountOp(AggregateOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + return _apply_window_if_present(column.count(), window) + + @property + def skips_nulls(self): + return False + + +class NuniqueOp(AggregateOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + return _apply_window_if_present(column.nunique(), window) + + @property + def skips_nulls(self): + return False + + +class RankOp(WindowOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + # Ibis produces 0-based ranks, while pandas creates 1-based ranks + return _apply_window_if_present(column.rank(), window) + 1 + + @property + def skips_nulls(self): + return False + + @property + def handles_ties(self): + return True + + +class DenseRankOp(WindowOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + # Ibis produces 0-based ranks, while pandas creates 1-based ranks + return _apply_window_if_present(column.dense_rank(), window) + 1 + + @property + def skips_nulls(self): + return False + + @property + def handles_ties(self): + return True + + +class FirstOp(WindowOp): + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present(column.first(), window) + + +class ShiftOp(WindowOp): + def __init__(self, periods: int): + self._periods = periods + + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + if self._periods == 0: # No-op + return column + if self._periods > 0: + return _apply_window_if_present(column.lag(self._periods), window) + return _apply_window_if_present(column.lead(-self._periods), window) + + @property + def skips_nulls(self): + return False + + +class AllOp(AggregateOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.BooleanValue: + # BQ will return null for empty column, result would be true in pandas. + result = _is_true(column).all() + return typing.cast( + ibis_types.BooleanScalar, + _apply_window_if_present(result, window).fillna(ibis_types.literal(True)), + ) + + +class AnyOp(AggregateOp): + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.BooleanValue: + # BQ will return null for empty column, result would be false in pandas. + result = _is_true(column).any() + return typing.cast( + ibis_types.BooleanScalar, + _apply_window_if_present(result, window).fillna(ibis_types.literal(True)), + ) + + +def _is_true(column: ibis_types.Column) -> ibis_types.BooleanColumn: + if column.type().is_boolean(): + return typing.cast(ibis_types.BooleanColumn, column) + elif column.type().is_numeric(): + result = typing.cast(ibis_types.NumericColumn, column).__ne__( + ibis_types.literal(0) + ) + return typing.cast(ibis_types.BooleanColumn, result) + elif column.type().is_string(): + result = typing.cast( + ibis_types.StringValue, column + ).length() > ibis_types.literal(0) + return typing.cast(ibis_types.BooleanColumn, result) + else: + # Time and geo values don't have a 'False' value + return typing.cast( + ibis_types.BooleanColumn, _map_to_literal(column, ibis_types.literal(True)) + ) + + +def _apply_window_if_present(value: ibis_types.Value, window): + return value.over(window) if (window is not None) else value + + +def _map_to_literal( + original: ibis_types.Value, literal: ibis_types.Scalar +) -> ibis_types.Column: + # Hack required to perform aggregations on literals in ibis, even though bigquery will let you directly aggregate literals (eg. 'SELECT COUNT(1) from table1') + return ibis.ifelse(original.isnull(), literal, literal) + + +sum_op = SumOp() +mean_op = MeanOp() +product_op = ProductOp() +max_op = MaxOp() +min_op = MinOp() +std_op = StdOp() +var_op = VarOp() +count_op = CountOp() +nunique_op = NuniqueOp() +rank_op = RankOp() +dense_rank_op = DenseRankOp() +all_op = AllOp() +any_op = AnyOp() +first_op = FirstOp() diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py new file mode 100644 index 00000000000..29017d2bbe3 --- /dev/null +++ b/bigframes/operations/base.py @@ -0,0 +1,198 @@ +# 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. + +from __future__ import annotations + +import typing + +import ibis.expr.types as ibis_types +import pandas as pd + +import bigframes.core.blocks as blocks +import bigframes.core.scalar as scalars +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.series as series +import bigframes.session +import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing + +# BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. +# TODO(tbergeron): Convert to bytes-based limit +MAX_INLINE_SERIES_SIZE = 5000 + + +class SeriesMethods: + def __init__( + self, + data=None, + index: vendored_pandas_typing.Axes | None = None, + dtype: typing.Optional[ + bigframes.dtypes.DtypeString | bigframes.dtypes.Dtype + ] = None, + name: str | None = None, + copy: typing.Optional[bool] = None, + *, + session: typing.Optional[bigframes.session.Session] = None, + ): + block = None + if copy is not None and not copy: + raise ValueError("Series constructor only supports copy=True") + if isinstance(data, blocks.Block): + assert len(data.value_columns) == 1 + assert len(data.column_labels) == 1 + block = data + + elif isinstance(data, SeriesMethods): + block = data._get_block() + + if block: + if name: + if not isinstance(name, str): + raise NotImplementedError( + "BigQuery DataFrames only supports string series names." + ) + block = block.with_column_labels([name]) + if index: + raise NotImplementedError( + "Series 'index' constructor parameter not supported when passing BigQuery-backed objects" + ) + if dtype: + block = block.multi_apply_unary_op( + block.value_columns, ops.AsTypeOp(dtype) + ) + self._block = block + + else: + import bigframes.pandas + + pd_series = pd.Series( + data=data, index=index, dtype=dtype, name=name # type:ignore + ) + pd_dataframe = pd_series.to_frame() + if pd_series.name is None: + # to_frame will set default numeric column label if unnamed, but we do not support int column label, so must rename + pd_dataframe = pd_dataframe.set_axis(["unnamed_col"], axis=1) + if pd_dataframe.size < MAX_INLINE_SERIES_SIZE: + self._block = blocks.block_from_local( + pd_dataframe, session or bigframes.pandas.get_global_session() + ) + elif session: + self._block = session.read_pandas(pd_dataframe)._get_block() + else: + # Uses default global session + self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block() + if pd_series.name is None: + self._block = self._block.with_column_labels([None]) + + @property + def _value(self) -> ibis_types.Value: + """Private property to get Ibis expression for the value column.""" + return self._block.expr.get_column(self._value_column) + + @property + def _value_column(self) -> str: + return self._block.value_columns[0] + + @property + def _name(self) -> blocks.Label: + return self._block.column_labels[0] + + @property + def _dtype(self): + return self._block.dtypes[0] + + def _set_block(self, block: blocks.Block): + self._block = block + + def _get_block(self) -> blocks.Block: + return self._block + + def _apply_unary_op( + self, + op: ops.UnaryOp, + ) -> series.Series: + """Applies a unary operator to the series.""" + block, result_id = self._block.apply_unary_op( + self._value_column, op, result_label=self._name + ) + return series.Series(block.select_column(result_id)) + + def _apply_binary_op( + self, + other: typing.Any, + op: ops.BinaryOp, + alignment: typing.Literal["outer", "left"] = "outer", + ) -> series.Series: + """Applies a binary operator to the series and other.""" + if isinstance(other, pd.Series): + # TODO: Convert to BigQuery DataFrames series + raise NotImplementedError( + "Pandas series not supported supported as operand." + ) + if isinstance(other, series.Series): + (left, right, block) = self._align(other, how=alignment) + + block, result_id = block.apply_binary_op( + left, right, op, self._value_column + ) + + name = self._name + if ( + isinstance(other, series.Series) + and other.name != self._name + and alignment == "outer" + ): + name = None + + return series.Series( + block.select_column(result_id).assign_label(result_id, name) + ) + else: + partial_op = ops.BinopPartialRight(op, other) + return self._apply_unary_op(partial_op) + + def _align(self, other: series.Series, how="outer") -> tuple[str, str, blocks.Block]: # type: ignore + """Aligns the series value with another scalar or series object. Returns new left column id, right column id and joined tabled expression.""" + values, block = self._align_n( + [ + other, + ], + how, + ) + return (values[0], values[1], block) + + def _align_n( + self, + others: typing.Sequence[typing.Union[series.Series, scalars.Scalar]], + how="outer", + ) -> tuple[typing.Sequence[str], blocks.Block]: + value_ids = [self._value_column] + block = self._block + for other in others: + if isinstance(other, series.Series): + combined_index, ( + get_column_left, + get_column_right, + ) = block.index.join(other._block.index, how=how) + value_ids = [ + *[get_column_left(value) for value in value_ids], + get_column_right(other._value_column), + ] + block = combined_index._block + else: + # Will throw if can't interpret as scalar. + dtype = typing.cast(bigframes.dtypes.Dtype, self._dtype) + block, constant_col_id = block.create_constant(other, dtype=dtype) + value_ids = [*value_ids, constant_col_id] + return (value_ids, block) diff --git a/bigframes/operations/datetimes.py b/bigframes/operations/datetimes.py new file mode 100644 index 00000000000..1b20c2d593d --- /dev/null +++ b/bigframes/operations/datetimes.py @@ -0,0 +1,66 @@ +# 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. + +from __future__ import annotations + +import bigframes.operations as ops +import bigframes.operations.base +import bigframes.series as series +import third_party.bigframes_vendored.pandas.core.indexes.accessor as vendordt + + +class DatetimeMethods( + bigframes.operations.base.SeriesMethods, vendordt.DatetimeProperties +): + __doc__ = vendordt.DatetimeProperties.__doc__ + + @property + def day(self) -> series.Series: + return self._apply_unary_op(ops.day_op) + + @property + def dayofweek(self) -> series.Series: + return self._apply_unary_op(ops.dayofweek_op) + + @property + def date(self) -> series.Series: + return self._apply_unary_op(ops.date_op) + + @property + def hour(self) -> series.Series: + return self._apply_unary_op(ops.hour_op) + + @property + def minute(self) -> series.Series: + return self._apply_unary_op(ops.minute_op) + + @property + def month(self) -> series.Series: + return self._apply_unary_op(ops.month_op) + + @property + def second(self) -> series.Series: + return self._apply_unary_op(ops.second_op) + + @property + def time(self) -> series.Series: + return self._apply_unary_op(ops.time_op) + + @property + def quarter(self) -> series.Series: + return self._apply_unary_op(ops.quarter_op) + + @property + def year(self) -> series.Series: + return self._apply_unary_op(ops.year_op) diff --git a/bigframes/operations/strings.py b/bigframes/operations/strings.py new file mode 100644 index 00000000000..57029ede416 --- /dev/null +++ b/bigframes/operations/strings.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. + +from __future__ import annotations + +from typing import Literal, Optional, Union + +import bigframes.operations as ops +import bigframes.operations.base +import bigframes.series as series +import third_party.bigframes_vendored.pandas.core.strings.accessor as vendorstr + + +class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods): + __doc__ = vendorstr.StringMethods.__doc__ + + def find( + self, + sub: str, + start: Optional[int] = None, + end: Optional[int] = None, + ) -> series.Series: + return self._apply_unary_op(ops.FindOp(sub, start, end)) + + def len(self) -> series.Series: + return self._apply_unary_op(ops.len_op) + + def lower(self) -> series.Series: + return self._apply_unary_op(ops.lower_op) + + def reverse(self) -> series.Series: + """Reverse strings in the Series.""" + # reverse method is in ibis, not pandas. + return self._apply_unary_op(ops.reverse_op) + + def slice( + self, + start: Optional[int] = None, + stop: Optional[int] = None, + ) -> series.Series: + return self._apply_unary_op(ops.SliceOp(start, stop)) + + def strip(self) -> series.Series: + return self._apply_unary_op(ops.strip_op) + + def upper(self) -> series.Series: + return self._apply_unary_op(ops.upper_op) + + def isnumeric(self) -> series.Series: + return self._apply_unary_op(ops.isnumeric_op) + + def rstrip(self) -> series.Series: + return self._apply_unary_op(ops.rstrip_op) + + def lstrip(self) -> series.Series: + return self._apply_unary_op(ops.lstrip_op) + + def repeat(self, repeats: int) -> series.Series: + return self._apply_unary_op(ops.RepeatOp(repeats)) + + def capitalize(self) -> series.Series: + return self._apply_unary_op(ops.capitalize_op) + + def cat( + self, + others: Union[str, series.Series], + *, + join: Literal["outer", "left"] = "left", + ) -> series.Series: + return self._apply_binary_op(others, ops.concat_op, alignment=join) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py new file mode 100644 index 00000000000..e7c14351513 --- /dev/null +++ b/bigframes/pandas/__init__.py @@ -0,0 +1,339 @@ +# 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. + +"""BigQuery DataFrames provides a DataFrame API backed by the BigQuery engine.""" + +from __future__ import annotations + +import inspect +import threading +import typing +from typing import ( + Any, + Callable, + Dict, + IO, + Iterable, + List, + Literal, + MutableSequence, + Optional, + Sequence, + Tuple, + TypeVar, + Union, +) + +from google.cloud import bigquery +import numpy +import pandas + +import bigframes._config as config +import bigframes.core.indexes +import bigframes.core.reshape +import bigframes.dataframe +import bigframes.series +import bigframes.session +import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat + + +# Include method definition so that the method appears in our docs for +# bigframes.pandas general functions. +@typing.overload +def concat( + objs: Iterable[bigframes.dataframe.DataFrame], *, join, ignore_index +) -> bigframes.dataframe.DataFrame: + ... + + +@typing.overload +def concat( + objs: Iterable[bigframes.series.Series], *, join, ignore_index +) -> bigframes.series.Series: + ... + + +def concat( + objs: Union[ + Iterable[bigframes.dataframe.DataFrame], Iterable[bigframes.series.Series] + ], + *, + join: Literal["inner", "outer"] = "outer", + ignore_index: bool = False, +) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: + return bigframes.core.reshape.concat( + objs=objs, join=join, ignore_index=ignore_index + ) + + +concat.__doc__ = vendored_pandas_concat.concat.__doc__ + + +options = config.options +"""Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames.""" + +_global_session: Optional[bigframes.session.Session] = None +_global_session_lock = threading.Lock() + + +def reset_session() -> None: + """Start a fresh session next time a function requires a session. + + Closes the current session if it was already started. + """ + global _global_session + + with _global_session_lock: + if _global_session is not None: + _global_session.close() + _global_session = None + + options.bigquery._session_started = False + + +def get_global_session(): + """Gets the global session. + + Creates the global session if it does not exist. + """ + global _global_session, _global_session_lock + + with _global_session_lock: + if _global_session is None: + _global_session = bigframes.session.connect(options.bigquery) + + return _global_session + + +_T = TypeVar("_T") + + +def _with_default_session(func: Callable[..., _T], *args, **kwargs) -> _T: + return func(get_global_session(), *args, **kwargs) + + +def _set_default_session_location_if_possible(query): + # If the default session has not started yet and this is the first API user + # is calling, then set the default location as per the query. + # If query is a table name, then it would be the location of the table. + # If query is a SQL with a table, then it would be table's location. + # If query is a SQL with no table, then it would be the BQ default location. + if options.bigquery._session_started or options.bigquery.use_regional_endpoints: + return + + bqclient, _, _ = bigframes.session._create_bq_clients( + project=options.bigquery.project, + location=options.bigquery.location, + use_regional_endpoints=options.bigquery.use_regional_endpoints, + credentials=options.bigquery.credentials, + ) + + if bigframes.session._is_query(query): + job = bqclient.query(query, bigquery.QueryJobConfig(dry_run=True)) + options.bigquery.location = job.location + else: + table = bqclient.get_table(query) + options.bigquery.location = table.location + + +# Note: the following methods are duplicated from Session. This duplication +# enables the following: +# +# 1. Static type checking knows the argument and return types, which is +# difficult to do with decorators. Aside: When we require Python 3.10, we +# can use Concatenate for generic typing in decorators. See: +# https://stackoverflow.com/a/68290080/101923 +# 2. docstrings get processed by static processing tools, such as VS Code's +# autocomplete. +# 3. Positional arguments function as expected. If we were to pull in the +# methods directly from Session, a Session object would need to be the first +# argument, even if we allow a default value. +# 4. Allows to set BigQuery options for the BigFrames session based on the +# method and its arguments. + + +def read_csv( + filepath_or_buffer: str | IO["bytes"], + *, + sep: Optional[str] = ",", + header: Optional[int] = 0, + names: Optional[ + Union[MutableSequence[Any], numpy.ndarray[Any, Any], Tuple[Any, ...], range] + ] = None, + index_col: Optional[ + Union[int, str, Sequence[Union[str, int]], Literal[False]] + ] = None, + usecols: Optional[ + Union[ + MutableSequence[str], + Tuple[str, ...], + Sequence[int], + pandas.Series, + pandas.Index, + numpy.ndarray[Any, Any], + Callable[[Any], bool], + ] + ] = None, + dtype: Optional[Dict] = None, + engine: Optional[ + Literal["c", "python", "pyarrow", "python-fwf", "bigquery"] + ] = None, + encoding: Optional[str] = None, + **kwargs, +) -> bigframes.dataframe.DataFrame: + return _with_default_session( + bigframes.session.Session.read_csv, + filepath_or_buffer=filepath_or_buffer, + sep=sep, + header=header, + names=names, + index_col=index_col, + usecols=usecols, + dtype=dtype, + engine=engine, + encoding=encoding, + **kwargs, + ) + + +read_csv.__doc__ = inspect.getdoc(bigframes.session.Session.read_csv) + + +def read_gbq( + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, +) -> bigframes.dataframe.DataFrame: + _set_default_session_location_if_possible(query) + return _with_default_session( + bigframes.session.Session.read_gbq, + query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + ) + + +read_gbq.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq) + + +def read_gbq_model(model_name: str): + return _with_default_session( + bigframes.session.Session.read_gbq_model, + model_name, + ) + + +read_gbq_model.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_model) + + +def read_gbq_query( + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, +) -> bigframes.dataframe.DataFrame: + _set_default_session_location_if_possible(query) + return _with_default_session( + bigframes.session.Session.read_gbq_query, + query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + ) + + +read_gbq_query.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_query) + + +def read_gbq_table( + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, +) -> bigframes.dataframe.DataFrame: + _set_default_session_location_if_possible(query) + return _with_default_session( + bigframes.session.Session.read_gbq_table, + query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + ) + + +read_gbq_table.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_table) + + +def read_pandas(pandas_dataframe: pandas.DataFrame) -> bigframes.dataframe.DataFrame: + return _with_default_session( + bigframes.session.Session.read_pandas, + pandas_dataframe, + ) + + +read_pandas.__doc__ = inspect.getdoc(bigframes.session.Session.read_pandas) + + +def read_parquet(path: str | IO["bytes"]) -> bigframes.dataframe.DataFrame: + return _with_default_session( + bigframes.session.Session.read_parquet, + path, + ) + + +read_parquet.__doc__ = inspect.getdoc(bigframes.session.Session.read_parquet) + + +def remote_function( + input_types: List[type], + output_type: type, + dataset: Optional[str] = None, + bigquery_connection: Optional[str] = None, + reuse: bool = True, +): + return _with_default_session( + bigframes.session.Session.remote_function, + input_types=input_types, + output_type=output_type, + dataset=dataset, + bigquery_connection=bigquery_connection, + reuse=reuse, + ) + + +remote_function.__doc__ = inspect.getdoc(bigframes.session.Session.remote_function) + + +# Other aliases +DataFrame = bigframes.dataframe.DataFrame +Index = bigframes.core.indexes.Index +Series = bigframes.series.Series + +# Use __all__ to let type checkers know what is part of the public API. +__all___ = [ + "concat", + "DataFrame", + "options", + "read_csv", + "read_gbq", + "read_gbq_model", + "read_pandas", + "remote_function", + "Series", +] diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py new file mode 100644 index 00000000000..3fe4e38d785 --- /dev/null +++ b/bigframes/remote_function.py @@ -0,0 +1,697 @@ +# 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. + +from __future__ import annotations + +import functools +import hashlib +import inspect +import logging +import os +import random +import shutil +import string +import subprocess +import sys +import tempfile +import textwrap +import time +import typing + +if typing.TYPE_CHECKING: + from bigframes.session import Session + +import cloudpickle +import google.api_core.exceptions +from google.cloud import bigquery, bigquery_connection_v1, functions_v2 +from ibis.backends.bigquery.compiler import compiles +from ibis.backends.bigquery.datatypes import BigQueryType +from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type +import ibis.expr.operations as ops +import ibis.expr.rules as rlz + +# TODO(shobs): Change the min log level to INFO after the development stabilizes +# before June 2023 +logging.basicConfig( + level=logging.INFO, format="[%(levelname)s][%(asctime)s][%(name)s] %(message)s" +) +logger = logging.getLogger(__name__) + +# Protocol version 4 is available in python version 3.4 and above +# https://docs.python.org/3/library/pickle.html#data-stream-format +_pickle_protocol_version = 4 + +# Input and output python types supported by BigQuery DataFrames remote functions. +# TODO(shobs): Extend the support to all types supported by BQ remote functions +# https://cloud.google.com/bigquery/docs/remote-functions#limitations +_supported_io_types = set((bool, float, int, str)) + + +def get_remote_function_locations(bq_location): + """Get BQ location and cloud functions region given a BQ client.""" + # TODO(shobs, b/274647164): Find the best way to determine default location. + # For now let's assume that if no BQ location is set in the client then it + # defaults to US multi region + bq_location = bq_location.lower() if bq_location else "us" + + # Cloud function should be in the same region as the bigquery remote function + cloud_function_region = bq_location + + # BigQuery has multi region but cloud functions does not. + # Any region in the multi region that supports cloud functions should work + # https://cloud.google.com/functions/docs/locations + if bq_location == "us": + cloud_function_region = "us-central1" + elif bq_location == "eu": + cloud_function_region = "europe-west1" + + return bq_location, cloud_function_region + + +def _get_hash(def_): + "Get hash of a function." + def_repr = cloudpickle.dumps(def_, protocol=_pickle_protocol_version) + return hashlib.md5(def_repr).hexdigest() + + +def _run_system_command(command): + program = subprocess.Popen( + [command], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True + ) + stdout, stderr = program.communicate() + exit_code = program.wait() + if exit_code: + raise RuntimeError( + f"Command: {command}\nOutput: {stdout.decode()}\nError: {stderr.decode()}" + ) + + +def get_cloud_function_name(def_, uniq_suffix=None): + """Get the name of the cloud function.""" + cf_name = _get_hash(def_) + cf_name = f"bigframes-{cf_name}" # for identification + if uniq_suffix: + cf_name = f"{cf_name}-{uniq_suffix}" + return cf_name + + +def get_remote_function_name(def_, uniq_suffix=None): + """Get the name for the BQ remote function.""" + bq_rf_name = _get_hash(def_) + bq_rf_name = f"bigframes_{bq_rf_name}" # for identification + if uniq_suffix: + bq_rf_name = f"{bq_rf_name}_{uniq_suffix}" + return bq_rf_name + + +class RemoteFunctionClient: + # Wait time (in seconds) for an IAM binding to take effect after creation + _iam_wait_seconds = 120 + + def __init__( + self, + gcp_project_id, + cloud_function_region, + bq_location, + bq_dataset, + bq_client, + bq_connection_client, + bq_connection_id, + ): + self._gcp_project_id = gcp_project_id + self._cloud_function_region = cloud_function_region + self._bq_location = bq_location + self._bq_dataset = bq_dataset + self._bq_client = bq_client + self._bq_connection_client = bq_connection_client + self._bq_connection_id = bq_connection_id + + def create_bq_remote_function( + self, input_args, input_types, output_type, endpoint, bq_function_name + ): + """Create a BigQuery remote function given the artifacts of a user defined + function and the http endpoint of a corresponding cloud function.""" + # TODO(shobs): The below command to enable BigQuery Connection API needs + # to be automated. Disabling for now since most target users would not + # have the privilege to enable API in a project. + # log("Making sure BigQuery Connection API is enabled") + # if os.system("gcloud services enable bigqueryconnection.googleapis.com"): + # raise ValueError("Failed to enable BigQuery Connection API") + + # If the intended connection does not exist then create it + if self.check_bq_connection_exists(): + logger.info(f"Connector {self._bq_connection_id} already exists") + else: + connection_name, service_account_id = self.create_bq_connection() + logger.info( + f"Created BQ connection {connection_name} with service account id: {service_account_id}" + ) + + # Set up access on the newly created BQ connection + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + # We would explicitly wait for 60+ seconds for the IAM binding to take effect + command_iam = ( + f"gcloud projects add-iam-policy-binding {self._gcp_project_id}" + + f' --member="serviceAccount:{service_account_id}"' + + ' --role="roles/run.invoker"' + ) + logger.info(f"Setting up IAM binding on the BQ connection: {command_iam}") + _run_system_command(command_iam) + + logger.info( + f"Waiting {self._iam_wait_seconds} seconds for IAM to take effect.." + ) + time.sleep(self._iam_wait_seconds) + + # Create BQ function + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2 + bq_function_args = [] + bq_function_return_type = BigQueryType.from_ibis(output_type) + # We are expecting the input type annotations to be 1:1 with the input args + for idx, name in enumerate(input_args): + bq_function_args.append( + f"{name} {BigQueryType.from_ibis(input_types[idx])}" + ) + create_function_ddl = f""" + CREATE OR REPLACE FUNCTION `{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name}({','.join(bq_function_args)}) + RETURNS {bq_function_return_type} + REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}` + OPTIONS ( + endpoint = "{endpoint}" + )""" + logger.info(f"Creating BQ remote function: {create_function_ddl}") + query_job = self._bq_client.query(create_function_ddl) # Make an API request. + query_job.result() # Wait for the job to complete. + logger.info(f"Created remote function {query_job.ddl_target_routine}") + + def get_remote_function_fully_qualified_name(self, name): + "Get the fully qualilfied name for a BQ remote function." + return "{}.{}.{}".format(self._gcp_project_id, self._bq_dataset, name) + + def get_cloud_function_fully_qualified_name(self, name): + "Get the fully qualilfied name for a cloud function." + return "projects/{}/locations/{}/functions/{}".format( + self._gcp_project_id, self._cloud_function_region, name + ) + + def get_cloud_function_endpoint(self, name): + """Get the http endpoint of a cloud function if it exists.""" + client = functions_v2.FunctionServiceClient() + fully_qualified_name = self.get_cloud_function_fully_qualified_name(name) + try: + response = client.get_function(name=fully_qualified_name) + return response.service_config.uri + except google.api_core.exceptions.NotFound: + pass + return None + + def create_bq_connection(self): + """Create the BigQuery Connection and returns corresponding service account id.""" + client = self._bq_connection_client + connection = bigquery_connection_v1.Connection( + cloud_resource=bigquery_connection_v1.CloudResourceProperties() + ) + request = bigquery_connection_v1.CreateConnectionRequest( + parent=client.common_location_path(self._gcp_project_id, self._bq_location), + connection_id=self._bq_connection_id, + connection=connection, + ) + connection = client.create_connection(request) + return connection.name, connection.cloud_resource.service_account_id + + def check_bq_connection_exists(self): + """Check if the BigQuery Connection exists.""" + client = self._bq_connection_client + request = bigquery_connection_v1.GetConnectionRequest( + name=client.connection_path( + self._gcp_project_id, self._bq_location, self._bq_connection_id + ) + ) + + try: + client.get_connection(request=request) + return True + except google.api_core.exceptions.NotFound: + pass + return False + + def generate_udf_code(self, def_, dir): + """Generate serialized bytecode using cloudpickle given a udf.""" + udf_code_file_name = "udf.py" + udf_bytecode_file_name = "udf.cloudpickle" + + # original code, only for debugging purpose + udf_code = textwrap.dedent(inspect.getsource(def_)) + udf_code_file_path = os.path.join(dir, udf_code_file_name) + with open(udf_code_file_path, "w") as f: + f.write(udf_code) + + # serialized bytecode + udf_bytecode_file_path = os.path.join(dir, udf_bytecode_file_name) + with open(udf_bytecode_file_path, "wb") as f: + cloudpickle.dump(def_, f, protocol=_pickle_protocol_version) + + return udf_code_file_name, udf_bytecode_file_name + + def generate_cloud_function_main_code(self, def_, dir): + """Get main.py code for the cloud function for the given user defined function.""" + + # Pickle the udf with all its dependencies + udf_code_file, udf_bytecode_file = self.generate_udf_code(def_, dir) + handler_func_name = "udf_http" + + # We want to build a cloud function that works for BQ remote functions, + # where we receive `calls` in json which is a batch of rows from BQ SQL. + # The number and the order of values in each row is expected to exactly + # match to the number and order of arguments in the udf , e.g. if the udf is + # def foo(x: int, y: str): + # ... + # then the http request body could look like + # { + # ... + # "calls" : [ + # [123, "hello"], + # [456, "world"] + # ] + # ... + # } + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#input_format + code_template = textwrap.dedent( + """\ + import cloudpickle + import json + + # original udf code is in {udf_code_file} + # serialized udf code is in {udf_bytecode_file} + with open("{udf_bytecode_file}", "rb") as f: + udf = cloudpickle.load(f) + + def {handler_func_name}(request): + request_json = request.get_json(silent=True) + print("[debug] received json request: " + str(request_json)) + calls = request_json["calls"] + replies = [] + for call in calls: + reply = udf(*call) + replies.append(reply) + return_json = json.dumps({{"replies" : replies}}) + return return_json + """ + ) + + code = code_template.format( + udf_code_file=udf_code_file, + udf_bytecode_file=udf_bytecode_file, + handler_func_name=handler_func_name, + ) + + main_py = os.path.join(dir, "main.py") + with open(main_py, "w") as f: + f.write(code) + logger.debug(f"Wrote {os.path.abspath(main_py)}:\n{open(main_py).read()}") + + return handler_func_name + + def generate_cloud_function_code(self, def_, dir): + """Generate the cloud function code for a given user defined function.""" + + # requirements.txt + requirements = ["cloudpickle >= 2.1.0"] + requirements_txt = os.path.join(dir, "requirements.txt") + with open(requirements_txt, "w") as f: + f.write("\n".join(requirements)) + + # main.py + entry_point = self.generate_cloud_function_main_code(def_, dir) + return entry_point + + def create_cloud_function(self, def_, cf_name): + """Create a cloud function from the given user defined function.""" + + # Build and deploy folder structure containing cloud function + with tempfile.TemporaryDirectory() as dir: + entry_point = self.generate_cloud_function_code(def_, dir) + + # We are creating cloud function source code from the currently running + # python version. Use the same version to deploy. This is necessary + # because cloudpickle serialization done in one python version and + # deserialization done in another python version doesn't work. + # TODO(shobs): Figure out how to achieve version compatibility, specially + # when pickle (internally used by cloudpickle) guarantees that: + # https://docs.python.org/3/library/pickle.html#:~:text=The%20pickle%20serialization%20format%20is,unique%20breaking%20change%20language%20boundary. + python_version = "python{}{}".format( + sys.version_info.major, sys.version_info.minor + ) + + # deploy/redeploy the cloud function + # TODO(shobs): Figure out a way to skip this step if a cloud function + # already exists with the same name and source code + command = ( + "gcloud functions deploy" + + f" {cf_name} --gen2" + + f" --runtime={python_version}" + + f" --project={self._gcp_project_id}" + + f" --region={self._cloud_function_region}" + + f" --source={dir}" + + f" --entry-point={entry_point}" + + " --trigger-http" + ) + + # If the cloud function is being created for the first time, then let's + # make it not allow unauthenticated calls. If it was previously created + # then this invocation will update it, in which case do not touch that + # aspect and let the previous policy hold. The reason we do this is to + # avoid an IAM permission needed to update the invocation policy. + # For example, when a cloud function is being created for the first + # time, i.e. + # $ gcloud functions deploy python-foo-http --gen2 --runtime=python310 + # --region=us-central1 + # --source=/source/code/dir + # --entry-point=foo_http + # --trigger-http + # --no-allow-unauthenticated + # It works. When an invocation of the same command is done for the + # second time, it may run into an error like: + # ERROR: (gcloud.functions.deploy) PERMISSION_DENIED: Permission + # 'run.services.setIamPolicy' denied on resource + # 'projects/my_project/locations/us-central1/services/python-foo-http' (or resource may not exist) + # But when --no-allow-unauthenticated is omitted then it goes through. + # It suggests that in the second invocation the command is trying to set + # the IAM policy of the service, and the user running BigQuery + # DataFrame may not have privilege to do so, so better avoid this + # if we can. + if self.get_cloud_function_endpoint(cf_name): + logger.info(f"Updating existing cloud function: {command}") + else: + command = f"{command} --no-allow-unauthenticated" + logger.info(f"Creating new cloud function: {command}") + + _run_system_command(command) + + # Fetch the endpoint of the just created function + endpoint = self.get_cloud_function_endpoint(cf_name) + if not endpoint: + raise ValueError("Couldn't fetch the http endpoint") + + logger.info( + f"Successfully created cloud function {cf_name} with uri ({endpoint})" + ) + return endpoint + + def provision_bq_remote_function( + self, def_, input_types, output_type, uniq_suffix=None + ): + """Provision a BigQuery remote function.""" + # Derive the name of the underlying cloud function and first create + # it if it does not exist + cloud_function_name = get_cloud_function_name(def_, uniq_suffix) + cf_endpoint = self.get_cloud_function_endpoint(cloud_function_name) + if not cf_endpoint: + self.check_cloud_function_tools_and_permissions() + cf_endpoint = self.create_cloud_function(def_, cloud_function_name) + else: + logger.info(f"Cloud function {cloud_function_name} already exists.") + + # Derive the name of the remote function and create/replace it if needed + remote_function_name = get_remote_function_name(def_, uniq_suffix) + rf_endpoint, rf_conn = self.get_remote_function_specs(remote_function_name) + if rf_endpoint != cf_endpoint or rf_conn != self._bq_connection_id: + input_args = inspect.getargs(def_.__code__).args + if len(input_args) != len(input_types): + raise ValueError( + "Exactly one type should be provided for every input arg." + ) + self.create_bq_remote_function( + input_args, input_types, output_type, cf_endpoint, remote_function_name + ) + else: + logger.info(f"Remote function {remote_function_name} already exists.") + + return remote_function_name, cloud_function_name + + def get_remote_function_specs(self, remote_function_name): + """Check whether a remote function already exists for the udf.""" + http_endpoint = None + bq_connection = None + routines = self._bq_client.list_routines( + f"{self._gcp_project_id}.{self._bq_dataset}" + ) + for routine in routines: + if routine.reference.routine_id == remote_function_name: + # TODO(shobs): Use first class properties when they are available + # https://github.com/googleapis/python-bigquery/issues/1552 + rf_options = routine._properties.get("remoteFunctionOptions") + if rf_options: + http_endpoint = rf_options.get("endpoint") + bq_connection = rf_options.get("connection") + if bq_connection: + bq_connection = os.path.basename(bq_connection) + break + return (http_endpoint, bq_connection) + + def check_cloud_function_tools_and_permissions(self): + """Check if the necessary tools and permissions are in place for creating remote function""" + # gcloud CLI comes with bq CLI and they are required for creating google + # cloud function and BigQuery remote function respectively + if not shutil.which("gcloud"): + raise ValueError( + "gcloud tool not installed, install it from https://cloud.google.com/sdk/docs/install" + ) + + # TODO(shobs): Check for permissions too + # I (shobs) tried the following method + # $ gcloud asset search-all-iam-policies \ + # --format=json \ + # --scope=projects/{gcp_project_id} \ + # --query='policy.role.permissions:cloudfunctions.functions.create' + # as a proxy to all the privilges necessary to create cloud function + # https://cloud.google.com/functions/docs/reference/iam/roles#cloudfunctions.developer + # but that itself required the runner to have the permission to enable + # `cloudasset.googleapis.com` + + +# Inspired by @udf decorator implemented in ibis-bigquery package +# https://github.com/ibis-project/ibis-bigquery/blob/main/ibis_bigquery/udf/__init__.py +# which has moved as @js to the ibis package +# https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/udf/__init__.py +def remote_function( + input_types: typing.Sequence[type], + output_type: type, + session: typing.Optional[Session] = None, + bigquery_client: typing.Optional[bigquery.Client] = None, + bigquery_connection_client: typing.Optional[ + bigquery_connection_v1.ConnectionServiceClient + ] = None, + dataset: typing.Optional[str] = None, + bigquery_connection: typing.Optional[str] = None, + reuse: bool = True, +): + """Decorator to turn a user defined function into a BigQuery remote function. + + .. deprecated:: 0.0.1 + Use :func:`bigframes.pandas.remote_function` instead. + + Args: + input_types : list(type). + List of input data types in the user defined function. + output_type : type. + Data type of the output in the user defined function. + session : bigframes.Session, Optional + BigQuery DataFrames session to use for getting default project, + dataset and BigQuery connection. + bigquery_client : google.cloud.bigquery.Client, Optional + Client to use for BigQuery operations. If this param is not provided + then bigquery client from the session would be used. + bigquery_connection_client : google.cloud.bigquery_connection_v1.ConnectionServiceClient, Optional + Client to use for BigQuery connection operations. If this param is + not provided then bigquery connection client from the session would + be used. + dataset : str, Optional. + Dataset to use to create a BigQuery function. It should be in + `.` or `` format. If this + param is not provided then session dataset id would be used. + bigquery_connection : str, Optional. + Name of the BigQuery connection. If this param is not provided then + the bigquery connection from the session would be used. If it is pre + created in the same location as the `bigquery_client.location` then + it would be used, otherwise it would be created dynamically using + the `bigquery_connection_client` assuming the user has necessary + priviliges. + reuse : bool, Optional. + Reuse the remote function if already exists. + `True` by default, which will result in reusing an existing remote + function (if any) that was previously created for the same udf. + Setting it to false would force creating a unique remote function. + If the required remote function does not exist then it would be + created irrespective of this param. + + Notes: + Please make sure following is setup before using this API: + + 1. Have the below APIs enabled for your project: + a. BigQuery Connection API + b. Cloud Functions API + c. Cloud Run API + d. Cloud Build API + e. Artifact Registry API + f. Cloud Resource Manager API + + This can be done from the cloud console (change PROJECT_ID to yours): + https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID + Or from the gcloud CLI: + $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com + + 2. Have following IAM roles enabled for you: + a. BigQuery Data Editor (roles/bigquery.dataEditor) + b. BigQuery Connection Admin (roles/bigquery.connectionAdmin) + c. Cloud Functions Developer (roles/cloudfunctions.developer) + d. Service Account User (roles/iam.serviceAccountUser) + e. Storage Object Viewer (roles/storage.objectViewer) + f. Project IAM Admin (roles/resourcemanager.projectIamAdmin) + (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) + + 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: + a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection + b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + Alternatively, the IAM could also be setup via the gcloud CLI: + $ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker" + + """ + + # A BigQuery client is required to perform BQ operations + if not bigquery_client: + if session: + bigquery_client = session.bqclient + if not bigquery_client: + raise ValueError( + "A bigquery client must be provided, either directly or via session" + ) + + # A BigQuery connection client is required to perform BQ connection operations + if not bigquery_connection_client: + if session: + bigquery_connection_client = session.bqconnectionclient + if not bigquery_connection_client: + raise ValueError( + "A bigquery connection client must be provided, either directly or via session" + ) + + # BQ remote function must be persisted, for which we need a dataset + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#:~:text=You%20cannot%20create%20temporary%20remote%20functions. + if dataset: + dataset_ref = bigquery.DatasetReference.from_string( + dataset, default_project=bigquery_client.project + ) + gcp_project_id = dataset_ref.project + bq_dataset = dataset_ref.dataset_id + else: + gcp_project_id = bigquery_client.project + if session: + bq_dataset = session._session_dataset_id + if not gcp_project_id: + raise ValueError("Project must be provided, either directly or via session") + if not bq_dataset: + raise ValueError("Dataset must be provided, either directly or via session") + + bq_location, cloud_function_region = get_remote_function_locations( + bigquery_client.location + ) + + # A connection is required for BQ remote function + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function + if not bigquery_connection and session: + bigquery_connection = session._remote_udf_connection # type: ignore + if not bigquery_connection: + raise ValueError( + "BigQuery connection must be provided, either directly or via session" + ) + + uniq_suffix = None + if not reuse: + uniq_suffix = "".join( + random.choices(string.ascii_lowercase + string.digits, k=8) + ) + + def wrapper(f): + if not callable(f): + raise TypeError("f must be callable, got {}".format(f)) + + signature = inspect.signature(f) + parameter_names = signature.parameters.keys() + + # Check supported python datatypes and convert to ibis datatypes + type_error_message_format = ( + "type {{}} not supported, supported types are {}.".format( + ", ".join([type_.__name__ for type_ in _supported_io_types]) + ) + ) + for type_ in input_types: + assert type_ in _supported_io_types, type_error_message_format.format(type_) + assert output_type in _supported_io_types, type_error_message_format.format( + output_type + ) + input_types_ibis = [ + python_type_to_bigquery_type(type_) for type_ in input_types + ] + output_type_ibis = python_type_to_bigquery_type(output_type) + + rf_node_fields = { + name: rlz.value(type) + for name, type in zip(parameter_names, input_types_ibis) + } + + try: + rf_node_fields["output_type"] = rlz.shape_like( + "args", dtype=output_type_ibis + ) + except TypeError: + rf_node_fields["output_dtype"] = property(lambda _: output_type_ibis) + rf_node_fields["output_shape"] = rlz.shape_like("args") + + remote_function_client = RemoteFunctionClient( + gcp_project_id, + cloud_function_region, + bq_location, + bq_dataset, + bigquery_client, + bigquery_connection_client, + bigquery_connection, + ) + rf_name, cf_name = remote_function_client.provision_bq_remote_function( + f, input_types_ibis, output_type_ibis, uniq_suffix + ) + rf_fully_qualified_name = f"`{gcp_project_id}.{bq_dataset}`.{rf_name}" + rf_node = type(rf_fully_qualified_name, (ops.ValueOp,), rf_node_fields) + + @compiles(rf_node) + def compiles_rf_node(t, op): + return "{}({})".format( + rf_node.__name__, ", ".join(map(t.translate, op.args)) + ) + + @functools.wraps(f) + def wrapped(*args, **kwargs): + node = rf_node(*args, **kwargs) + return node.to_expr() + + wrapped.__signature__ = signature + wrapped.bigframes_remote_function = ( + remote_function_client.get_remote_function_fully_qualified_name(rf_name) + ) + wrapped.bigframes_cloud_function = ( + remote_function_client.get_cloud_function_fully_qualified_name(cf_name) + ) + return wrapped + + return wrapper diff --git a/bigframes/series.py b/bigframes/series.py new file mode 100644 index 00000000000..f008e0fb060 --- /dev/null +++ b/bigframes/series.py @@ -0,0 +1,964 @@ +# 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. + +"""Series is a 1 dimensional data structure.""" + +from __future__ import annotations + +import textwrap +import typing +from typing import Any, Optional, Union + +import google.cloud.bigquery as bigquery +import ibis.expr.types as ibis_types +import numpy +import pandas +import pandas.core.dtypes.common +import typing_extensions + +import bigframes.core +from bigframes.core import WindowSpec +import bigframes.core.block_transforms as block_ops +import bigframes.core.blocks as blocks +import bigframes.core.groupby as groupby +import bigframes.core.indexers +import bigframes.core.indexes as indexes +from bigframes.core.ordering import OrderingColumnReference, OrderingDirection +import bigframes.core.scalar as scalars +import bigframes.core.window +import bigframes.dataframe +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops +import bigframes.operations.base +import bigframes.operations.datetimes as dt +import bigframes.operations.strings as strings +import third_party.bigframes_vendored.pandas.core.series as vendored_pandas_series + +LevelsType = typing.Union[str, int, typing.Sequence[typing.Union[str, int]]] + + +class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series): + def __init__(self, *args, **kwargs): + self._query_job: Optional[bigquery.QueryJob] = None + super().__init__(*args, **kwargs) + + @property + def dt(self) -> dt.DatetimeMethods: + return dt.DatetimeMethods(self._block) + + @property + def dtype(self): + return self._dtype + + @property + def dtypes(self): + return self._dtype + + @property + def index(self) -> indexes.Index: + return indexes.Index(self) + + @property + def loc(self) -> bigframes.core.indexers.LocSeriesIndexer: + return bigframes.core.indexers.LocSeriesIndexer(self) + + @property + def iloc(self) -> bigframes.core.indexers.IlocSeriesIndexer: + return bigframes.core.indexers.IlocSeriesIndexer(self) + + @property + def name(self) -> Optional[str]: + return self._name + + @property + def shape(self) -> typing.Tuple[int]: + return (self._block.shape[0],) + + @property + def size(self) -> int: + return self.shape[0] + + @property + def empty(self) -> bool: + return self.shape[0] == 0 + + @property + def values(self) -> numpy.ndarray: + return self.to_numpy() + + @property + def query_job(self) -> Optional[bigquery.QueryJob]: + return self._query_job + + def __len__(self): + return self.shape[0] + + def copy(self) -> Series: + return Series(self._block) + + def rename(self, index: Optional[str], **kwargs) -> Series: + if len(kwargs) != 0: + raise NotImplementedError( + "rename does not currently support any keyword arguments." + ) + block = self._block.with_column_labels([index]) + return Series(block) + + def rename_axis( + self, + mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]], + **kwargs, + ) -> Series: + if len(kwargs) != 0: + raise NotImplementedError( + "rename_axis does not currently support any keyword arguments." + ) + # limited implementation: the new index name is simply the 'mapper' parameter + if _is_list_like(mapper): + labels = mapper + else: + labels = [mapper] + return Series(self._block.with_index_labels(labels)) + + def reset_index( + self, + *, + name: typing.Optional[str] = None, + drop: bool = False, + ) -> bigframes.dataframe.DataFrame | Series: + block = self._block.reset_index(drop) + if drop: + return Series(block) + else: + if name: + block = block.assign_label(self._value_column, name) + return bigframes.dataframe.DataFrame(block) + + def __repr__(self) -> str: + # TODO(swast): Add a timeout here? If the query is taking a long time, + # maybe we just print the job metadata that we have so far? + # TODO(swast): Avoid downloading the whole series by using job + # metadata, like we do with DataFrame. + preview = self.compute() + return repr(preview) + + def _to_ibis_expr(self): + """Creates an Ibis table expression representing the Series.""" + expr = self._block.expr.projection([self._value]) + ibis_expr = expr.to_ibis_expr()[self._value_column] + if self._name: + return ibis_expr.name(self._name) + return ibis_expr + + def astype( + self, + dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], + ) -> Series: + return self._apply_unary_op(bigframes.operations.AsTypeOp(dtype)) + + def compute(self) -> pandas.Series: + """Executes deferred operations and downloads the results.""" + df, query_job = self._block.compute((self._value_column,)) + self._query_job = query_job + series = df[self._value_column] + series.name = self._name + return series + + def drop(self, labels: blocks.Label | typing.Sequence[blocks.Label] = None): + block = self._block + index_column = block.index_columns[0] + + if _is_list_like(labels): + block, inverse_condition_id = block.apply_unary_op( + index_column, ops.partial_right(ops.isin_op, labels) + ) + block, condition_id = block.apply_unary_op( + inverse_condition_id, ops.invert_op + ) + + else: + block, condition_id = block.apply_unary_op( + index_column, ops.partial_right(ops.ne_op, labels) + ) + block = block.filter(condition_id) + block = block.drop_columns([condition_id]) + return Series(block.select_column(self._value_column)) + + def droplevel(self, level: LevelsType): + resolved_level_ids = self._resolve_levels(level) + return Series(self._block.drop_levels(resolved_level_ids)) + + def reorder_levels(self, order: LevelsType): + resolved_level_ids = self._resolve_levels(order) + return Series(self._block.reorder_levels(resolved_level_ids)) + + def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: + if _is_list_like(level): + levels = list(level) + else: + levels = [level] + resolved_level_ids = [] + for level_ref in levels: + if isinstance(level_ref, int): + resolved_level_ids.append(self._block.index_columns[level_ref]) + elif isinstance(level_ref, str): + matching_ids = self._block.index_name_to_col_id.get(level_ref, []) + if len(matching_ids) != 1: + raise ValueError("level name cannot be found or is ambiguous") + resolved_level_ids.append(matching_ids[0]) + else: + raise ValueError(f"Unexpected level: {level_ref}") + return resolved_level_ids + + def between(self, left, right, inclusive="both"): + if inclusive not in ["both", "neither", "left", "right"]: + raise ValueError( + "Must set 'inclusive' to one of 'both', 'neither', 'left', or 'right'" + ) + left_op = ops.ge_op if (inclusive in ["left", "both"]) else ops.gt_op + right_op = ops.le_op if (inclusive in ["right", "both"]) else ops.lt_op + return self._apply_binary_op(left, left_op).__and__( + self._apply_binary_op(right, right_op) + ) + + def cumsum(self) -> Series: + return self._apply_window_op( + agg_ops.sum_op, bigframes.core.WindowSpec(following=0) + ) + + def cummax(self) -> Series: + return self._apply_window_op( + agg_ops.max_op, bigframes.core.WindowSpec(following=0) + ) + + def cummin(self) -> Series: + return self._apply_window_op( + agg_ops.min_op, bigframes.core.WindowSpec(following=0) + ) + + def shift(self, periods: int = 1) -> Series: + window = bigframes.core.WindowSpec( + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.ShiftOp(periods), window) + + def diff(self) -> Series: + return self - self.shift(1) + + def rank( + self, + axis=0, + method: str = "average", + numeric_only=False, + na_option: str = "keep", + ascending: bool = True, + ) -> Series: + return Series(block_ops.rank(self._block, method, na_option, ascending)) + + def fillna(self, value=None) -> "Series" | None: + return self._apply_binary_op(value, ops.fillna_op) + + def head(self, n: int = 5) -> Series: + return typing.cast(Series, self.iloc[0:n]) + + def tail(self, n: int = 5) -> Series: + return typing.cast(Series, self.iloc[-n:]) + + def nlargest(self, n: int = 5, keep: str = "first") -> Series: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + block = self._block + if keep == "last": + block = block.reversed() + ordering = ( + OrderingColumnReference( + self._value_column, direction=OrderingDirection.DESC + ), + ) + block = block.order_by(ordering, stable=True) + if keep in ("first", "last"): + return Series(block.slice(0, n)) + else: # keep == "all": + block, counter = block.apply_window_op( + self._value_column, + agg_ops.rank_op, + window_spec=WindowSpec(ordering=ordering), + ) + block, condition = block.apply_unary_op( + counter, ops.partial_right(ops.le_op, n) + ) + block = block.filter(condition) + block = block.select_column(self._value_column) + return Series(block) + + def nsmallest(self, n: int = 5, keep: str = "first") -> Series: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + block = self._block + if keep == "last": + block = block.reversed() + ordering = (OrderingColumnReference(self._value_column),) + block = block.order_by(ordering, stable=True) + if keep in ("first", "last"): + return Series(block.slice(0, n)) + else: # keep == "all": + block, counter = block.apply_window_op( + self._value_column, + agg_ops.rank_op, + window_spec=WindowSpec(ordering=ordering), + ) + block, condition = block.apply_unary_op( + counter, ops.partial_right(ops.le_op, n) + ) + block = block.filter(condition) + block = block.select_column(self._value_column) + return Series(block) + + def isna(self) -> "Series": + return self._apply_unary_op(ops.isnull_op) + + isnull = isna + + def notna(self) -> "Series": + return self._apply_unary_op(ops.notnull_op) + + notnull = notna + + def __and__(self, other: bool | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.and_op) + + __rand__ = __and__ + + def __or__(self, other: bool | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.or_op) + + __ror__ = __or__ + + def __add__(self, other: float | int | Series | pandas.Series) -> Series: + return self.add(other) + + def __radd__(self, other: float | int | Series | pandas.Series) -> Series: + return self.radd(other) + + def add(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.add_op) + + def radd(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.add_op)) + + def __sub__(self, other: float | int | Series | pandas.Series) -> Series: + return self.sub(other) + + def __rsub__(self, other: float | int | Series | pandas.Series) -> Series: + return self.rsub(other) + + def sub(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.sub_op) + + def rsub(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.sub_op)) + + def __mul__(self, other: float | int | Series | pandas.Series) -> Series: + return self.mul(other) + + def __rmul__(self, other: float | int | Series | pandas.Series) -> Series: + return self.rmul(other) + + def mul(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.mul_op) + + def rmul(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.mul_op)) + + multiply = mul + + def __truediv__(self, other: float | int | Series | pandas.Series) -> Series: + return self.truediv(other) + + def __rtruediv__(self, other: float | int | Series | pandas.Series) -> Series: + return self.rtruediv(other) + + def truediv(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.div_op) + + def rtruediv(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.div_op)) + + div = truediv + + divide = truediv + + rdiv = rtruediv + + def __floordiv__(self, other: float | int | Series | pandas.Series) -> Series: + return self.floordiv(other) + + def __rfloordiv__(self, other: float | int | Series | pandas.Series) -> Series: + return self.rfloordiv(other) + + def floordiv(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.floordiv_op) + + def rfloordiv(self, other: float | int | Series | pandas.Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.floordiv_op)) + + def __lt__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + return self.lt(other) + + def __le__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + return self.le(other) + + def lt(self, other) -> Series: + return self._apply_binary_op(other, ops.lt_op) + + def le(self, other) -> Series: + return self._apply_binary_op(other, ops.le_op) + + def __gt__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + return self.gt(other) + + def __ge__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + return self.ge(other) + + def gt(self, other) -> Series: + return self._apply_binary_op(other, ops.gt_op) + + def ge(self, other) -> Series: + return self._apply_binary_op(other, ops.ge_op) + + def __mod__(self, other) -> Series: # type: ignore + return self.mod(other) + + def __rmod__(self, other) -> Series: # type: ignore + return self.rmod(other) + + def mod(self, other) -> Series: # type: ignore + return self._apply_binary_op(other, ops.mod_op) + + def rmod(self, other) -> Series: # type: ignore + return self._apply_binary_op(other, ops.reverse(ops.mod_op)) + + def __matmul__(self, other): + return (self * other).sum() + + dot = __matmul__ + + def abs(self) -> Series: + return self._apply_unary_op(ops.abs_op) + + def round(self, decimals=0) -> "Series": + def round_op(x: ibis_types.Value, y: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).round( + digits=typing.cast(ibis_types.IntegerValue, y) + ) + + return self._apply_binary_op(decimals, round_op) + + def all(self) -> bool: + return typing.cast(bool, self._apply_aggregation(agg_ops.all_op)) + + def any(self) -> bool: + return typing.cast(bool, self._apply_aggregation(agg_ops.any_op)) + + def count(self) -> int: + return typing.cast(int, self._apply_aggregation(agg_ops.count_op)) + + def nunique(self) -> int: + return typing.cast(int, self._apply_aggregation(agg_ops.nunique_op)) + + def max(self) -> scalars.Scalar: + return self._apply_aggregation(agg_ops.max_op) + + def min(self) -> scalars.Scalar: + return self._apply_aggregation(agg_ops.min_op) + + def std(self) -> float: + return typing.cast(float, self._apply_aggregation(agg_ops.std_op)) + + def var(self) -> float: + return typing.cast(float, self._apply_aggregation(agg_ops.var_op)) + + def _central_moment(self, n: int) -> float: + """Useful helper for calculating central moment statistics""" + # Nth central moment is mean((x-mean(x))^n) + # See: https://en.wikipedia.org/wiki/Moment_(mathematics) + mean = self.mean() + mean_deltas = self - mean + delta_power = mean_deltas + # TODO(tbergeron): Replace with pow once implemented + for i in range(1, n): + delta_power = delta_power * mean_deltas + return delta_power.mean() + + def kurt(self) -> float: + # TODO(tbergeron): Cache intermediate count/moment/etc. statistics at block level + count = self.count() + moment4 = self._central_moment(4) + moment2 = self._central_moment(2) # AKA: Population Variance + + # Kurtosis is often defined as the second standardize moment: moment(4)/moment(2)**2 + # Pandas however uses Fisher’s estimator, implemented below + numerator = (count + 1) * (count - 1) * moment4 + denominator = (count - 2) * (count - 3) * moment2**2 + adjustment = 3 * (count - 1) ** 2 / ((count - 2) * (count - 3)) + + return (numerator / denominator) - adjustment + + kurtosis = kurt + + def mode(self) -> Series: + block = self._block + # Approach: Count each value, return each value for which count(x) == max(counts)) + block, agg_ids = block.aggregate( + [self._value_column], + ((self._value_column, agg_ops.count_op),), + as_index=False, + ) + value_count_col_id = agg_ids[0] + block, max_value_count_col_id = block.apply_window_op( + value_count_col_id, + agg_ops.max_op, + window_spec=WindowSpec(), + ) + block, is_mode_col_id = block.apply_binary_op( + value_count_col_id, + max_value_count_col_id, + ops.eq_op, + ) + block = block.filter(is_mode_col_id) + mode_values_series = Series( + block.select_column(self._value_column).assign_label( + self._value_column, self.name + ) + ) + return typing.cast( + Series, mode_values_series.sort_values().reset_index(drop=True) + ) + + def mean(self) -> float: + return typing.cast(float, self._apply_aggregation(agg_ops.mean_op)) + + def sum(self) -> float: + return typing.cast(float, self._apply_aggregation(agg_ops.sum_op)) + + def prod(self) -> float: + return typing.cast(float, self._apply_aggregation(agg_ops.product_op)) + + product = prod + + def __eq__(self, other: object) -> Series: # type: ignore + return self.eq(other) + + def __ne__(self, other: object) -> Series: # type: ignore + return self.ne(other) + + def __invert__(self) -> Series: + return self._apply_unary_op(ops.invert_op) + + def eq(self, other: object) -> Series: + # TODO: enforce stricter alignment + return self._apply_binary_op(other, ops.eq_op) + + def ne(self, other: object) -> Series: + # TODO: enforce stricter alignment + return self._apply_binary_op(other, ops.ne_op) + + def where(self, cond, other=None): + value_id, cond_id, other_id, block = self._align3(cond, other) + block, result_id = block.apply_ternary_op( + value_id, cond_id, other_id, ops.where_op + ) + return Series(block.select_column(result_id).with_column_labels([self.name])) + + def clip(self, lower, upper): + if lower is None and upper is None: + return self + if lower is None: + return self._apply_binary_op(upper, ops.clip_upper, alignment="left") + if upper is None: + return self._apply_binary_op(lower, ops.clip_lower, alignment="left") + value_id, lower_id, upper_id, block = self._align3(lower, upper) + block, result_id = block.apply_ternary_op( + value_id, lower_id, upper_id, ops.clip_op + ) + return Series(block.select_column(result_id).with_column_labels([self.name])) + + def argmax(self) -> scalars.Scalar: + block, row_nums = self._block.promote_offsets() + block = block.order_by( + [ + OrderingColumnReference( + self._value_column, direction=OrderingDirection.DESC + ), + OrderingColumnReference(row_nums), + ] + ) + return typing.cast( + scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] + ) + + def argmin(self) -> scalars.Scalar: + block, row_nums = self._block.promote_offsets() + block = block.order_by( + [ + OrderingColumnReference(self._value_column), + OrderingColumnReference(row_nums), + ] + ) + return typing.cast( + scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] + ) + + def __getitem__(self, indexer: Series): + # TODO: enforce stricter alignment, should fail if indexer is missing any keys. + (left, right, block) = self._align(indexer, "left") + block = block.filter(right) + block = block.select_column(left) + return Series(block) + + def __getattr__(self, key: str): + if hasattr(pandas.Series, key): + raise NotImplementedError( + textwrap.dedent( + f""" + BigQuery DataFrames has not yet implemented an equivalent to + 'pandas.Series.{key}'. Please check + https://github.com/googleapis/python-bigquery-dataframes/issues for + existing feature requests, or file your own. + Please include information about your use case, as well as + relevant code snippets. + """ + ) + ) + else: + raise AttributeError(key) + + def _align3(self, other1: Series | scalars.Scalar, other2: Series | scalars.Scalar, how="left") -> tuple[str, str, str, blocks.Block]: # type: ignore + """Aligns the series value with 2 other scalars or series objects. Returns new values and joined tabled expression.""" + values, index = self._align_n([other1, other2], how) + return (values[0], values[1], values[2], index) + + def _apply_aggregation(self, op: agg_ops.AggregateOp) -> Any: + aggregation_result = typing.cast( + ibis_types.Scalar, op._as_ibis(self[self.notnull()]._to_ibis_expr()) + ) + return bigframes.core.scalar.DeferredScalar( + aggregation_result, self._block._expr._session + ).compute() + + def _apply_window_op( + self, + op: agg_ops.WindowOp, + window_spec: bigframes.core.WindowSpec, + ): + block = self._block + block, result_id = block.apply_window_op( + self._value_column, op, window_spec=window_spec, result_label=self.name + ) + return Series(block.select_column(result_id)) + + def value_counts( + self, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + *, + dropna: bool = True, + ): + block = block_ops.value_counts( + self._block, + [self._value_column], + normalize=normalize, + ascending=ascending, + dropna=dropna, + ) + return Series(block) + + def sort_values(self, *, axis=0, ascending=True, na_position="last") -> Series: + if na_position not in ["first", "last"]: + raise ValueError("Param na_position must be one of 'first' or 'last'") + direction = OrderingDirection.ASC if ascending else OrderingDirection.DESC + block = self._block.order_by( + [ + OrderingColumnReference( + self._value_column, + direction=direction, + na_last=(na_position == "last"), + ) + ] + ) + return Series(block) + + def sort_index(self, *, axis=0, ascending=True, na_position="last") -> Series: + # TODO(tbergeron): Support level parameter once multi-index introduced. + if na_position not in ["first", "last"]: + raise ValueError("Param na_position must be one of 'first' or 'last'") + block = self._block + direction = OrderingDirection.ASC if ascending else OrderingDirection.DESC + na_last = na_position == "last" + ordering = [ + OrderingColumnReference(column, direction=direction, na_last=na_last) + for column in block.index_columns + ] + block = block.order_by(ordering) + return Series(block) + + def rolling(self, window: int, min_periods=None) -> bigframes.core.window.Window: + # To get n size window, need current row and n-1 preceding rows. + window_spec = WindowSpec( + preceding=window - 1, following=0, min_periods=min_periods or window + ) + return bigframes.core.window.Window( + self._block, window_spec, self._value_column + ) + + def expanding(self, min_periods: int = 1) -> bigframes.core.window.Window: + window_spec = WindowSpec(following=0, min_periods=min_periods) + return bigframes.core.window.Window( + self._block, window_spec, self._value_column + ) + + def groupby( + self, + by: typing.Union[ + blocks.Label, Series, typing.Sequence[typing.Union[blocks.Label, Series]] + ] = None, + axis=0, + level: typing.Optional[ + int | str | typing.Sequence[int] | typing.Sequence[str] + ] = None, + as_index: bool = True, + *, + dropna: bool = True, + ) -> bigframes.core.groupby.SeriesGroupBy: + if (by is not None) and (level is not None): + raise ValueError("Do not specify both 'by' and 'level'") + if not as_index: + raise ValueError("as_index=False only valid with DataFrame") + if axis: + raise ValueError("No axis named {} for object type Series".format(level)) + if by is not None: + return self._groupby_values(by, dropna) + if level is not None: + return self._groupby_level(level, dropna) + else: + raise TypeError("You have to supply one of 'by' and 'level'") + + def _groupby_level( + self, + level: int | str | typing.Sequence[int] | typing.Sequence[str], + dropna: bool = True, + ) -> bigframes.core.groupby.SeriesGroupBy: + return groupby.SeriesGroupBy( + self._block, + self._value_column, + self._resolve_levels(level), + value_name=self.name, + dropna=dropna, + ) + + def _groupby_values( + self, + by: typing.Union[ + blocks.Label, Series, typing.Sequence[typing.Union[blocks.Label, Series]] + ], + dropna: bool = True, + ) -> bigframes.core.groupby.SeriesGroupBy: + if not isinstance(by, Series) and _is_list_like(by): + by = list(by) + else: + by = [typing.cast(typing.Union[blocks.Label, Series], by)] + + block = self._block + grouping_cols: typing.Sequence[str] = [] + value_col = self._value_column + for key in by: + if isinstance(key, Series): + combined_index, ( + get_column_left, + get_column_right, + ) = block.index.join( + key._block.index, how="inner" if dropna else "left" + ) + + value_col = get_column_left(self._value_column) + grouping_cols = [ + *[get_column_left(value) for value in grouping_cols], + get_column_right(key._value_column), + ] + block = combined_index._block + else: + # Interpret as index level + matches = block.index_name_to_col_id.get(key, []) + if len(matches) != 1: + raise ValueError( + f"GroupBy key {key} does not map to unambiguous index level" + ) + grouping_cols = [*grouping_cols, matches[0]] + + return groupby.SeriesGroupBy( + block, + value_col, + grouping_cols, + value_name=self.name, + dropna=dropna, + ) + + def apply(self, func) -> Series: + # TODO(shobs, b/274645634): Support convert_dtype, args, **kwargs + # is actually a ternary op + return self._apply_unary_op(ops.RemoteFunctionOp(func)) + + def add_prefix(self, prefix: str, axis: int | str | None = None) -> Series: + return Series(self._get_block().add_prefix(prefix)) + + def add_suffix(self, suffix: str, axis: int | str | None = None) -> Series: + return Series(self._get_block().add_suffix(suffix)) + + def drop_duplicates(self, *, keep: str = "first") -> Series: + block = block_ops.drop_duplicates(self._block, (self._value_column,), keep) + return Series(block) + + def unique(self) -> Series: + return self.drop_duplicates() + + def duplicated(self, keep: str = "first") -> Series: + block, indicator = block_ops.indicate_duplicates( + self._block, (self._value_column,), keep + ) + return Series( + block.select_column( + indicator, + ).with_column_labels([self.name]) + ) + + def mask(self, cond, other=None) -> Series: + if callable(cond): + cond = self.apply(cond) + + if not isinstance(cond, Series): + raise TypeError( + f"Only bigframes series condition is supported, received {type(cond).__name__}" + ) + return self.where(~cond, other) + + def to_frame(self) -> bigframes.dataframe.DataFrame: + # To be consistent with Pandas, it assigns 0 as the column name if missing. 0 is the first element of RangeIndex. + block = self._block.with_column_labels([self.name] if self.name else ["0"]) + return bigframes.dataframe.DataFrame(block) + + def to_csv(self, path_or_buf=None, **kwargs) -> typing.Optional[str]: + # TODO(b/280651142): Implement version that leverages bq export native csv support to bypass local pandas step. + return self.compute().to_csv(path_or_buf, **kwargs) + + def to_dict(self, into: type[dict] = dict) -> typing.Mapping: + return typing.cast(dict, self.compute().to_dict(into)) + + def to_excel(self, excel_writer, sheet_name="Sheet1", **kwargs) -> None: + return self.compute().to_excel(excel_writer, sheet_name, **kwargs) + + def to_json( + self, + path_or_buf=None, + orient: typing.Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + **kwargs, + ) -> typing.Optional[str]: + # TODO(b/280651142): Implement version that leverages bq export native csv support to bypass local pandas step. + return self.compute().to_json(path_or_buf, **kwargs) + + def to_latex( + self, buf=None, columns=None, header=True, index=True, **kwargs + ) -> typing.Optional[str]: + return self.compute().to_latex( + buf, columns=columns, header=header, index=index, **kwargs + ) + + def tolist(self) -> list: + return self.compute().to_list() + + to_list = tolist + + def to_markdown( + self, + buf: typing.IO[str] | None = None, + mode: str = "wt", + index: bool = True, + **kwargs, + ) -> typing.Optional[str]: + return self.compute().to_markdown(buf, mode=mode, index=index, **kwargs) # type: ignore + + def to_numpy( + self, dtype=None, copy=False, na_value=None, **kwargs + ) -> numpy.ndarray: + return self.compute().to_numpy(dtype, copy, na_value, **kwargs) + + __array__ = to_numpy + + def to_pickle(self, path, **kwargs) -> None: + return self.compute().to_pickle(path, **kwargs) + + def to_string( + self, + buf=None, + na_rep="NaN", + float_format=None, + header=True, + index=True, + length=False, + dtype=False, + name=False, + max_rows=None, + min_rows=None, + ) -> typing.Optional[str]: + return self.compute().to_string( + buf, + na_rep, + float_format, + header, + index, + length, + dtype, + name, + max_rows, + min_rows, + ) + + def to_xarray(self): + return self.compute().to_xarray() + + # Keep this at the bottom of the Series class to avoid + # confusing type checker by overriding str + @property + def str(self) -> strings.StringMethods: + return strings.StringMethods(self._block) + + def _slice( + self, + start: typing.Optional[int] = None, + stop: typing.Optional[int] = None, + step: typing.Optional[int] = None, + ) -> bigframes.series.Series: + return bigframes.series.Series( + self._block.slice(start=start, stop=stop, step=step).select_column( + self._value_column + ), + ) + + +def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: + return pandas.api.types.is_list_like(obj) diff --git a/bigframes/session.py b/bigframes/session.py new file mode 100644 index 00000000000..9590ecb8c77 --- /dev/null +++ b/bigframes/session.py @@ -0,0 +1,1055 @@ +# 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. + +"""Session manages the connection to BigQuery.""" + +from __future__ import annotations + +import logging +import os +import re +import textwrap +import typing +from typing import ( + Any, + Callable, + Dict, + IO, + Iterable, + List, + Literal, + MutableSequence, + Optional, + Sequence, + Tuple, + Union, +) +import uuid +import warnings + +import google.api_core.client_info +import google.api_core.client_options +import google.api_core.exceptions +import google.api_core.gapic_v1.client_info +import google.auth.credentials +import google.cloud.bigquery as bigquery +import google.cloud.bigquery_connection_v1 +import google.cloud.bigquery_storage_v1 +import google.cloud.storage as storage # type: ignore +import ibis +import ibis.backends.bigquery as ibis_bigquery +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types +import numpy as np +import pandas +import pydata_google_auth + +import bigframes._config.bigquery_options as bigquery_options +import bigframes.core as core +import bigframes.core.blocks as blocks +import bigframes.core.guid as guid +from bigframes.core.ordering import OrderingColumnReference +import bigframes.dataframe as dataframe +import bigframes.formatting_helpers as formatting_helpers +import bigframes.ml.loader +from bigframes.remote_function import remote_function as bigframes_rf +import bigframes.version +import third_party.bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq +import third_party.bigframes_vendored.pandas.io.parquet as third_party_pandas_parquet +import third_party.bigframes_vendored.pandas.io.parsers.readers as third_party_pandas_readers + +_ENV_DEFAULT_PROJECT = "GOOGLE_CLOUD_PROJECT" +_APPLICATION_NAME = f"bigframes/{bigframes.version.__version__}" +_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] + +# BigQuery is a REST API, which requires the protocol as part of the URL. +_BIGQUERY_REGIONAL_ENDPOINT = "https://{location}-bigquery.googleapis.com" + +# BigQuery Connection and Storage are gRPC APIs, which don't support the +# https:// protocol in the API endpoint URL. +_BIGQUERYCONNECTION_REGIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com" +_BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com" + +# TODO(swast): Need to connect to regional endpoints when performing remote +# functions operations (BQ Connection API, Cloud Run / Cloud Functions). + +# pydata-google-auth credentials in case auth credentials are not available +# otherwise +_pydata_google_auth_credentials: Optional[google.auth.credentials.Credentials] = None +_pydata_google_auth_project: Optional[str] = None + +logger = logging.getLogger(__name__) + + +def _is_query(query_or_table: str) -> bool: + """Determine if `query_or_table` is a table ID or a SQL string""" + return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None + + +# TODO(shobs): Remove it after the same is available via pydata-google-auth +# after https://github.com/pydata/pydata-google-auth/pull/71 is merged, released +# and upgraded in the google colab image. +def _ensure_application_default_credentials_in_colab_environment(): + # This is a special handling for google colab environment where we want to + # use the colab specific authentication flow + # https://github.com/googlecolab/colabtools/blob/3c8772efd332289e1c6d1204826b0915d22b5b95/google/colab/auth.py#L209 + try: + from google.colab import auth + + auth.authenticate_user() + except Exception: + # We are catching a broad exception class here because we want to be + # agnostic to anything that could internally go wrong in the google + # colab auth. Some of the known exception we want to pass on are: + # + # ModuleNotFoundError: No module named 'google.colab' + # ImportError: cannot import name 'auth' from 'google.cloud' + # MessageError: Error: credential propagation was unsuccessful + # + # The MessageError happens on Vertex Colab when it fails to resolve auth + # from the Compute Engine Metadata server. + pass + + +pydata_google_auth.auth._ensure_application_default_credentials_in_colab_environment = ( + _ensure_application_default_credentials_in_colab_environment +) + + +def _get_default_credentials_with_project(): + global _pydata_google_auth_credentials, _pydata_google_auth_project + if not _pydata_google_auth_credentials or not _pydata_google_auth_credentials.valid: + # We want to initiate auth via a non-local web server which + # particularly helps in a cloud notebook environment where the + # machine running the notebook UI and the VM running the notebook + # runtime are not the same. + # TODO(shobs, b/278903498): Use BigQuery DataFrames's own client id + # and secret + ( + _pydata_google_auth_credentials, + _pydata_google_auth_project, + ) = pydata_google_auth.default(_SCOPES, use_local_webserver=False) + return _pydata_google_auth_credentials, _pydata_google_auth_project + + +def _create_bq_clients( + project: Optional[str], + location: Optional[str], + use_regional_endpoints: Optional[bool], + credentials: Optional[google.auth.credentials.Credentials], +) -> typing.Tuple[ + bigquery.Client, + google.cloud.bigquery_connection_v1.ConnectionServiceClient, + google.cloud.bigquery_storage_v1.BigQueryReadClient, +]: + """Create and initialize BigQuery client objects.""" + + credentials_project = None + if credentials is None: + credentials, credentials_project = _get_default_credentials_with_project() + + # Prefer the project in this order: + # 1. Project explicitly specified by the user + # 2. Project set in the environment + # 3. Project associated with the default credentials + project = ( + project + or os.getenv(_ENV_DEFAULT_PROJECT) + or typing.cast(Optional[str], credentials_project) + ) + + if not project: + raise ValueError("Project must be set to initialize BigQuery client.") + + if use_regional_endpoints: + bq_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERY_REGIONAL_ENDPOINT.format(location=location), + ) + bqstorage_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERYSTORAGE_REGIONAL_ENDPOINT.format(location=location) + ) + bqconnection_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERYCONNECTION_REGIONAL_ENDPOINT.format(location=location) + ) + else: + bq_options = None + bqstorage_options = None + bqconnection_options = None + + bq_info = google.api_core.client_info.ClientInfo(user_agent=_APPLICATION_NAME) + bqclient = bigquery.Client( + client_info=bq_info, + client_options=bq_options, + credentials=credentials, + project=project, + ) + + bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + bqconnectionclient = google.cloud.bigquery_connection_v1.ConnectionServiceClient( + client_info=bqconnection_info, + client_options=bqconnection_options, + credentials=credentials, + ) + + bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + bqstorageclient = google.cloud.bigquery_storage_v1.BigQueryReadClient( + client_info=bqstorage_info, + client_options=bqstorage_options, + credentials=credentials, + ) + + return bqclient, bqconnectionclient, bqstorageclient + + +class Session( + third_party_pandas_gbq.GBQIOMixin, + third_party_pandas_parquet.ParquetIOMixin, + third_party_pandas_readers.ReaderIOMixin, +): + """Establishes a BigQuery connection to capture a group of job activities related to + DataFrames.""" + + def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): + if context is None: + context = bigquery_options.BigQueryOptions() + + # TODO(swast): Get location from the environment. + if context is None or context.location is None: + self._location = "US" + warnings.warn( + f"No explicit location is set, so using location {self._location} for the session.", + stacklevel=2, + ) + else: + self._location = context.location + + ( + self.bqclient, + self.bqconnectionclient, + self.bqstorageclient, + ) = _create_bq_clients( + project=context.project, + location=self._location, + use_regional_endpoints=context.use_regional_endpoints, + credentials=context.credentials, + ) + + self._create_and_bind_bq_session() + self.ibis_client = typing.cast( + ibis_bigquery.Backend, + ibis.bigquery.connect( + project_id=context.project, + client=self.bqclient, + storage_client=self.bqstorageclient, + ), + ) + + self._remote_udf_connection = context.remote_udf_connection + + # Now that we're starting the session, don't allow the options to be + # changed. + context._session_started = True + + @property + def _session_dataset_id(self): + """A dataset for storing temporary objects local to the session + This is a workaround for BQML models and remote functions that do not + yet support session-temporary instances.""" + return self._session_dataset.dataset_id + + def _create_and_bind_bq_session(self): + """Create a BQ session and bind the session id with clients to capture BQ activities: + go/bigframes-transient-data""" + job_config = bigquery.QueryJobConfig(create_session=True) + query_job = self.bqclient.query( + "SELECT 1", job_config=job_config, location=self._location + ) + query_job.result() # blocks until finished + self._session_id = query_job.session_info.session_id + + self.bqclient.default_query_job_config = bigquery.QueryJobConfig( + connection_properties=[ + bigquery.ConnectionProperty("session_id", self._session_id) + ] + ) + self.bqclient.default_load_job_config = bigquery.LoadJobConfig( + connection_properties=[ + bigquery.ConnectionProperty("session_id", self._session_id) + ] + ) + + # Dataset for storing BQML models and remote functions, which don't yet + # support proper session temporary storage yet + self._session_dataset = bigquery.Dataset( + f"{self.bqclient.project}.bigframes_temp_{self._location.lower().replace('-', '_')}" + ) + self._session_dataset.location = self._location + self._session_dataset.default_table_expiration_ms = 24 * 60 * 60 * 1000 + + # TODO: handle case when the dataset does not exist and the user does + # not have permission to create one (bigquery.datasets.create IAM) + self.bqclient.create_dataset(self._session_dataset, exists_ok=True) + + def close(self): + """Terminated the BQ session, otherwises the session will be terminated automatically after + 24 hours of inactivity or after 7 days.""" + if self._session_id is not None and self.bqclient is not None: + abort_session_query = "CALL BQ.ABORT_SESSION('{}')".format(self._session_id) + query_job = self.bqclient.query(abort_session_query) + query_job.result() # blocks until finished + self._session_id = None + + def read_gbq( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + # Add a verify index argument that fails if the index is not unique. + ) -> dataframe.DataFrame: + # TODO(b/281571214): Generate prompt to show the progress of read_gbq. + if _is_query(query): + return self.read_gbq_query( + query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + ) + else: + # TODO(swast): Query the snapshot table but mark it as a + # deterministic query so we can avoid serializing if we have a + # unique index. + return self.read_gbq_table( + query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + ) + + def read_gbq_query( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + ) -> dataframe.DataFrame: + """Turn a SQL query into a DataFrame. + + Note: Because the results are written to a temporary table, ordering by + ``ORDER BY`` is not preserved. A unique `index_col` is recommended. Use + ``row_number() over ()`` if there is no natural unique index or you + want to preserve ordering. + + See also: :meth:`Session.read_gbq`. + """ + # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so + # these docstrings are inline. + + if isinstance(index_col, str): + index_cols = [index_col] + else: + index_cols = list(index_col) + + # Make sure we cluster by the index column so that subsequent + # operations are as speedy as they can be. + if index_cols: + destination: bigquery.Table | bigquery.TableReference = ( + self._query_to_session_table(query, index_cols) + ) + else: + _, query_job = self._start_query(query) + query_job.result() # Wait for job to finish. + destination = query_job.destination + + # If there was no destination table, that means the query must have + # been DDL or DML. Return some job metadata, instead. + if not destination: + return dataframe.DataFrame( + data=pandas.DataFrame( + { + "statement_type": [query_job.statement_type], + "job_id": [query_job.job_id], + "location": [query_job.location], + } + ), + session=self, + ) + + return self.read_gbq_table( + f"{destination.project}.{destination.dataset_id}.{destination.table_id}", + index_col=index_cols, + col_order=col_order, + max_results=max_results, + ) + + def read_gbq_table( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + ) -> dataframe.DataFrame: + """Turn a BigQuery table into a DataFrame. + + See also: :meth:`Session.read_gbq`. + """ + # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so + # these docstrings are inline. + # TODO(swast): Can we re-use the temp table from other reads in the + # session, if the original table wasn't modified? + table_ref = bigquery.table.TableReference.from_string( + query, default_project=self.bqclient.project + ) + + if table_ref.dataset_id.upper() == "_SESSION": + # _SESSION tables aren't supported by the tables.get REST API. + table_expression = self.ibis_client.sql( + f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" + ) + else: + # TODO(swast): Read from a table snapshot so that reads are consistent. + table_expression = self.ibis_client.table( + table_ref.table_id, + database=f"{table_ref.project}.{table_ref.dataset_id}", + ) + + for key in col_order: + if key not in table_expression.columns: + raise ValueError( + f"Column '{key}' of `col_order` not found in this table." + ) + + if isinstance(index_col, str): + index_cols: List[str] = [index_col] + else: + index_cols = list(index_col) + + for key in index_cols: + if key not in table_expression.columns: + raise ValueError( + f"Column `{key}` of `index_col` not found in this table." + ) + + # If the index is unique and sortable, then we don't need to generate + # an ordering column. + ordering = None + is_total_ordering = False + + if len(index_cols) != 0: + index_labels = typing.cast(List[Optional[str]], index_cols) + distinct_table = table_expression.select(*index_cols).distinct() + is_unique_sql = f"""WITH full_table AS ( + {self.ibis_client.compile(table_expression)} + ), + distinct_table AS ( + {self.ibis_client.compile(distinct_table)} + ) + + SELECT (SELECT COUNT(*) FROM full_table) AS total_count, + (SELECT COUNT(*) FROM distinct_table) AS distinct_count + """ + results, _ = self._start_query(is_unique_sql) + row = next(iter(results)) + + total_count = row["total_count"] + distinct_count = row["distinct_count"] + is_total_ordering = total_count == distinct_count + ordering = core.ExpressionOrdering( + ordering_value_columns=[ + core.OrderingColumnReference(column_id) for column_id in index_cols + ], + ) + + if not is_total_ordering: + # Make sure when we generate an ordering, the row_number() + # coresponds to the index columns. + table_expression = table_expression.order_by(index_cols) + warnings.warn( + textwrap.dedent( + f""" + Got a non-unique index. A consistent ordering is not + guaranteed. DataFrame has {total_count} rows, + but only {distinct_count} distinct index values. + """, + ) + ) + else: + table_expression, ordering = self._create_sequential_ordering( + table_expression + ) + ordering_id_column = ordering.ordering_id + assert ordering_id_column is not None + is_total_ordering = True + index_cols = [ordering_id_column] + index_labels = [None] + + if max_results is not None: + if max_results <= 0: + raise ValueError("`max_results` should be a positive number.") + table_expression = table_expression.limit(max_results) + + return self._read_gbq_with_ordering( + table_expression=table_expression, + col_order=col_order, + index_cols=index_cols, + index_labels=index_labels, + ordering=ordering, + is_total_ordering=is_total_ordering, + ) + + def _read_gbq_with_ordering( + self, + table_expression: ibis_types.Table, + *, + col_order: Iterable[str] = (), + index_cols: Sequence[str] = (), + index_labels: Sequence[Optional[str]] = (), + ordering: core.ExpressionOrdering, + is_total_ordering: bool = False, + ) -> dataframe.DataFrame: + """Internal helper method that loads DataFrame from Google BigQuery given an ordering column. + + Args: + table_expression: an ibis table expression to be executed in BigQuery. + col_order: List of BigQuery column names in the desired order for results DataFrame. + index_cols: List of column names to use as the index or multi-index. + ordering: Column name to be used for ordering. If not supplied, a default ordering is generated. + + Returns: + A DataFrame representing results of the query or table. + """ + if len(index_cols) != len(index_labels): + raise ValueError( + "Needs same number of index labels are there are index columns. " + f"Got {len(index_labels)}, expected {len(index_cols)}." + ) + + if not index_cols: + raise ValueError("Need at least 1 index column.") + + # Logic: + # no total ordering, index -> create sequential order, ordered by index, use for both ordering and index + # total ordering, index -> use ordering as ordering, index as index + + # This code block ensures the existence of a total ordering. + if not is_total_ordering: + # Rows are not ordered, we need to generate a default ordering and materialize it + table_expression, ordering = self._create_sequential_ordering( + table_expression, index_cols + ) + + index_col_values = [table_expression[index_id] for index_id in index_cols] + + column_keys = list(col_order) + if len(column_keys) == 0: + non_columns = set(index_cols) + if ordering.ordering_id is not None: + non_columns.add(ordering.ordering_id) + column_keys = [ + key for key in table_expression.columns if key not in non_columns + ] + return self._read_ibis( + table_expression, + index_col_values, + index_labels, + column_keys, + ordering=ordering, + ) + + def _read_bigquery_load_job( + self, + filepath_or_buffer: str | IO["bytes"], + table: bigquery.Table, + *, + job_config: bigquery.LoadJobConfig, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + ) -> dataframe.DataFrame: + if isinstance(index_col, str): + index_cols = [index_col] + else: + index_cols = list(index_col) + + if not job_config.clustering_fields and index_cols: + job_config.clustering_fields = index_cols + + if isinstance(filepath_or_buffer, str): + if filepath_or_buffer.startswith("gs://"): + load_job = self.bqclient.load_table_from_uri( + filepath_or_buffer, table, job_config=job_config + ) + else: + with open(filepath_or_buffer, "rb") as source_file: + load_job = self.bqclient.load_table_from_file( + source_file, table, job_config=job_config + ) + else: + load_job = self.bqclient.load_table_from_file( + filepath_or_buffer, table, job_config=job_config + ) + + load_job.result() # Wait for the job to complete + + # The BigQuery REST API for tables.get doesn't take a session ID, so we + # can't get the schema for a temp table that way. + return self.read_gbq_table( + f"{table.project}.{table.dataset_id}.{table.table_id}", + index_col=index_col, + col_order=col_order, + ) + + def _read_ibis( + self, + table_expression: ibis_types.Table, + index_cols: Sequence[ibis_types.Value], + index_labels: Sequence[Optional[str]], + column_keys: Sequence[str], + ordering: Optional[core.ExpressionOrdering] = None, + ): + """Turns a table expression (plus index column) into a DataFrame.""" + hidden_ordering_columns = None + if ordering is not None and ordering.ordering_id is not None: + hidden_ordering_columns = (table_expression[ordering.ordering_id],) + + columns = list(index_cols) + for key in column_keys: + if key not in table_expression.columns: + raise ValueError(f"Column '{key}' not found in this table.") + columns.append(table_expression[key]) + + block = blocks.Block( + core.ArrayValue( + self, table_expression, columns, hidden_ordering_columns, ordering + ), + [index_col.get_name() for index_col in index_cols], + index_labels=index_labels, + ) + + return dataframe.DataFrame(block) + + def read_gbq_model(self, model_name: str): + """Loads a BQML model from Google BigQuery. + + Args: + model_name : the model's name in BigQuery in the format + `project_id.dataset_id.model_id`, or just `dataset_id.model_id` + to load from the default project. + + Returns: + A bigframes.ml Model wrapping the model. + """ + model_ref = bigquery.ModelReference.from_string( + model_name, default_project=self.bqclient.project + ) + model = self.bqclient.get_model(model_ref) + return bigframes.ml.loader.from_bq(self, model) + + def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame: + """Loads DataFrame from a Pandas DataFrame. + + The Pandas DataFrame will be persisted as a temporary BigQuery table, which can be + automatically recycled after the Session is closed. + + Args: + pandas_dataframe: a Pandas DataFrame object to be loaded. + + Returns: + A BigQuery DataFrames. + """ + # Add order column to pandas DataFrame to preserve order in BigQuery + ordering_col = "rowid" + columns = frozenset(pandas_dataframe.columns) + suffix = 2 + while ordering_col in columns: + ordering_col = f"rowid_{suffix}" + suffix += 1 + + pandas_dataframe_copy = pandas_dataframe.copy() + pandas_dataframe_copy[ordering_col] = np.arange(pandas_dataframe_copy.shape[0]) + + # Specify the datetime dtypes, which is auto-detected as timestamp types. + schema = [] + for column, dtype in zip(pandas_dataframe.columns, pandas_dataframe.dtypes): + if dtype == "timestamp[us][pyarrow]": + schema.append( + bigquery.SchemaField(column, bigquery.enums.SqlTypeNames.DATETIME) + ) + + # Unnamed are not copied to BigQuery when load_table_from_dataframe + # executes. + index_cols = list( + filter(lambda name: name is not None, pandas_dataframe_copy.index.names) + ) + index_labels = typing.cast(List[Optional[str]], index_cols) + cluster_cols = index_cols + [ordering_col] + + if len(index_cols) == 0: + index_cols = [ordering_col] + index_labels = [None] + + job_config = bigquery.LoadJobConfig(schema=schema) + job_config.clustering_fields = cluster_cols + + # TODO(swast): Rename the unnamed index columns and restore them after + # the load job completes. + # Column values will be loaded as null if the column name has spaces. + # https://github.com/googleapis/python-bigquery/issues/1566 + load_table_destination = self._create_session_table() + load_job = self.bqclient.load_table_from_dataframe( + pandas_dataframe_copy, + load_table_destination, + job_config=job_config, + ) + load_job.result() # Wait for the job to complete + + ordering = core.ExpressionOrdering( + ordering_id_column=OrderingColumnReference(ordering_col), is_sequential=True + ) + table_expression = self.ibis_client.sql( + f"SELECT * FROM `{load_table_destination.table_id}`" + ) + + return self._read_gbq_with_ordering( + table_expression=table_expression, + index_cols=index_cols, + index_labels=index_labels, + ordering=ordering, + is_total_ordering=True, + ) + + def read_csv( + self, + filepath_or_buffer: str | IO["bytes"], + *, + sep: Optional[str] = ",", + header: Optional[int] = 0, + names: Optional[ + Union[MutableSequence[Any], np.ndarray[Any, Any], Tuple[Any, ...], range] + ] = None, + index_col: Optional[ + Union[int, str, Sequence[Union[str, int]], Literal[False]] + ] = None, + usecols: Optional[ + Union[ + MutableSequence[str], + Tuple[str, ...], + Sequence[int], + pandas.Series, + pandas.Index, + np.ndarray[Any, Any], + Callable[[Any], bool], + ] + ] = None, + dtype: Optional[Dict] = None, + engine: Optional[ + Literal["c", "python", "pyarrow", "python-fwf", "bigquery"] + ] = None, + encoding: Optional[str] = None, + **kwargs, + ) -> dataframe.DataFrame: + table = bigquery.Table(self._create_session_table()) + + if engine is not None and engine == "bigquery": + if any(param is not None for param in (dtype, names)): + not_supported = ("dtype", "names") + raise NotImplementedError( + f"BigQuery engine does not support these arguments: {not_supported}" + ) + + if index_col is not None and ( + not index_col or not isinstance(index_col, str) + ): + raise NotImplementedError( + "BigQuery engine only supports a single column name for `index_col`." + ) + + # None value for index_col cannot be passed to read_gbq + if index_col is None: + index_col = () + + # usecols should only be an iterable of strings (column names) for use as col_order in read_gbq. + col_order: Tuple[Any, ...] = tuple() + if usecols is not None: + if isinstance(usecols, Iterable) and all( + isinstance(col, str) for col in usecols + ): + col_order = tuple(col for col in usecols) + else: + raise NotImplementedError( + "BigQuery engine only supports an iterable of strings for `usecols`." + ) + + valid_encodings = {"UTF-8", "ISO-8859-1"} + if encoding is not None and encoding not in valid_encodings: + raise NotImplementedError( + f"BigQuery engine only supports the following encodings: {valid_encodings}" + ) + + job_config = bigquery.LoadJobConfig() + job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED + job_config.source_format = bigquery.SourceFormat.CSV + job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY + job_config.autodetect = True + job_config.field_delimiter = sep + job_config.encoding = encoding + + # We want to match pandas behavior. If header is 0, no rows should be skipped, so we + # do not need to set `skip_leading_rows`. If header is None, then there is no header. + # Setting skip_leading_rows to 0 does that. If header=N and N>0, we want to skip N rows. + # `skip_leading_rows` skips N-1 rows, so we set it to header+1. + if header is not None and header > 0: + job_config.skip_leading_rows = header + 1 + elif header is None: + job_config.skip_leading_rows = 0 + + return self._read_bigquery_load_job( + filepath_or_buffer, + table, + job_config=job_config, + index_col=index_col, + col_order=col_order, + ) + else: + if any(arg in kwargs for arg in ("chunksize", "iterator")): + raise NotImplementedError( + "'chunksize' and 'iterator' arguments are not supported." + ) + + if isinstance(filepath_or_buffer, str): + self._check_file_size(filepath_or_buffer) + pandas_df = pandas.read_csv( + filepath_or_buffer, + sep=sep, + header=header, + names=names, + index_col=index_col, + usecols=usecols, + dtype=dtype, + engine=engine, + encoding=encoding, + **kwargs, + ) + return self.read_pandas(pandas_df) + + def read_parquet( + self, + path: str | IO["bytes"], + ) -> dataframe.DataFrame: + # Note: "engine" is omitted because it is redundant. Loading a table + # from a pandas DataFrame will just create another parquet file + load + # job anyway. + table = bigquery.Table(self._create_session_table()) + + job_config = bigquery.LoadJobConfig() + job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED + job_config.source_format = bigquery.SourceFormat.PARQUET + job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY + + return self._read_bigquery_load_job(path, table, job_config=job_config) + + def _check_file_size(self, filepath: str): + max_size = 1024 * 1024 * 1024 # 1 GB in bytes + if filepath.startswith("gs://"): # GCS file path + client = storage.Client() + bucket_name, blob_name = filepath.split("/", 3)[2:] + bucket = client.bucket(bucket_name) + blob = bucket.blob(blob_name) + blob.reload() + file_size = blob.size + else: # local file path + file_size = os.path.getsize(filepath) + + if file_size > max_size: + # Convert to GB + file_size = round(file_size / (1024**3), 1) + max_size = int(max_size / 1024**3) + logger.warning( + f"File size {file_size}GB exceeds {max_size}GB. " + "It is recommended to use engine='bigquery' " + "for large files to avoid loading the file into local memory." + ) + + def _create_session_table(self) -> bigquery.TableReference: + table_name = f"{uuid.uuid4().hex}" + dataset = bigquery.Dataset( + bigquery.DatasetReference(self.bqclient.project, "_SESSION") + ) + return dataset.table(table_name) + + def _create_sequential_ordering( + self, table: ibis_types.Table, index_cols: Iterable[str] = () + ) -> Tuple[ibis_types.Table, core.ExpressionOrdering]: + # Since this might also be used as the index, don't use the default + # "ordering ID" name. + default_ordering_name = guid.generate_guid("bigframes_ordering_") + default_ordering_col = ( + ibis.row_number().cast(ibis_dtypes.int64).name(default_ordering_name) + ) + table = table.mutate(**{default_ordering_name: default_ordering_col}) + table_ref = self._query_to_session_table( + self.ibis_client.compile(table), + cluster_cols=list(index_cols) + [default_ordering_name], + ) + table = self.ibis_client.sql(f"SELECT * FROM `{table_ref.table_id}`") + ordering_reference = core.OrderingColumnReference(default_ordering_name) + ordering = core.ExpressionOrdering( + ordering_id_column=ordering_reference, is_sequential=True + ) + return table, ordering + + def _query_to_session_table( + self, query_text: str, cluster_cols: Iterable[str] + ) -> bigquery.TableReference: + # Can't set a table in _SESSION as destination via query job API, so we + # run DDL, instead. + table = self._create_session_table() + cluster_cols_sql = ", ".join(f"`{cluster_col}`" for cluster_col in cluster_cols) + + # TODO(swast): This might not support multi-statement SQL queries. + ddl_text = f""" + CREATE TEMP TABLE `_SESSION`.`{table.table_id}` + CLUSTER BY {cluster_cols_sql} + AS {query_text} + """ + query_job = self.bqclient.query(ddl_text) + try: + query_job.result() # Wait for the job to complete + except google.api_core.exceptions.Conflict: + # Allow query retry to succeed. + pass + return table + + def remote_function( + self, + input_types: List[type], + output_type: type, + dataset: Optional[str] = None, + bigquery_connection: Optional[str] = None, + reuse: bool = True, + ): + """Decorator to turn a user defined function into a BigQuery remote function. + + Args: + input_types (list(type)): + List of input data types in the user defined function. + output_type (type): + Data type of the output in the user defined function. + dataset (str, Optional): + Dataset to use to create a BigQuery function. It should be in + `.` or `` format. If this + param is not provided then session dataset id would be used. + bigquery_connection (str, Optional): + Name of the BigQuery connection. If it is pre created in the same + location as the `bigquery_client.location` then it would be used, + otherwise it would be created dynamically assuming the user has + necessary priviliges. If this param is not provided then the + bigquery connection from the session would be used. + reuse (bool, Optional): + Reuse the remote function if already exists. + `True` by default, which will result in reusing an existing remote + function (if any) that was previously created for the same udf. + Setting it to false would force creating a unique remote function. + If the required remote function does not exist then it would be + created irrespective of this param. + + Notes: + Please make sure following is setup before using this API: + + 1. Have the below APIs enabled for your project: + a. BigQuery Connection API + b. Cloud Functions API + c. Cloud Run API + d. Cloud Build API + e. Artifact Registry API + f. Cloud Resource Manager API + + This can be done from the cloud console (change PROJECT_ID to yours): + https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID + Or from the gcloud CLI: + $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com + + 2. Have following IAM roles enabled for you: + a. BigQuery Data Editor (roles/bigquery.dataEditor) + b. BigQuery Connection Admin (roles/bigquery.connectionAdmin) + c. Cloud Functions Developer (roles/cloudfunctions.developer) + d. Service Account User (roles/iam.serviceAccountUser) + e. Storage Object Viewer (roles/storage.objectViewer) + f. Project IAM Admin (roles/resourcemanager.projectIamAdmin) + (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) + + 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: + a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection + b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + Alternatively, the IAM could also be setup via the gcloud CLI: + $ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker" + + """ + return bigframes_rf( + input_types, + output_type, + session=self, + dataset=dataset, + bigquery_connection=bigquery_connection, + reuse=reuse, + ) + + def _start_query( + self, + sql: str, + job_config: Optional[bigquery.job.QueryJobConfig] = None, + max_results: Optional[int] = None, + ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: + if job_config is not None: + query_job = self.bqclient.query(sql, job_config=job_config) + else: + query_job = self.bqclient.query(sql) + + opts = bigframes.options.display + if opts.progress_bar is not None: + results_iterator = formatting_helpers.wait_for_job( + query_job, max_results, opts.progress_bar + ) + else: + results_iterator = query_job.result(max_results=max_results) + return results_iterator, query_job + + def _extract_table(self, source_table, destination_uris, job_config): + extract_job = self.bqclient.extract_table( + source=source_table, + destination_uris=destination_uris, + job_config=job_config, + ) + extract_job.result() + return extract_job + + def _rows_to_dataframe( + self, row_iterator: bigquery.table.RowIterator + ) -> pandas.DataFrame: + return row_iterator.to_dataframe( + bool_dtype=pandas.BooleanDtype(), + int_dtype=pandas.Int64Dtype(), + float_dtype=pandas.Float64Dtype(), + string_dtype=pandas.StringDtype(storage="pyarrow"), + ) + + +def connect(context: Optional[bigquery_options.BigQueryOptions] = None) -> Session: + return Session(context) diff --git a/bigframes/version.py b/bigframes/version.py new file mode 100644 index 00000000000..df6dc3c17c8 --- /dev/null +++ b/bigframes/version.py @@ -0,0 +1,15 @@ +# Copyright 2022 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. + +__version__ = "0.1.0" diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000000..d4bb2cbb9ed --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 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/_static/custom.css b/docs/_static/custom.css new file mode 100644 index 00000000000..c663a6ef337 --- /dev/null +++ b/docs/_static/custom.css @@ -0,0 +1,15 @@ +/* Ensure minimum width for 'Parameters' / 'Returns' column */ +dl.field-list > dt { + min-width: 100px +} + +/* Insert space between methods for readability */ +dl.method { + padding-top: 10px; + padding-bottom: 10px +} + +/* Insert empty space between classes */ +dl.class { + padding-bottom: 50px +} diff --git a/docs/_templates/layout.html b/docs/_templates/layout.html new file mode 100644 index 00000000000..bfdcc4759c5 --- /dev/null +++ b/docs/_templates/layout.html @@ -0,0 +1,45 @@ + +{% extends "!layout.html" %} +{%- block content %} +{%- if theme_fixed_sidebar|lower == 'true' %} +
+ {{ sidebar() }} + {%- block document %} +
+ {%- if render_sidebar %} +
+ {%- endif %} + + {%- block relbar_top %} + {%- if theme_show_relbar_top|tobool %} +
+   + {{- rellink_markup () }} +
+ {%- endif %} + {% endblock %} + +
+ {% block body %} {% endblock %} +
+ + {%- block relbar_bottom %} + {%- if theme_show_relbar_bottom|tobool %} +
+   + {{- rellink_markup () }} +
+ {%- endif %} + {% endblock %} + + {%- if render_sidebar %} +
+ {%- endif %} +
+ {%- endblock %} +
+
+{%- else %} +{{ super() }} +{%- endif %} +{%- endblock %} 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/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000000..3ab5be51035 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,382 @@ +# -*- coding: utf-8 -*- +# Copyright 2021 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. +# +# bigframes documentation build configuration file +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +import os +import sys +from typing import Any, Dict + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +sys.path.insert(0, os.path.abspath("..")) + +# For plugins that can not read conf.py. +# See also: https://github.com/docascode/sphinx-docfx-yaml/issues/85 +sys.path.insert(0, os.path.abspath(".")) + +__version__ = "" + +# -- General configuration ------------------------------------------------ + +# If your documentation needs a minimal Sphinx version, state it here. +needs_sphinx = "1.5.5" + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", + "sphinx.ext.coverage", + "sphinx.ext.doctest", + "sphinx.ext.napoleon", + "sphinx.ext.todo", + "sphinx.ext.viewcode", + "recommonmark", +] + +# autodoc/autosummary flags +autoclass_content = "both" +autodoc_default_options = {"members": True} +autosummary_generate = True + + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# source_suffix = ['.rst', '.md'] +source_suffix = [".rst", ".md"] + +# The encoding of source files. +# source_encoding = 'utf-8-sig' + +# The root toctree document. +root_doc = "index" + +# General information about the project. +project = "BigQuery DataFrames" +copyright = "2022-2023 Google LLC" +author = "Google APIs" + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. +# +# The full version, including alpha/beta/rc tags. +release = __version__ +# The short X.Y version. +version = ".".join(release.split(".")[0:2]) + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# There are two options for replacing |today|: either, you set today to some +# non-false value, then it is used: +# today = '' +# Else, today_fmt is used as the format for a strftime call. +# today_fmt = '%B %d, %Y' + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +exclude_patterns = [ + "_build", + "**/.nox/**/*", + "samples/AUTHORING_GUIDE.md", + "samples/CONTRIBUTING.md", + "samples/snippets/README.rst", +] + +# The reST default role (used for this markup: `text`) to use for all +# documents. +# default_role = None + +# If true, '()' will be appended to :func: etc. cross-reference text. +# add_function_parentheses = True + +# If true, the current module name will be prepended to all description +# unit titles (such as .. function::). +# add_module_names = True + +# If true, sectionauthor and moduleauthor directives will be shown in the +# output. They are ignored by default. +# show_authors = False + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + +# A list of ignored prefixes for module index sorting. +# modindex_common_prefix = [] + +# If true, keep warnings as "system message" paragraphs in the built documents. +# keep_warnings = False + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = True + + +# -- Options for HTML output ---------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +html_theme = "alabaster" + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +html_theme_options = { + "description": "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine.", + # "github_user": "googleapis", + # "github_repo": "python-bigquery-storage", + # "github_banner": True, + "font_family": "'Roboto', Georgia, sans", + "head_font_family": "'Roboto', Georgia, serif", + "code_font_family": "'Roboto Mono', 'Consolas', monospace", +} + +# Add any paths that contain custom themes here, relative to this directory. +# html_theme_path = [] + +# The name for this set of Sphinx documents. If None, it defaults to +# " v documentation". +# html_title = None + +# A shorter title for the navigation bar. Default is the same as html_title. +# html_short_title = None + +# The name of an image file (relative to this directory) to place at the top +# of the sidebar. +# html_logo = None + +# The name of an image file (within the static path) to use as favicon of the +# docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 +# pixels large. +# html_favicon = None + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +# html_static_path = ["_static"] + +# Add any extra paths that contain custom files (such as robots.txt or +# .htaccess) here, relative to this directory. These files are copied +# directly to the root of the documentation. +# html_extra_path = [] + +# If not '', a 'Last updated on:' timestamp is inserted at every page bottom, +# using the given strftime format. +# html_last_updated_fmt = '%b %d, %Y' + +# If true, SmartyPants will be used to convert quotes and dashes to +# typographically correct entities. +# html_use_smartypants = True + +# Custom sidebar templates, maps document names to template names. +# html_sidebars = {} + +# Additional templates that should be rendered to pages, maps page names to +# template names. +# html_additional_pages = {} + +# If false, no module index is generated. +# html_domain_indices = True + +# If false, no index is generated. +# html_use_index = True + +# If true, the index is split into individual pages for each letter. +# html_split_index = False + +# If true, links to the reST sources are added to the pages. +# html_show_sourcelink = True + +# If true, "Created using Sphinx" is shown in the HTML footer. Default is True. +# html_show_sphinx = True + +# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. +# html_show_copyright = True + +# If true, an OpenSearch description file will be output, and all pages will +# contain a tag referring to it. The value of this option must be the +# base URL from which the finished HTML is served. +# html_use_opensearch = '' + +# This is the file name suffix for HTML files (e.g. ".xhtml"). +# html_file_suffix = None + +# Language to be used for generating the HTML full-text search index. +# Sphinx supports the following languages: +# 'da', 'de', 'en', 'es', 'fi', 'fr', 'hu', 'it', 'ja' +# 'nl', 'no', 'pt', 'ro', 'ru', 'sv', 'tr' +# html_search_language = 'en' + +# A dictionary with options for the search language support, empty by default. +# Now only 'ja' uses this config value +# html_search_options = {'type': 'default'} + +# The name of a javascript file (relative to the configuration directory) that +# implements a search results scorer. If empty, the default will be used. +# html_search_scorer = 'scorer.js' + +# Output file base name for HTML help builder. +htmlhelp_basename = "bigframes-doc" + +# -- Options for warnings ------------------------------------------------------ + + +suppress_warnings = [ + # Temporarily suppress this to avoid "more than one target found for + # cross-reference" warning, which are intractable for us to avoid while in + # a mono-repo. + # See https://github.com/sphinx-doc/sphinx/blob + # /2a65ffeef5c107c19084fabdd706cdff3f52d93c/sphinx/domains/python.py#L843 + "ref.python" +] + +# -- Options for LaTeX output --------------------------------------------- + +latex_elements: Dict[str, Any] = { + # Avoid "too deeply nested" error by using enumitem package. + # See: https://stackoverflow.com/a/28454426/101923 + "preamble": r""" +\usepackage{enumitem} +\setlistdepth{99} +""" +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + ( + root_doc, + "bigframes.tex", + "bigframes Documentation", + author, + "manual", + ) +] + +# The name of an image file (relative to this directory) to place at the top of +# the title page. +# latex_logo = None + +# For "manual" documents, if this is true, then toplevel headings are parts, +# not chapters. +# latex_use_parts = False + +# If true, show page references after internal links. +# latex_show_pagerefs = False + +# If true, show URL addresses after external links. +# latex_show_urls = False + +# Documents to append as an appendix to all manuals. +# latex_appendices = [] + +# If false, no module index is generated. +# latex_domain_indices = True + + +# -- Options for manual page output --------------------------------------- + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + ( + root_doc, + "bigframes", + "bigframes Documentation", + [author], + 1, + ) +] + +# If true, show URL addresses after external links. +# man_show_urls = False + + +# -- Options for Texinfo output ------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ( + root_doc, + "bigframes", + "bigframes Documentation", + author, + "bigframes", + "bigframes Library", + "APIs", + ) +] + +# Documents to append as an appendix to all manuals. +# texinfo_appendices = [] + +# If false, no module index is generated. +# texinfo_domain_indices = True + +# How to display URL addresses: 'footnote', 'no', or 'inline'. +# texinfo_show_urls = 'footnote' + +# If true, do not generate a @detailmenu in the "Top" node's menu. +# texinfo_no_detailmenu = False + + +# Example configuration for intersphinx: refer to the Python standard library. +intersphinx_mapping = { + "python": ("https://python.readthedocs.org/en/latest/", None), + "google-auth": ("https://googleapis.dev/python/google-auth/latest/", None), + "google.api_core": ( + "https://googleapis.dev/python/google-api-core/latest/", + None, + ), + "grpc": ("https://grpc.github.io/grpc/python/", None), + "proto-plus": ("https://proto-plus-python.readthedocs.io/en/latest/", None), + "protobuf": ("https://googleapis.dev/python/protobuf/latest/", None), +} + + +# Napoleon settings +napoleon_google_docstring = True +napoleon_numpy_docstring = True +napoleon_include_private_with_doc = False +napoleon_include_special_with_doc = True +napoleon_use_admonition_for_examples = False +napoleon_use_admonition_for_notes = False +napoleon_use_admonition_for_references = False +napoleon_use_ivar = False +napoleon_use_param = True +napoleon_use_rtype = True diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst new file mode 100644 index 00000000000..cde2a70a3fb --- /dev/null +++ b/docs/getting_started/index.rst @@ -0,0 +1,27 @@ +Getting started +=============== + +Prerequisites +-------------- + +* Install the ``bigframes`` package. +* Create a GCP project and billing account. +* When running locally, authenticate with application default credentials. See + the `gcloud auth application-default login + `_ + reference. + +Code sample +----------- + +.. literalinclude:: ../samples/snippets/quickstart.py + :language: python + :dedent: 4 + :start-after: [START bigquery_bigframes_quickstart] + :end-before: [END bigquery_bigframes_quickstart] + +Learn more +---------- + +* See :ref:`user_guide` for more samples. +* See :ref:`bigframes_ml` for more ML samples. diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000000..eb721fbf293 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,18 @@ +.. include:: README.rst + +.. toctree:: + :maxdepth: 3 + + getting_started/index + user_guide/index + reference/index + +Changelog +--------- + +For a list of all BigQuery DataFrames releases: + +.. toctree:: + :maxdepth: 2 + + changelog diff --git a/docs/reference/bigframes.ml/README.rst b/docs/reference/bigframes.ml/README.rst new file mode 100644 index 00000000000..80a1fe97b73 --- /dev/null +++ b/docs/reference/bigframes.ml/README.rst @@ -0,0 +1,125 @@ +BigQuery DataFrames ML +====================== + +As BigQuery DataFrames implements the Pandas API over top of BigQuery, BigQuery +DataFrame ML implements the SKLearn API over top of BigQuery Machine Learning. + +Tutorial +-------- + +Start a session and initialize a dataframe for a BigQuery table + +.. code-block:: python + + import bigframes.pandas + + df = bigframes.pandas.read_gbq("bigquery-public-data.ml_datasets.penguins") + df + +Clean and prepare the data + +.. code-block:: python + + # 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() + + # take a peek at the training data + training_data + +.. code-block:: python + + # 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']] + +Use train_test_split to create train and test datasets + +.. code-block:: python + + from bigframes.ml.model_selection import train_test_split + + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2) + +Define the model training pipeline + +.. code-block:: python + + from bigframes.ml.linear_model import LinearRegression + from bigframes.ml.pipeline import Pipeline + from bigframes.ml.compose import ColumnTransformer + from bigframes.ml.preprocessing import StandardScaler, OneHotEncoder + + preprocessing = ColumnTransformer([ + ("onehot", OneHotEncoder(), ["island", "species", "sex"]), + ("scaler", StandardScaler(), ["culmen_depth_mm", "culmen_length_mm", "flipper_length_mm"]), + ]) + + model = LinearRegression(fit_intercept=False) + + pipeline = Pipeline([ + ('preproc', preprocessing), + ('linreg', model) + ]) + + # view the pipeline + pipeline + +Train the pipeline + +.. code-block:: python + + pipeline.fit(X_train, y_train) + +Evaluate the model's performance on the test data + +.. code-block:: python + + from bigframes.ml.metrics import r2_score + + y_pred = pipeline.predict(X_test) + + r2_score(y_test, y_pred) + +Make predictions on new data + +.. code-block:: python + + import pandas + + new_penguins = bigframes.pandas.read_pandas( + pandas.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + ) + + # view the new data + new_penguins + +.. code-block:: python + + pipeline.predict(new_penguins) + +Save the trained model to BigQuery, so we can load it later + +.. code-block:: python + + pipeline.to_gbq("bqml_tutorial.penguins_model", replace=True) diff --git a/docs/reference/bigframes.ml/cluster.rst b/docs/reference/bigframes.ml/cluster.rst new file mode 100644 index 00000000000..5456b728b67 --- /dev/null +++ b/docs/reference/bigframes.ml/cluster.rst @@ -0,0 +1,5 @@ +bigframes.ml.cluster +==================== + +.. automodule:: bigframes.ml.cluster + :members: diff --git a/docs/reference/bigframes.ml/compose.rst b/docs/reference/bigframes.ml/compose.rst new file mode 100644 index 00000000000..e41f17a33b1 --- /dev/null +++ b/docs/reference/bigframes.ml/compose.rst @@ -0,0 +1,5 @@ +bigframes.ml.compose +==================== + +.. automodule:: bigframes.ml.compose + :members: diff --git a/docs/reference/bigframes.ml/decomposition.rst b/docs/reference/bigframes.ml/decomposition.rst new file mode 100644 index 00000000000..122dbc2b168 --- /dev/null +++ b/docs/reference/bigframes.ml/decomposition.rst @@ -0,0 +1,5 @@ +bigframes.ml.decomposition +========================== + +.. automodule:: bigframes.ml.decomposition + :members: diff --git a/docs/reference/bigframes.ml/forecasting.rst b/docs/reference/bigframes.ml/forecasting.rst new file mode 100644 index 00000000000..037bd7e82ef --- /dev/null +++ b/docs/reference/bigframes.ml/forecasting.rst @@ -0,0 +1,5 @@ +bigframes.ml.forecasting +======================== + +.. automodule:: bigframes.ml.forecasting + :members: diff --git a/docs/reference/bigframes.ml/imported.rst b/docs/reference/bigframes.ml/imported.rst new file mode 100644 index 00000000000..54120534c02 --- /dev/null +++ b/docs/reference/bigframes.ml/imported.rst @@ -0,0 +1,5 @@ +bigframes.ml.imported +===================== + +.. automodule:: bigframes.ml.imported + :members: diff --git a/docs/reference/bigframes.ml/index.rst b/docs/reference/bigframes.ml/index.rst new file mode 100644 index 00000000000..e17edc2ee98 --- /dev/null +++ b/docs/reference/bigframes.ml/index.rst @@ -0,0 +1,30 @@ +.. _bigframes_ml: +.. include:: README.rst + +API Reference +------------- + +.. toctree:: + :maxdepth: 3 + + cluster + + compose + + decomposition + + forecasting + + imported + + linear_model + + llm + + metrics + + model_selection + + pipeline + + preprocessing diff --git a/docs/reference/bigframes.ml/linear_model.rst b/docs/reference/bigframes.ml/linear_model.rst new file mode 100644 index 00000000000..5936858e64b --- /dev/null +++ b/docs/reference/bigframes.ml/linear_model.rst @@ -0,0 +1,5 @@ +bigframes.ml.linear_model +========================= + +.. automodule:: bigframes.ml.linear_model + :members: diff --git a/docs/reference/bigframes.ml/llm.rst b/docs/reference/bigframes.ml/llm.rst new file mode 100644 index 00000000000..4debef74ffb --- /dev/null +++ b/docs/reference/bigframes.ml/llm.rst @@ -0,0 +1,5 @@ +bigframes.ml.llm +================ + +.. automodule:: bigframes.ml.llm + :members: diff --git a/docs/reference/bigframes.ml/metrics.rst b/docs/reference/bigframes.ml/metrics.rst new file mode 100644 index 00000000000..dc951ebed42 --- /dev/null +++ b/docs/reference/bigframes.ml/metrics.rst @@ -0,0 +1,5 @@ +bigframes.ml.metrics +==================== + +.. automodule:: bigframes.ml.metrics + :members: diff --git a/docs/reference/bigframes.ml/model_selection.rst b/docs/reference/bigframes.ml/model_selection.rst new file mode 100644 index 00000000000..b0d4e5ed3ee --- /dev/null +++ b/docs/reference/bigframes.ml/model_selection.rst @@ -0,0 +1,5 @@ +bigframes.ml.model_selection +============================ + +.. automodule:: bigframes.ml.model_selection + :members: diff --git a/docs/reference/bigframes.ml/pipeline.rst b/docs/reference/bigframes.ml/pipeline.rst new file mode 100644 index 00000000000..239ee59c431 --- /dev/null +++ b/docs/reference/bigframes.ml/pipeline.rst @@ -0,0 +1,5 @@ +bigframes.ml.pipeline +===================== + +.. automodule:: bigframes.ml.pipeline + :members: diff --git a/docs/reference/bigframes.ml/preprocessing.rst b/docs/reference/bigframes.ml/preprocessing.rst new file mode 100644 index 00000000000..4d2279fc8ea --- /dev/null +++ b/docs/reference/bigframes.ml/preprocessing.rst @@ -0,0 +1,5 @@ +bigframes.ml.preprocessing +========================== + +.. automodule:: bigframes.ml.preprocessing + :members: diff --git a/docs/reference/bigframes.pandas/frame.rst b/docs/reference/bigframes.pandas/frame.rst new file mode 100644 index 00000000000..a49bcc8f7cb --- /dev/null +++ b/docs/reference/bigframes.pandas/frame.rst @@ -0,0 +1,9 @@ + +========= +DataFrame +========= + +.. autoclass:: bigframes.dataframe.DataFrame + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.pandas/general_functions.rst b/docs/reference/bigframes.pandas/general_functions.rst new file mode 100644 index 00000000000..4fff9aabf83 --- /dev/null +++ b/docs/reference/bigframes.pandas/general_functions.rst @@ -0,0 +1,8 @@ + +================= +General functions +================= + +.. automodule:: bigframes.pandas + :members: + :undoc-members: diff --git a/docs/reference/bigframes.pandas/groupby.rst b/docs/reference/bigframes.pandas/groupby.rst new file mode 100644 index 00000000000..483340f3487 --- /dev/null +++ b/docs/reference/bigframes.pandas/groupby.rst @@ -0,0 +1,20 @@ + +======= +GroupBy +======= + +DataFrameGroupBy +---------------- + +.. autoclass:: bigframes.core.groupby.DataFrameGroupBy + :members: + :inherited-members: + :undoc-members: + +SeriesGroupBy +------------- + +.. autoclass:: bigframes.core.groupby.SeriesGroupBy + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.pandas/index.rst b/docs/reference/bigframes.pandas/index.rst new file mode 100644 index 00000000000..c7ff586884e --- /dev/null +++ b/docs/reference/bigframes.pandas/index.rst @@ -0,0 +1,15 @@ + +============================ +BigQuery DataFrames (pandas) +============================ + +.. toctree:: + :maxdepth: 2 + + general_functions + series + frame + indexing + window + groupby + options diff --git a/docs/reference/bigframes.pandas/indexing.rst b/docs/reference/bigframes.pandas/indexing.rst new file mode 100644 index 00000000000..8f7f1947401 --- /dev/null +++ b/docs/reference/bigframes.pandas/indexing.rst @@ -0,0 +1,9 @@ + +============= +Index objects +============= + +.. autoclass:: bigframes.core.indexes.index.Index + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.pandas/options.rst b/docs/reference/bigframes.pandas/options.rst new file mode 100644 index 00000000000..60af8c826a4 --- /dev/null +++ b/docs/reference/bigframes.pandas/options.rst @@ -0,0 +1,6 @@ + +==================== +Options and settings +==================== + +``bigframes.pandas.options`` is an alias for :data:`bigframes.options`. diff --git a/docs/reference/bigframes.pandas/series.rst b/docs/reference/bigframes.pandas/series.rst new file mode 100644 index 00000000000..b179da9ca13 --- /dev/null +++ b/docs/reference/bigframes.pandas/series.rst @@ -0,0 +1,36 @@ + +====== +Series +====== + +.. contents:: Table of Contents + :depth: 2 + :local: + :backlinks: none + +Series +------ + +.. autoclass:: bigframes.series.Series + :members: + :inherited-members: + :undoc-members: + +Accessors +--------- + +Datetime properties +^^^^^^^^^^^^^^^^^^^ + +.. automodule:: bigframes.operations.datetimes + :members: + :inherited-members: + :undoc-members: + +String handling +^^^^^^^^^^^^^^^ + +.. automodule:: bigframes.operations.strings + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.pandas/window.rst b/docs/reference/bigframes.pandas/window.rst new file mode 100644 index 00000000000..55d911ecf4f --- /dev/null +++ b/docs/reference/bigframes.pandas/window.rst @@ -0,0 +1,9 @@ + +====== +Window +====== + +.. autoclass:: bigframes.core.window.Window + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes/index.rst b/docs/reference/bigframes/index.rst new file mode 100644 index 00000000000..76d64444faf --- /dev/null +++ b/docs/reference/bigframes/index.rst @@ -0,0 +1,19 @@ + +Core objects +============ + +.. toctree:: + :maxdepth: 2 + + options + + +Session +------- + +.. autofunction:: bigframes.connect + +.. autoclass:: bigframes.session.Session + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes/options.rst b/docs/reference/bigframes/options.rst new file mode 100644 index 00000000000..914e4976d54 --- /dev/null +++ b/docs/reference/bigframes/options.rst @@ -0,0 +1,12 @@ +Options and settings +==================== + +.. currentmodule:: bigframes + +.. autodata:: options + +.. autoclass:: bigframes._config.Options + +.. autoclass:: bigframes._config.bigquery_options.BigQueryOptions + +.. autoclass:: bigframes._config.display_options.DisplayOptions diff --git a/docs/reference/index.rst b/docs/reference/index.rst new file mode 100644 index 00000000000..c790831db18 --- /dev/null +++ b/docs/reference/index.rst @@ -0,0 +1,12 @@ +API Reference +============= + +Refer to these pages for details about the public objects in the ``bigframes`` +packages. + +.. toctree:: + :maxdepth: 2 + + bigframes/index + bigframes.pandas/index + bigframes.ml/index diff --git a/docs/samples b/docs/samples new file mode 120000 index 00000000000..e804737ed3a --- /dev/null +++ b/docs/samples @@ -0,0 +1 @@ +../samples \ No newline at end of file diff --git a/docs/user_guide/bigframes.pandas/remote_functions.rst b/docs/user_guide/bigframes.pandas/remote_functions.rst new file mode 100644 index 00000000000..3a1a9d342f0 --- /dev/null +++ b/docs/user_guide/bigframes.pandas/remote_functions.rst @@ -0,0 +1,40 @@ + +Using the Remote Functions +========================== + +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 on it via `help` command. + +.. code-block:: python + + import bigframes.pandas as pd + help(pd.remote_function) + +Read a table and inspect the column of interest. + +.. code-block:: python + + df = pd.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. + +.. code-block:: python + + @pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn') + def get_bucket(num): + if not num: return "NA" + boundary = 4000 + return "at_or_above_4000" if num >= boundary else "below_4000" + +Run the custom function on the column of interest to create a new column. + +.. code-block:: python + + df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket)) + df[['body_mass_g', 'body_mass_bucket']].head(10) diff --git a/docs/user_guide/index.rst b/docs/user_guide/index.rst new file mode 100644 index 00000000000..915e172159c --- /dev/null +++ b/docs/user_guide/index.rst @@ -0,0 +1,9 @@ +.. _user_guide: + +User Guide +========== + +.. toctree:: + :maxdepth: 2 + + bigframes.pandas/remote_functions diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 00000000000..ce78c4686ed --- /dev/null +++ b/mypy.ini @@ -0,0 +1,25 @@ +# https://mypy.readthedocs.io/en/stable/config_file.html#config-file + +[mypy] +exclude = ^third_party/ + +[mypy-google.auth.*] +ignore_missing_imports = True + +[mypy-cloudpickle.*] +ignore_missing_imports = True + +[mypy-pydata_google_auth] +ignore_missing_imports = True + +[mypy-google.colab] +ignore_missing_imports = True + +[mypy-pytz] +ignore_missing_imports = True + +[mypy-pyarrow] +ignore_missing_imports = True + +[mypy-ipywidgets] +ignore_missing_imports = True diff --git a/notebooks/00 - Summary.ipynb b/notebooks/00 - Summary.ipynb new file mode 100644 index 00000000000..66ac9a8de84 --- /dev/null +++ b/notebooks/00 - Summary.ipynb @@ -0,0 +1,2060 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the BigQuery DataFrames API" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Set BigQuery DataFrames options" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import bigframes.pandas\n", + "\n", + "bigframes.pandas.options.bigquery.project = \"bigframes-dev\"\n", + "bigframes.pandas.options.bigquery.location = \"us\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize a dataframe for a BigQuery table" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "df = bigframes.pandas.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## View the DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Adelie Penguin (Pygoscelis adeliae)Dream36.618.4184.03475.0FEMALE
1Adelie Penguin (Pygoscelis adeliae)Dream39.819.1184.04650.0MALE
2Adelie Penguin (Pygoscelis adeliae)Dream40.918.9184.03900.0MALE
3Chinstrap penguin (Pygoscelis antarctica)Dream46.517.9192.03500.0FEMALE
4Adelie Penguin (Pygoscelis adeliae)Dream37.316.8192.03000.0FEMALE
5Adelie Penguin (Pygoscelis adeliae)Dream43.218.5192.04100.0MALE
6Chinstrap penguin (Pygoscelis antarctica)Dream46.916.6192.02700.0FEMALE
7Chinstrap penguin (Pygoscelis antarctica)Dream50.518.4200.03400.0FEMALE
8Chinstrap penguin (Pygoscelis antarctica)Dream49.519.0200.03800.0MALE
9Adelie Penguin (Pygoscelis adeliae)Dream40.220.1200.03975.0MALE
10Adelie Penguin (Pygoscelis adeliae)Dream40.818.9208.04300.0MALE
11Adelie Penguin (Pygoscelis adeliae)Dream39.018.7185.03650.0MALE
12Adelie Penguin (Pygoscelis adeliae)Dream37.016.9185.03000.0FEMALE
13Chinstrap penguin (Pygoscelis antarctica)Dream47.017.3185.03700.0FEMALE
14Adelie Penguin (Pygoscelis adeliae)Dream34.017.1185.03400.0FEMALE
15Adelie Penguin (Pygoscelis adeliae)Dream37.016.5185.03400.0FEMALE
16Chinstrap penguin (Pygoscelis antarctica)Dream45.717.3193.03600.0FEMALE
17Chinstrap penguin (Pygoscelis antarctica)Dream50.619.4193.03800.0MALE
18Adelie Penguin (Pygoscelis adeliae)Dream39.717.9193.04250.0MALE
19Adelie Penguin (Pygoscelis adeliae)Dream37.818.1193.03750.0MALE
\n", + "
[344 rows x 7 columns in total]" + ], + "text/plain": [ + " species island culmen_length_mm \\\n", + "0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 \n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 \n", + "2 Adelie Penguin (Pygoscelis adeliae) Dream 40.9 \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream 46.5 \n", + "4 Adelie Penguin (Pygoscelis adeliae) Dream 37.3 \n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream 43.2 \n", + "6 Chinstrap penguin (Pygoscelis antarctica) Dream 46.9 \n", + "7 Chinstrap penguin (Pygoscelis antarctica) Dream 50.5 \n", + "8 Chinstrap penguin (Pygoscelis antarctica) Dream 49.5 \n", + "9 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream 40.8 \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream 39.0 \n", + "12 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "13 Chinstrap penguin (Pygoscelis antarctica) Dream 47.0 \n", + "14 Adelie Penguin (Pygoscelis adeliae) Dream 34.0 \n", + "15 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "16 Chinstrap penguin (Pygoscelis antarctica) Dream 45.7 \n", + "17 Chinstrap penguin (Pygoscelis antarctica) Dream 50.6 \n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream 39.7 \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.8 \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream 46.6 \n", + "21 Chinstrap penguin (Pygoscelis antarctica) Dream 51.3 \n", + "22 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream 36.8 \n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "0 18.4 184.0 3475.0 FEMALE \n", + "1 19.1 184.0 4650.0 MALE \n", + "2 18.9 184.0 3900.0 MALE \n", + "3 17.9 192.0 3500.0 FEMALE \n", + "4 16.8 192.0 3000.0 FEMALE \n", + "5 18.5 192.0 4100.0 MALE \n", + "6 16.6 192.0 2700.0 FEMALE \n", + "7 18.4 200.0 3400.0 FEMALE \n", + "8 19.0 200.0 3800.0 MALE \n", + "9 20.1 200.0 3975.0 MALE \n", + "10 18.9 208.0 4300.0 MALE \n", + "11 18.7 185.0 3650.0 MALE \n", + "12 16.9 185.0 3000.0 FEMALE \n", + "13 17.3 185.0 3700.0 FEMALE \n", + "14 17.1 185.0 3400.0 FEMALE \n", + "15 16.5 185.0 3400.0 FEMALE \n", + "16 17.3 193.0 3600.0 FEMALE \n", + "17 19.4 193.0 3800.0 MALE \n", + "18 17.9 193.0 4250.0 MALE \n", + "19 18.1 193.0 3750.0 MALE \n", + "20 17.8 193.0 3800.0 FEMALE \n", + "21 19.2 193.0 3650.0 MALE \n", + "22 17.1 193.0 3400.0 FEMALE \n", + "23 18.5 193.0 3500.0 FEMALE \n", + "24 18.2 193.0 3775.0 MALE \n", + "...\n", + "\n", + "[344 rows x 7 columns]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### View the column names in the dataframe (aka columns names in the table)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Index(['species', 'island', 'culmen_length_mm', 'culmen_depth_mm',\n", + " 'flipper_length_mm', 'body_mass_g', 'sex'],\n", + " dtype='object')" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.columns" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### View the table schema" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "species string[pyarrow]\n", + "island string[pyarrow]\n", + "culmen_length_mm Float64\n", + "culmen_depth_mm Float64\n", + "flipper_length_mm Float64\n", + "body_mass_g Float64\n", + "sex string[pyarrow]\n", + "dtype: object" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.dtypes" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Select a subset of columns" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandbody_mass_g
0Adelie Penguin (Pygoscelis adeliae)Dream3475.0
1Adelie Penguin (Pygoscelis adeliae)Dream4650.0
2Adelie Penguin (Pygoscelis adeliae)Dream3900.0
3Chinstrap penguin (Pygoscelis antarctica)Dream3500.0
4Adelie Penguin (Pygoscelis adeliae)Dream3000.0
5Adelie Penguin (Pygoscelis adeliae)Dream4100.0
6Chinstrap penguin (Pygoscelis antarctica)Dream2700.0
7Chinstrap penguin (Pygoscelis antarctica)Dream3400.0
8Chinstrap penguin (Pygoscelis antarctica)Dream3800.0
9Adelie Penguin (Pygoscelis adeliae)Dream3975.0
10Adelie Penguin (Pygoscelis adeliae)Dream4300.0
11Adelie Penguin (Pygoscelis adeliae)Dream3650.0
12Adelie Penguin (Pygoscelis adeliae)Dream3000.0
13Chinstrap penguin (Pygoscelis antarctica)Dream3700.0
14Adelie Penguin (Pygoscelis adeliae)Dream3400.0
15Adelie Penguin (Pygoscelis adeliae)Dream3400.0
16Chinstrap penguin (Pygoscelis antarctica)Dream3600.0
17Chinstrap penguin (Pygoscelis antarctica)Dream3800.0
18Adelie Penguin (Pygoscelis adeliae)Dream4250.0
19Adelie Penguin (Pygoscelis adeliae)Dream3750.0
\n", + "
[344 rows x 3 columns in total]" + ], + "text/plain": [ + " species island body_mass_g\n", + "0 Adelie Penguin (Pygoscelis adeliae) Dream 3475.0\n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream 4650.0\n", + "2 Adelie Penguin (Pygoscelis adeliae) Dream 3900.0\n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream 3500.0\n", + "4 Adelie Penguin (Pygoscelis adeliae) Dream 3000.0\n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream 4100.0\n", + "6 Chinstrap penguin (Pygoscelis antarctica) Dream 2700.0\n", + "7 Chinstrap penguin (Pygoscelis antarctica) Dream 3400.0\n", + "8 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", + "9 Adelie Penguin (Pygoscelis adeliae) Dream 3975.0\n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream 3650.0\n", + "12 Adelie Penguin (Pygoscelis adeliae) Dream 3000.0\n", + "13 Chinstrap penguin (Pygoscelis antarctica) Dream 3700.0\n", + "14 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", + "15 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", + "16 Chinstrap penguin (Pygoscelis antarctica) Dream 3600.0\n", + "17 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream 4250.0\n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream 3750.0\n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", + "21 Chinstrap penguin (Pygoscelis antarctica) Dream 3650.0\n", + "22 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream 3500.0\n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream 3775.0\n", + "...\n", + "\n", + "[344 rows x 3 columns]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df[[\n", + " \"species\",\n", + " \"island\",\n", + " \"body_mass_g\",\n", + "]]\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### View the first ten values of a series" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0 3475.0\n", + "1 4650.0\n", + "2 3900.0\n", + "3 3500.0\n", + "4 3000.0\n", + "5 4100.0\n", + "6 2700.0\n", + "7 3400.0\n", + "8 3800.0\n", + "9 3975.0\n", + "Name: body_mass_g, dtype: Float64" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df['body_mass_g'].head(10)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Compute the mean of a series" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "4201.7543859649095" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df['body_mass_g'].mean()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Filter the DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandbody_mass_g
1Adelie Penguin (Pygoscelis adeliae)Dream4650.0
5Adelie Penguin (Pygoscelis adeliae)Dream4100.0
10Adelie Penguin (Pygoscelis adeliae)Dream4300.0
18Adelie Penguin (Pygoscelis adeliae)Dream4250.0
25Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
26Adelie Penguin (Pygoscelis adeliae)Dream4000.0
27Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
28Chinstrap penguin (Pygoscelis antarctica)Dream4300.0
30Chinstrap penguin (Pygoscelis antarctica)Dream4450.0
36Adelie Penguin (Pygoscelis adeliae)Dream4450.0
44Chinstrap penguin (Pygoscelis antarctica)Dream4100.0
45Chinstrap penguin (Pygoscelis antarctica)Dream4800.0
57Chinstrap penguin (Pygoscelis antarctica)Dream4400.0
61Chinstrap penguin (Pygoscelis antarctica)Dream4150.0
66Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
67Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
68Chinstrap penguin (Pygoscelis antarctica)Dream4100.0
74Adelie Penguin (Pygoscelis adeliae)Dream4350.0
77Adelie Penguin (Pygoscelis adeliae)Dream4150.0
78Adelie Penguin (Pygoscelis adeliae)Dream4400.0
\n", + "
[177 rows x 3 columns in total]" + ], + "text/plain": [ + " species island body_mass_g\n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream 4650.0\n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream 4100.0\n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream 4250.0\n", + "25 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", + "26 Adelie Penguin (Pygoscelis adeliae) Dream 4000.0\n", + "27 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", + "28 Chinstrap penguin (Pygoscelis antarctica) Dream 4300.0\n", + "30 Chinstrap penguin (Pygoscelis antarctica) Dream 4450.0\n", + "36 Adelie Penguin (Pygoscelis adeliae) Dream 4450.0\n", + "44 Chinstrap penguin (Pygoscelis antarctica) Dream 4100.0\n", + "45 Chinstrap penguin (Pygoscelis antarctica) Dream 4800.0\n", + "57 Chinstrap penguin (Pygoscelis antarctica) Dream 4400.0\n", + "61 Chinstrap penguin (Pygoscelis antarctica) Dream 4150.0\n", + "66 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", + "67 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", + "68 Chinstrap penguin (Pygoscelis antarctica) Dream 4100.0\n", + "74 Adelie Penguin (Pygoscelis adeliae) Dream 4350.0\n", + "77 Adelie Penguin (Pygoscelis adeliae) Dream 4150.0\n", + "78 Adelie Penguin (Pygoscelis adeliae) Dream 4400.0\n", + "80 Chinstrap penguin (Pygoscelis antarctica) Dream 4300.0\n", + "90 Chinstrap penguin (Pygoscelis antarctica) Dream 4150.0\n", + "92 Chinstrap penguin (Pygoscelis antarctica) Dream 4500.0\n", + "93 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", + "94 Chinstrap penguin (Pygoscelis antarctica) Dream 4550.0\n", + "...\n", + "\n", + "[177 rows x 3 columns]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df[df['body_mass_g'] >= 4000.0]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the Remote Functions" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### BigQuery DataFrames gives you the ability to turn your custom scalar functions into a BigQuery remote function.\n", + "\n", + "It requires the GCP project to be set up appropriately and the user having sufficient privileges to use them. One can find more details on it via `help` command." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Help on function remote_function in module bigframes.pandas:\n", + "\n", + "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True)\n", + " Decorator to turn a user defined function into a BigQuery remote function.\n", + " \n", + " Args:\n", + " input_types (list(type)):\n", + " List of input data types in the user defined function.\n", + " output_type (type):\n", + " Data type of the output in the user defined function.\n", + " dataset (str, Optional):\n", + " Dataset to use to create a BigQuery function. It should be in\n", + " `.` or `` format. If this\n", + " param is not provided then session dataset id would be used.\n", + " bigquery_connection (str, Optional):\n", + " Name of the BigQuery connection. If it is pre created in the same\n", + " location as the `bigquery_client.location` then it would be used,\n", + " otherwise it would be created dynamically assuming the user has\n", + " necessary priviliges. If this param is not provided then the\n", + " bigquery connection from the session would be used.\n", + " reuse (bool, Optional):\n", + " Reuse the remote function if already exists.\n", + " `True` by default, which will result in reusing an existing remote\n", + " function (if any) that was previously created for the same udf.\n", + " Setting it to false would force creating a unique remote function.\n", + " If the required remote function does not exist then it would be\n", + " created irrespective of this param.\n", + " \n", + " Notes:\n", + " Please make sure following is setup before using this API:\n", + " \n", + " 1. Have the below APIs enabled for your project:\n", + " a. BigQuery Connection API\n", + " b. Cloud Functions API\n", + " c. Cloud Run API\n", + " d. Cloud Build API\n", + " e. Artifact Registry API\n", + " f. Cloud Resource Manager API\n", + " \n", + " This can be done from the cloud console (change PROJECT_ID to yours):\n", + " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", + " Or from the gcloud CLI:\n", + " $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com\n", + " \n", + " 2. Have following IAM roles enabled for you:\n", + " a. BigQuery Data Editor (roles/bigquery.dataEditor)\n", + " b. BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", + " c. Cloud Functions Developer (roles/cloudfunctions.developer)\n", + " d. Service Account User (roles/iam.serviceAccountUser)\n", + " e. Storage Object Viewer (roles/storage.objectViewer)\n", + " f. Project IAM Admin (roles/resourcemanager.projectIamAdmin)\n", + " (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", + " \n", + " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", + " a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", + " b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", + " Alternatively, the IAM could also be setup via the gcloud CLI:\n", + " $ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"\n", + "\n" + ] + } + ], + "source": [ + "import bigframes.pandas as pd\n", + "help(pd.remote_function)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Define a custom function, and specify the intent to turn it into a remote function.\n", + "\n", + "It requires a BigQuery connection. If the connection is not already created,\n", + "the BigQuery DataFrames package attempts to create one assuming the necessary\n", + "APIs and IAM permissions are setup in the project." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[INFO][2023-06-28 23:31:49,355][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-b3fab64f5997ad6a516379defe8d4202 --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmp9w5e89lh --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", + "Preparing function...\n", + ".done.\n", + "Deploying function...\n", + "[Build]..........................................................................................................................................................................................................................................................................................................................................................................................................................done\n", + "[Service].........................................................................................................................................................................................................done\n", + "Done.\n", + "You can view your function in the Cloud Console here: https://console.cloud.google.com/functions/details/us-central1/bigframes-b3fab64f5997ad6a516379defe8d4202?project=bigframes-dev\n", + "\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "buildConfig:\n", + " build: projects/1084210331973/locations/us-central1/builds/780b1780-9b38-4515-ae60-89d05454ef83\n", + " entryPoint: udf_http\n", + " runtime: python310\n", + " source:\n", + " storageSource:\n", + " bucket: gcf-v2-sources-1084210331973-us-central1\n", + " object: bigframes-b3fab64f5997ad6a516379defe8d4202/function-source.zip\n", + " sourceProvenance:\n", + " resolvedStorageSource:\n", + " bucket: gcf-v2-sources-1084210331973-us-central1\n", + " generation: '1687995112300727'\n", + " object: bigframes-b3fab64f5997ad6a516379defe8d4202/function-source.zip\n", + "environment: GEN_2\n", + "labels:\n", + " deployment-tool: cli-gcloud\n", + "name: projects/bigframes-dev/locations/us-central1/functions/bigframes-b3fab64f5997ad6a516379defe8d4202\n", + "serviceConfig:\n", + " allTrafficOnLatestRevision: true\n", + " availableCpu: '0.1666'\n", + " availableMemory: 256M\n", + " ingressSettings: ALLOW_ALL\n", + " maxInstanceCount: 100\n", + " maxInstanceRequestConcurrency: 1\n", + " revision: bigframes-b3fab64f5997ad6a516379defe8d4202-00001-tut\n", + " service: projects/bigframes-dev/locations/us-central1/services/bigframes-b3fab64f5997ad6a516379defe8d4202\n", + " serviceAccountEmail: 1084210331973-compute@developer.gserviceaccount.com\n", + " timeoutSeconds: 60\n", + " uri: https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app\n", + "state: ACTIVE\n", + "updateTime: '2023-06-28T23:32:51.911131997Z'\n", + "url: https://us-central1-bigframes-dev.cloudfunctions.net/bigframes-b3fab64f5997ad6a516379defe8d4202\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[INFO][2023-06-28 23:32:55,330][bigframes.remote_function] Successfully created cloud function bigframes-b3fab64f5997ad6a516379defe8d4202 with uri (https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app)\n", + "[INFO][2023-06-28 23:32:59,378][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", + "[INFO][2023-06-28 23:32:59,379][bigframes.remote_function] Creating BQ remote function: \n", + " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_b3fab64f5997ad6a516379defe8d4202(num FLOAT64)\n", + " RETURNS STRING\n", + " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", + " OPTIONS (\n", + " endpoint = \"https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app\"\n", + " )\n", + "[INFO][2023-06-28 23:33:00,338][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_b3fab64f5997ad6a516379defe8d4202\n" + ] + } + ], + "source": [ + "@pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", + "def get_bucket(num):\n", + " if not num: return \"NA\"\n", + " boundary = 4000\n", + " return \"at_or_above_4000\" if num >= boundary else \"below_4000\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run the custom function on the BigQuery-backed dataframe" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
body_mass_gbody_mass_bucket
03475.0below_4000
14650.0at_or_above_4000
23900.0below_4000
33500.0below_4000
43000.0below_4000
54100.0at_or_above_4000
62700.0below_4000
73400.0below_4000
83800.0below_4000
93975.0below_4000
\n", + "
[10 rows x 2 columns in total]" + ], + "text/plain": [ + " body_mass_g body_mass_bucket\n", + "0 3475.0 below_4000\n", + "1 4650.0 at_or_above_4000\n", + "2 3900.0 below_4000\n", + "3 3500.0 below_4000\n", + "4 3000.0 below_4000\n", + "5 4100.0 at_or_above_4000\n", + "6 2700.0 below_4000\n", + "7 3400.0 below_4000\n", + "8 3800.0 below_4000\n", + "9 3975.0 below_4000\n", + "\n", + "[10 rows x 2 columns]" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket))\n", + "df[['body_mass_g', 'body_mass_bucket']].head(10)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the ML API" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize a DataFrame from a BigQuery table" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Adelie Penguin (Pygoscelis adeliae)Dream36.618.4184.03475.0FEMALE
1Adelie Penguin (Pygoscelis adeliae)Dream39.819.1184.04650.0MALE
2Adelie Penguin (Pygoscelis adeliae)Dream40.918.9184.03900.0MALE
3Chinstrap penguin (Pygoscelis antarctica)Dream46.517.9192.03500.0FEMALE
4Adelie Penguin (Pygoscelis adeliae)Dream37.316.8192.03000.0FEMALE
5Adelie Penguin (Pygoscelis adeliae)Dream43.218.5192.04100.0MALE
6Chinstrap penguin (Pygoscelis antarctica)Dream46.916.6192.02700.0FEMALE
7Chinstrap penguin (Pygoscelis antarctica)Dream50.518.4200.03400.0FEMALE
8Chinstrap penguin (Pygoscelis antarctica)Dream49.519.0200.03800.0MALE
9Adelie Penguin (Pygoscelis adeliae)Dream40.220.1200.03975.0MALE
10Adelie Penguin (Pygoscelis adeliae)Dream40.818.9208.04300.0MALE
11Adelie Penguin (Pygoscelis adeliae)Dream39.018.7185.03650.0MALE
12Adelie Penguin (Pygoscelis adeliae)Dream37.016.9185.03000.0FEMALE
13Chinstrap penguin (Pygoscelis antarctica)Dream47.017.3185.03700.0FEMALE
14Adelie Penguin (Pygoscelis adeliae)Dream34.017.1185.03400.0FEMALE
15Adelie Penguin (Pygoscelis adeliae)Dream37.016.5185.03400.0FEMALE
16Chinstrap penguin (Pygoscelis antarctica)Dream45.717.3193.03600.0FEMALE
17Chinstrap penguin (Pygoscelis antarctica)Dream50.619.4193.03800.0MALE
18Adelie Penguin (Pygoscelis adeliae)Dream39.717.9193.04250.0MALE
19Adelie Penguin (Pygoscelis adeliae)Dream37.818.1193.03750.0MALE
\n", + "
[344 rows x 7 columns in total]" + ], + "text/plain": [ + " species island culmen_length_mm \\\n", + "0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 \n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 \n", + "2 Adelie Penguin (Pygoscelis adeliae) Dream 40.9 \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream 46.5 \n", + "4 Adelie Penguin (Pygoscelis adeliae) Dream 37.3 \n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream 43.2 \n", + "6 Chinstrap penguin (Pygoscelis antarctica) Dream 46.9 \n", + "7 Chinstrap penguin (Pygoscelis antarctica) Dream 50.5 \n", + "8 Chinstrap penguin (Pygoscelis antarctica) Dream 49.5 \n", + "9 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream 40.8 \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream 39.0 \n", + "12 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "13 Chinstrap penguin (Pygoscelis antarctica) Dream 47.0 \n", + "14 Adelie Penguin (Pygoscelis adeliae) Dream 34.0 \n", + "15 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "16 Chinstrap penguin (Pygoscelis antarctica) Dream 45.7 \n", + "17 Chinstrap penguin (Pygoscelis antarctica) Dream 50.6 \n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream 39.7 \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.8 \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream 46.6 \n", + "21 Chinstrap penguin (Pygoscelis antarctica) Dream 51.3 \n", + "22 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream 36.8 \n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "0 18.4 184.0 3475.0 FEMALE \n", + "1 19.1 184.0 4650.0 MALE \n", + "2 18.9 184.0 3900.0 MALE \n", + "3 17.9 192.0 3500.0 FEMALE \n", + "4 16.8 192.0 3000.0 FEMALE \n", + "5 18.5 192.0 4100.0 MALE \n", + "6 16.6 192.0 2700.0 FEMALE \n", + "7 18.4 200.0 3400.0 FEMALE \n", + "8 19.0 200.0 3800.0 MALE \n", + "9 20.1 200.0 3975.0 MALE \n", + "10 18.9 208.0 4300.0 MALE \n", + "11 18.7 185.0 3650.0 MALE \n", + "12 16.9 185.0 3000.0 FEMALE \n", + "13 17.3 185.0 3700.0 FEMALE \n", + "14 17.1 185.0 3400.0 FEMALE \n", + "15 16.5 185.0 3400.0 FEMALE \n", + "16 17.3 193.0 3600.0 FEMALE \n", + "17 19.4 193.0 3800.0 MALE \n", + "18 17.9 193.0 4250.0 MALE \n", + "19 18.1 193.0 3750.0 MALE \n", + "20 17.8 193.0 3800.0 FEMALE \n", + "21 19.2 193.0 3650.0 MALE \n", + "22 17.1 193.0 3400.0 FEMALE \n", + "23 18.5 193.0 3500.0 FEMALE \n", + "24 18.2 193.0 3775.0 MALE \n", + "...\n", + "\n", + "[344 rows x 7 columns]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = bigframes.pandas.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Clean and prepare the data" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
islandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Dream36.618.4184.03475.0FEMALE
1Dream39.819.1184.04650.0MALE
2Dream40.918.9184.03900.0MALE
4Dream37.316.8192.03000.0FEMALE
5Dream43.218.5192.04100.0MALE
9Dream40.220.1200.03975.0MALE
10Dream40.818.9208.04300.0MALE
11Dream39.018.7185.03650.0MALE
12Dream37.016.9185.03000.0FEMALE
14Dream34.017.1185.03400.0FEMALE
15Dream37.016.5185.03400.0FEMALE
18Dream39.717.9193.04250.0MALE
19Dream37.818.1193.03750.0MALE
22Dream40.217.1193.03400.0FEMALE
23Dream36.818.5193.03500.0FEMALE
26Dream41.518.5201.04000.0MALE
31Dream33.116.1178.02900.0FEMALE
32Dream37.218.1178.03900.0MALE
33Dream39.516.7178.03250.0FEMALE
35Dream36.018.5186.03100.0FEMALE
\n", + "
[146 rows x 6 columns in total]" + ], + "text/plain": [ + " island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", + "0 Dream 36.6 18.4 184.0 3475.0 \n", + "1 Dream 39.8 19.1 184.0 4650.0 \n", + "2 Dream 40.9 18.9 184.0 3900.0 \n", + "4 Dream 37.3 16.8 192.0 3000.0 \n", + "5 Dream 43.2 18.5 192.0 4100.0 \n", + "9 Dream 40.2 20.1 200.0 3975.0 \n", + "10 Dream 40.8 18.9 208.0 4300.0 \n", + "11 Dream 39.0 18.7 185.0 3650.0 \n", + "12 Dream 37.0 16.9 185.0 3000.0 \n", + "14 Dream 34.0 17.1 185.0 3400.0 \n", + "15 Dream 37.0 16.5 185.0 3400.0 \n", + "18 Dream 39.7 17.9 193.0 4250.0 \n", + "19 Dream 37.8 18.1 193.0 3750.0 \n", + "22 Dream 40.2 17.1 193.0 3400.0 \n", + "23 Dream 36.8 18.5 193.0 3500.0 \n", + "26 Dream 41.5 18.5 201.0 4000.0 \n", + "31 Dream 33.1 16.1 178.0 2900.0 \n", + "32 Dream 37.2 18.1 178.0 3900.0 \n", + "33 Dream 39.5 16.7 178.0 3250.0 \n", + "35 Dream 36.0 18.5 186.0 3100.0 \n", + "36 Dream 39.6 18.1 186.0 4450.0 \n", + "38 Dream 41.3 20.3 194.0 3550.0 \n", + "41 Dream 35.7 18.0 202.0 3550.0 \n", + "51 Dream 38.1 17.6 187.0 3425.0 \n", + "53 Dream 36.0 17.1 187.0 3700.0 \n", + "\n", + " sex \n", + "0 FEMALE \n", + "1 MALE \n", + "2 MALE \n", + "4 FEMALE \n", + "5 MALE \n", + "9 MALE \n", + "10 MALE \n", + "11 MALE \n", + "12 FEMALE \n", + "14 FEMALE \n", + "15 FEMALE \n", + "18 MALE \n", + "19 MALE \n", + "22 FEMALE \n", + "23 FEMALE \n", + "26 MALE \n", + "31 FEMALE \n", + "32 MALE \n", + "33 FEMALE \n", + "35 FEMALE \n", + "36 MALE \n", + "38 MALE \n", + "41 FEMALE \n", + "51 FEMALE \n", + "53 FEMALE \n", + "...\n", + "\n", + "[146 rows x 6 columns]" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# filter down to the data we want to analyze\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# drop the columns we don't care about\n", + "adelie_data = adelie_data.drop(columns=[\"species\"])\n", + "\n", + "# drop rows with nulls to get our training data\n", + "training_data = adelie_data.dropna()\n", + "\n", + "# take a peek at the training data\n", + "training_data" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "# pick feature columns and label column\n", + "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "label_columns = training_data[['body_mass_g']]\n", + "\n", + "# also get the rows that we want to make predictions for (i.e. where the feature column is null)\n", + "missing_body_mass = adelie_data[adelie_data.body_mass_g.isnull()]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Train and evaluate a linear regression model using the ML API" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.87876378553.6016340.005614181.3309110.6239510.623951
\n", + "
[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 223.878763 78553.601634 0.005614 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 181.330911 0.623951 0.623951 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "\n", + "# as in scikit-learn, a newly created model is just a bundle of parameters\n", + "# default parameters are fine here\n", + "model = LinearRegression()\n", + "\n", + "# this will train a temporary model in BigQuery Machine Learning\n", + "model.fit(feature_columns, label_columns)\n", + "\n", + "# check how the model performed\n", + "model.score(feature_columns, label_columns)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Make predictions using the model" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
2923603.735118
\n", + "
[1 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "292 3603.735118\n", + "\n", + "[1 rows x 1 columns]" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.predict(missing_body_mass)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Save the trained model to BigQuery, so we can load it later" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.to_gbq(\"bqml_tutorial.penguins_model\", replace=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/01 - Getting Started.ipynb b/notebooks/01 - Getting Started.ipynb new file mode 100644 index 00000000000..132cdf8db0d --- /dev/null +++ b/notebooks/01 - Getting Started.ipynb @@ -0,0 +1,965 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "d7a03de2-c0ef-4f80-9cd5-f96e87cf2d54", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# On the instance where you are running jupyter,\n", + "# authenticate with gcloud first:\n", + "#\n", + "# gcloud auth application-default login\n", + "\n", + "import bigframes.pandas as bpd\n", + "\n", + "# Change this location to the location of your datasets.\n", + "# We use \"us\" as that is the location of the sample data.\n", + "bpd.options.bigquery.location = \"us\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "941cb6c3-8c54-42ce-a945-4fa604176b2e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "df = bpd.read_gbq(\n", + " \"bigquery-public-data.chicago_taxi_trips.taxi_trips\",\n", + " # For large tables, a unique index or multiindex is highly recommended.\n", + " # This prevents BigQuery DataFrames from creating a default sequential index.\n", + " index_col=[\"trip_start_timestamp\", \"unique_key\"],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "4c3686d1-8b12-4c6b-8b19-6ef40a3b8ef5", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "bigframes.dataframe.DataFrame" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(df)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "586da7b2-8825-4b1e-87cd-512df60775e1", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
taxi_idtrip_end_timestamptrip_secondstrip_milespickup_census_tractdropoff_census_tractpickup_community_areadropoff_community_areafaretips...extrastrip_totalpayment_typecompanypickup_latitudepickup_longitudepickup_locationdropoff_latitudedropoff_longitudedropoff_location
trip_start_timestampunique_key
2013-01-01 00:00:00+00:00006d53094ff5b3c81140c00aa91fdb5467db08024d924c1ad8126ad7efcbe49964b63ce4e62e11223b0f41...2013-01-01 00:15:00+00:003602.3170313301001703132010033327.850.0...2.09.85Cash<NA>41.85935-87.617358POINT (-87.6173580061 41.859349715)41.884987-87.620993POINT (-87.6209929134 41.8849871918)
\n", + "

1 rows × 21 columns

\n", + "
[1 rows x 21 columns in total]" + ], + "text/plain": [ + " taxi_id \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 4d924c1ad8126ad7efcbe49964b63ce4e62e11223b0f41... \n", + "\n", + " trip_end_timestamp \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 2013-01-01 00:15:00+00:00 \n", + "\n", + " trip_seconds \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 360 \n", + "\n", + " trip_miles \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 2.3 \n", + "\n", + " pickup_census_tract \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 17031330100 \n", + "\n", + " dropoff_census_tract \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 17031320100 \n", + "\n", + " pickup_community_area \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 33 \n", + "\n", + " dropoff_community_area \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 32 \n", + "\n", + " fare \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 7.85 \n", + "\n", + " tips ... \\\n", + "trip_start_timestamp unique_key ... \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 0.0 ... \n", + "\n", + " extras \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 2.0 \n", + "\n", + " trip_total \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 9.85 \n", + "\n", + " payment_type \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 Cash \n", + "\n", + " company \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 \n", + "\n", + " pickup_latitude \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 41.85935 \n", + "\n", + " pickup_longitude \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 -87.617358 \n", + "\n", + " pickup_location \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 POINT (-87.6173580061 41.859349715) \n", + "\n", + " dropoff_latitude \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 41.884987 \n", + "\n", + " dropoff_longitude \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 -87.620993 \n", + "\n", + " dropoff_location \n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 POINT (-87.6209929134 41.8849871918) \n", + "\n", + "[1 rows x 21 columns]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.head(n=1)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "75a1ce88", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
taxi_idtrip_end_timestamptrip_secondstrip_milespickup_census_tractdropoff_census_tractpickup_community_areadropoff_community_areafaretips...extrastrip_totalpayment_typecompanypickup_latitudepickup_longitudepickup_locationdropoff_latitudedropoff_longitudedropoff_location
trip_start_timestampunique_key
2023-07-01 00:00:00+00:00fb004dfe98302ccf34e5e11bf1081568a94843a275cf3a53aae5e5858361a7ca64f75d3407dc0a44d7bc42...2023-07-01 00:00:00+00:005621.82<NA><NA>22229.52.53...0.012.03MobileCity Service41.922761-87.699155POINT (-87.6991553432 41.9227606205)41.922761-87.699155POINT (-87.6991553432 41.9227606205)
\n", + "

1 rows × 21 columns

\n", + "
[1 rows x 21 columns in total]" + ], + "text/plain": [ + " taxi_id \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 75cf3a53aae5e5858361a7ca64f75d3407dc0a44d7bc42... \n", + "\n", + " trip_end_timestamp \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 2023-07-01 00:00:00+00:00 \n", + "\n", + " trip_seconds \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 562 \n", + "\n", + " trip_miles \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 1.82 \n", + "\n", + " pickup_census_tract \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 \n", + "\n", + " dropoff_census_tract \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 \n", + "\n", + " pickup_community_area \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 22 \n", + "\n", + " dropoff_community_area \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 22 \n", + "\n", + " fare \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 9.5 \n", + "\n", + " tips ... \\\n", + "trip_start_timestamp unique_key ... \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 2.53 ... \n", + "\n", + " extras \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 0.0 \n", + "\n", + " trip_total \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 12.03 \n", + "\n", + " payment_type \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 Mobile \n", + "\n", + " company \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 City Service \n", + "\n", + " pickup_latitude \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 41.922761 \n", + "\n", + " pickup_longitude \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 -87.699155 \n", + "\n", + " pickup_location \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 POINT (-87.6991553432 41.9227606205) \n", + "\n", + " dropoff_latitude \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 41.922761 \n", + "\n", + " dropoff_longitude \\\n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 -87.699155 \n", + "\n", + " dropoff_location \n", + "trip_start_timestamp unique_key \n", + "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 POINT (-87.6991553432 41.9227606205) \n", + "\n", + "[1 rows x 21 columns]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.tail(n=1)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f1eee1c6-7214-440e-bc17-6839d53a6718", + "metadata": {}, + "outputs": [], + "source": [ + "df = df[[\n", + " \"company\",\n", + " \"trip_miles\",\n", + " \"fare\",\n", + " \"tips\",\n", + "]]" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "66071984-e371-4161-8d7e-00d810b5cfab", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
companytrip_milesfaretips
trip_start_timestampunique_key
2013-01-01 00:00:00+00:00006d53094ff5b3c81140c00aa91fdb5467db0802<NA>2.37.850.0
00a723cf08cf29c56faa8ef77be519f6590b3405<NA>1.05.650.0
01396f0debefe45ccce23e48dd471edd67bd32e8Dispatch Taxi Affiliation0.08.850.0
01e9a03fd793670ed35ef7195eeb99775895611fDispatch Taxi Affiliation1.87.450.0
024f0f183c43ad2be33f232fd8dca0a86a3e4925<NA>4.914.050.0
04373ef74229c2122ff90d92642fcd94c88302e4Dispatch Taxi Affiliation0.06.650.0
0501a22fba176594d6f974360e20737af1b6cb66<NA>2.27.650.0
05faff951f2a2703fe3eb24efac026c2a1669a92<NA>3.39.650.0
0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbeBlue Ribbon Taxi Association Inc.0.06.250.0
06b4dc6dcd6295e69284e1c9d5a98015b0588ea6Choice Taxi Association0.013.850.0
07483a0baa6df0af68696b6f2483b7f429ff1827Northwest Management LLC0.64.650.0
07585d085f65cd512b94de09684253fe0ffb31a2Taxi Affiliation Services7.04.650.0
090ea9dc5dc26374ee06cdab1b809895beb1befcDispatch Taxi Affiliation0.010.050.0
091eda4ec41520d6ba96003f05eb10744f73aa82Northwest Management LLC0.015.250.0
097f2efed248cb5900792a89bba96a64918345ceBlue Ribbon Taxi Association Inc.0.012.250.0
09ffa3f40d36dbf8d51ca903738d9bba92557ed7Taxi Affiliation Services0.06.450.0
0a3446fb4f72d6485548883cd7f6cfe068d44829Taxi Affiliation Services0.014.250.0
0a7ea707a8700555e898489200a66b546a834170<NA>5.0413.852.0
0be545d9934dd343a2d929fca2e53d3d2851137bDispatch Taxi Affiliation1.15.650.0
0c2b3705c9748a5949bf5b7d727b2d1f2e18805f<NA>12.326.250.0
0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4Taxi Affiliation Services0.015.250.0
0dafcce34426b5377fd3925418a4c24b3e45c7d9<NA>1.26.250.0
0f5f7162b289dd9743b9d2e7fc75bf69696bb7bcTaxi Affiliation Services0.09.850.0
0f63007e437d589c9139acb026f03a832da3240fNorthwest Management LLC1.56.650.0
10508745a1069e9cdaad8760a6a243dc924cc0f2<NA>5.513.250.0
\n", + "

25 rows × 4 columns

\n", + "
[208395018 rows x 4 columns in total]" + ], + "text/plain": [ + " company \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 \n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 \n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 Dispatch Taxi Affiliation \n", + " 01e9a03fd793670ed35ef7195eeb99775895611f Dispatch Taxi Affiliation \n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 \n", + " 04373ef74229c2122ff90d92642fcd94c88302e4 Dispatch Taxi Affiliation \n", + " 0501a22fba176594d6f974360e20737af1b6cb66 \n", + " 05faff951f2a2703fe3eb24efac026c2a1669a92 \n", + " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe Blue Ribbon Taxi Association Inc. \n", + " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 Choice Taxi Association \n", + " 07483a0baa6df0af68696b6f2483b7f429ff1827 Northwest Management LLC \n", + " 07585d085f65cd512b94de09684253fe0ffb31a2 Taxi Affiliation Services \n", + " 090ea9dc5dc26374ee06cdab1b809895beb1befc Dispatch Taxi Affiliation \n", + " 091eda4ec41520d6ba96003f05eb10744f73aa82 Northwest Management LLC \n", + " 097f2efed248cb5900792a89bba96a64918345ce Blue Ribbon Taxi Association Inc. \n", + " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 Taxi Affiliation Services \n", + " 0a3446fb4f72d6485548883cd7f6cfe068d44829 Taxi Affiliation Services \n", + " 0a7ea707a8700555e898489200a66b546a834170 \n", + " 0be545d9934dd343a2d929fca2e53d3d2851137b Dispatch Taxi Affiliation \n", + " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f \n", + " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 Taxi Affiliation Services \n", + " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 \n", + " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc Taxi Affiliation Services \n", + " 0f63007e437d589c9139acb026f03a832da3240f Northwest Management LLC \n", + " 10508745a1069e9cdaad8760a6a243dc924cc0f2 \n", + "\n", + " trip_miles \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 2.3 \n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 1.0 \n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 0.0 \n", + " 01e9a03fd793670ed35ef7195eeb99775895611f 1.8 \n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 4.9 \n", + " 04373ef74229c2122ff90d92642fcd94c88302e4 0.0 \n", + " 0501a22fba176594d6f974360e20737af1b6cb66 2.2 \n", + " 05faff951f2a2703fe3eb24efac026c2a1669a92 3.3 \n", + " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 0.0 \n", + " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 0.0 \n", + " 07483a0baa6df0af68696b6f2483b7f429ff1827 0.6 \n", + " 07585d085f65cd512b94de09684253fe0ffb31a2 7.0 \n", + " 090ea9dc5dc26374ee06cdab1b809895beb1befc 0.0 \n", + " 091eda4ec41520d6ba96003f05eb10744f73aa82 0.0 \n", + " 097f2efed248cb5900792a89bba96a64918345ce 0.0 \n", + " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 0.0 \n", + " 0a3446fb4f72d6485548883cd7f6cfe068d44829 0.0 \n", + " 0a7ea707a8700555e898489200a66b546a834170 5.04 \n", + " 0be545d9934dd343a2d929fca2e53d3d2851137b 1.1 \n", + " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 12.3 \n", + " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 0.0 \n", + " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 1.2 \n", + " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 0.0 \n", + " 0f63007e437d589c9139acb026f03a832da3240f 1.5 \n", + " 10508745a1069e9cdaad8760a6a243dc924cc0f2 5.5 \n", + "\n", + " fare \\\n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 7.85 \n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 5.65 \n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 8.85 \n", + " 01e9a03fd793670ed35ef7195eeb99775895611f 7.45 \n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 14.05 \n", + " 04373ef74229c2122ff90d92642fcd94c88302e4 6.65 \n", + " 0501a22fba176594d6f974360e20737af1b6cb66 7.65 \n", + " 05faff951f2a2703fe3eb24efac026c2a1669a92 9.65 \n", + " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 6.25 \n", + " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 13.85 \n", + " 07483a0baa6df0af68696b6f2483b7f429ff1827 4.65 \n", + " 07585d085f65cd512b94de09684253fe0ffb31a2 4.65 \n", + " 090ea9dc5dc26374ee06cdab1b809895beb1befc 10.05 \n", + " 091eda4ec41520d6ba96003f05eb10744f73aa82 15.25 \n", + " 097f2efed248cb5900792a89bba96a64918345ce 12.25 \n", + " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 6.45 \n", + " 0a3446fb4f72d6485548883cd7f6cfe068d44829 14.25 \n", + " 0a7ea707a8700555e898489200a66b546a834170 13.85 \n", + " 0be545d9934dd343a2d929fca2e53d3d2851137b 5.65 \n", + " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 26.25 \n", + " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 15.25 \n", + " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 6.25 \n", + " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 9.85 \n", + " 0f63007e437d589c9139acb026f03a832da3240f 6.65 \n", + " 10508745a1069e9cdaad8760a6a243dc924cc0f2 13.25 \n", + "\n", + " tips \n", + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 0.0 \n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 0.0 \n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 0.0 \n", + " 01e9a03fd793670ed35ef7195eeb99775895611f 0.0 \n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 0.0 \n", + " 04373ef74229c2122ff90d92642fcd94c88302e4 0.0 \n", + " 0501a22fba176594d6f974360e20737af1b6cb66 0.0 \n", + " 05faff951f2a2703fe3eb24efac026c2a1669a92 0.0 \n", + " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 0.0 \n", + " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 0.0 \n", + " 07483a0baa6df0af68696b6f2483b7f429ff1827 0.0 \n", + " 07585d085f65cd512b94de09684253fe0ffb31a2 0.0 \n", + " 090ea9dc5dc26374ee06cdab1b809895beb1befc 0.0 \n", + " 091eda4ec41520d6ba96003f05eb10744f73aa82 0.0 \n", + " 097f2efed248cb5900792a89bba96a64918345ce 0.0 \n", + " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 0.0 \n", + " 0a3446fb4f72d6485548883cd7f6cfe068d44829 0.0 \n", + " 0a7ea707a8700555e898489200a66b546a834170 2.0 \n", + " 0be545d9934dd343a2d929fca2e53d3d2851137b 0.0 \n", + " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 0.0 \n", + " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 0.0 \n", + " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 0.0 \n", + " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 0.0 \n", + " 0f63007e437d589c9139acb026f03a832da3240f 0.0 \n", + " 10508745a1069e9cdaad8760a6a243dc924cc0f2 0.0 \n", + "...\n", + "\n", + "[208395018 rows x 4 columns]" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "e51b687e-0282-459c-8fd0-0ca22cfb153d", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 7.85\n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 5.65\n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 8.85\n", + " 01e9a03fd793670ed35ef7195eeb99775895611f 7.45\n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 14.05\n", + " 04373ef74229c2122ff90d92642fcd94c88302e4 6.65\n", + " 0501a22fba176594d6f974360e20737af1b6cb66 7.65\n", + " 05faff951f2a2703fe3eb24efac026c2a1669a92 9.65\n", + " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 6.25\n", + " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 13.85\n", + "Name: fare, dtype: Float64" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df['fare'].head(n=10)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "ef9d466e-1138-41a7-ad4a-51e09c05870c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "13.669183161039143" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df['fare'].mean()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "eb800cbb-d7d9-4562-8c23-4dd6c06759cc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "trip_start_timestamp unique_key \n", + "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 7.85\n", + " 00a723cf08cf29c56faa8ef77be519f6590b3405 5.65\n", + " 01396f0debefe45ccce23e48dd471edd67bd32e8 8.85\n", + " 01e9a03fd793670ed35ef7195eeb99775895611f 7.45\n", + " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 14.05\n", + "dtype: Float64" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "subtotal = df['fare'] + df['tips']\n", + "subtotal.head()" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "d2ffa53f-e158-4b0b-898e-825e3bd99881", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "2844629.7599999933" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "subtotal[df['company'] == 'Norshore Cab'].sum()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "737dca2c", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/02 - DataFrame.ipynb b/notebooks/02 - DataFrame.ipynb new file mode 100644 index 00000000000..d23c5d01ee9 --- /dev/null +++ b/notebooks/02 - DataFrame.ipynb @@ -0,0 +1,4158 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "47439dbd-4e54-4954-8b16-edc4bcd4f855", + "metadata": {}, + "source": [ + "# Demo to Show DataFrame Operations" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "72ebb083-f06b-4408-b24d-f349bd0851e3", + "metadata": {}, + "outputs": [], + "source": [ + "# On the instance where you are running jupyter,\n", + "# authenticate with gcloud first:\n", + "#\n", + "# gcloud auth application-default login\n", + "\n", + "import bigframes.pandas as bpd\n", + "\n", + "bpd.options.bigquery.location = \"US\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "96757c59-fc22-420e-a42f-c6cb956110ec", + "metadata": {}, + "outputs": [], + "source": [ + "df = bpd.read_gbq(\"bigquery-public-data.baseball.schedules\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "13861abc-120c-4db6-ad0c-e414b85d3443", + "metadata": {}, + "source": [ + "### Select a subset of the DF" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "05cb36e9-bb75-4f6f-8eb6-e4219df6e1d2", + "metadata": {}, + "outputs": [], + "source": [ + "df = df[[\"gameId\", \"year\", \"homeTeamName\", \"awayTeamName\", \"duration_minutes\"]]" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "f124234c-d292-4520-b9cc-a162e2b16d91", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df\n", + "# Here starts real execution, may take a while" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "a6b8b3ac-1df8-46ff-ac4f-d6e7657fc80c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(2431, 5)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.shape" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "efc9ccb5-4dee-4fb8-a95f-3a654421cec4", + "metadata": {}, + "source": [ + "### Retrieve properties of the DF." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "34457cc7-e734-4e3f-9f2b-34cdd4e2aba4", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "gameId string[pyarrow]\n", + "year Int64\n", + "homeTeamName string[pyarrow]\n", + "awayTeamName string[pyarrow]\n", + "duration_minutes Int64\n", + "dtype: object" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.dtypes" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "b4f4383f-f596-41d8-aad2-2fd68d261cfd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Index(['gameId', 'year', 'homeTeamName', 'awayTeamName', 'duration_minutes'], dtype='object')" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.columns" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "70e0a637-52ac-47cd-93e2-75d06d7a6170", + "metadata": {}, + "source": [ + "### Add a new column" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "c7017f3d-869d-42e3-bbd8-b3fbc408c2d0", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutestitle
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187Marlins vs Cubs
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189Marlins vs Cubs
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165Braves vs Cubs
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222Braves vs Cubs
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164Phillies vs Cubs
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201Diamondbacks vs Cubs
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173Athletics vs Cubs
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176Athletics vs Cubs
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180Rockies vs Cubs
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157Cardinals vs Cubs
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218Cardinals vs Cubs
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160Cardinals vs Cubs
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178Pirates vs Cubs
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205Pirates vs Cubs
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197Giants vs Cubs
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198Reds vs Cubs
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188Reds vs Cubs
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188Reds vs Cubs
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194Reds vs Cubs
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175Padres vs Cubs
\n", + "
[2431 rows x 6 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes title \n", + "0 187 Marlins vs Cubs \n", + "1 189 Marlins vs Cubs \n", + "2 165 Braves vs Cubs \n", + "3 222 Braves vs Cubs \n", + "4 164 Phillies vs Cubs \n", + "5 201 Diamondbacks vs Cubs \n", + "6 173 Athletics vs Cubs \n", + "7 176 Athletics vs Cubs \n", + "8 180 Rockies vs Cubs \n", + "9 157 Cardinals vs Cubs \n", + "10 218 Cardinals vs Cubs \n", + "11 160 Cardinals vs Cubs \n", + "12 178 Pirates vs Cubs \n", + "13 205 Pirates vs Cubs \n", + "14 197 Giants vs Cubs \n", + "15 198 Reds vs Cubs \n", + "16 188 Reds vs Cubs \n", + "17 188 Reds vs Cubs \n", + "18 194 Reds vs Cubs \n", + "19 175 Padres vs Cubs \n", + "20 257 Nationals vs Cubs \n", + "21 178 Brewers vs Cubs \n", + "22 171 Brewers vs Cubs \n", + "23 248 Brewers vs Cubs \n", + "24 174 Astros vs Cubs \n", + "...\n", + "\n", + "[2431 rows x 6 columns]" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.assign(title=df['homeTeamName'] + \" vs \" + df['awayTeamName'])\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a4f9c18c-865e-4057-9e9a-85564f940c82", + "metadata": {}, + "source": [ + "### Manipulate the column" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "8bbe000a-36f0-4b6f-b403-b9ec28dd608b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutesheadline
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187Marlins vs Cubs
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189Marlins vs Cubs
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165Braves vs Cubs
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222Braves vs Cubs
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164Phillies vs Cubs
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201Diamondbacks vs Cubs
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173Athletics vs Cubs
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176Athletics vs Cubs
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180Rockies vs Cubs
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157Cardinals vs Cubs
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218Cardinals vs Cubs
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160Cardinals vs Cubs
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178Pirates vs Cubs
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205Pirates vs Cubs
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197Giants vs Cubs
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198Reds vs Cubs
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188Reds vs Cubs
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188Reds vs Cubs
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194Reds vs Cubs
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175Padres vs Cubs
\n", + "
[2431 rows x 6 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes headline \n", + "0 187 Marlins vs Cubs \n", + "1 189 Marlins vs Cubs \n", + "2 165 Braves vs Cubs \n", + "3 222 Braves vs Cubs \n", + "4 164 Phillies vs Cubs \n", + "5 201 Diamondbacks vs Cubs \n", + "6 173 Athletics vs Cubs \n", + "7 176 Athletics vs Cubs \n", + "8 180 Rockies vs Cubs \n", + "9 157 Cardinals vs Cubs \n", + "10 218 Cardinals vs Cubs \n", + "11 160 Cardinals vs Cubs \n", + "12 178 Pirates vs Cubs \n", + "13 205 Pirates vs Cubs \n", + "14 197 Giants vs Cubs \n", + "15 198 Reds vs Cubs \n", + "16 188 Reds vs Cubs \n", + "17 188 Reds vs Cubs \n", + "18 194 Reds vs Cubs \n", + "19 175 Padres vs Cubs \n", + "20 257 Nationals vs Cubs \n", + "21 178 Brewers vs Cubs \n", + "22 171 Brewers vs Cubs \n", + "23 248 Brewers vs Cubs \n", + "24 174 Astros vs Cubs \n", + "...\n", + "\n", + "[2431 rows x 6 columns]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.rename(columns={\"title\": \"headline\"})\n", + "df" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "87eee643-28ac-4f4b-ac61-1f3de9c08a9d", + "metadata": {}, + "outputs": [], + "source": [ + "df = df.drop(columns=\"headline\")" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "fad6d3da-1f40-4c5f-94ec-0bdfe21ca5b6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "218b3d0f-ee0d-4a8a-8c7d-cf24825f3d3e", + "metadata": {}, + "source": [ + "### Drop Nan values" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "67a7c35f-80cf-4482-80f9-7f01c7743807", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.dropna()\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "80e9a2e2-c4c9-4c17-bbd0-06882d7657fe", + "metadata": {}, + "source": [ + "### Join two DFs" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "3f09ff32-ef43-4fab-a86b-8868afc34363", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlins
11f32b347-cbcb-4c31-a145-0e685306d168Marlins
20c2292d1-7398-48be-bf8e-b41dad5e1a43Braves
38fbec734-a15a-42ab-8d51-60790de7750bBraves
489e514d5-fbf5-4b9d-bdac-6ca45bfd18ddPhillies
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52Diamondbacks
676ea8662-c7e6-4c38-8f2a-efe373e428ceAthletics
766fad23d-6e89-4f99-be29-d49b6e94f95dAthletics
8d977367c-cf0c-4687-95a0-eb4542efcb01Rockies
9a87070ff-1084-43ca-a7ba-69278f93ecbaCardinals
10ea6b350d-3c1d-4737-878d-4465f66999f6Cardinals
1146463c50-0f5c-4dca-a661-dd194464e791Cardinals
1259134e6d-9d13-49aa-978e-c3c2300eb90fPirates
13387630a3-a894-4327-baa1-b24ec1a654d9Pirates
145d084e13-94fd-4995-b95a-4801ea3ed556Giants
1534444c94-03ec-4d12-96af-68b8f399a22fReds
169580bffe-22e1-4975-978b-1b13e7505193Reds
17645e6a08-afd6-4677-a5c9-01ef446b0cf3Reds
1808981bd8-d1d7-48e1-8668-9098b8f7fe90Reds
19303703bb-b55f-476d-8faf-bf582169fb1dPadres
\n", + "
[2431 rows x 2 columns in total]" + ], + "text/plain": [ + " gameId homeTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins\n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 Braves\n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b Braves\n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd Phillies\n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 Diamondbacks\n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce Athletics\n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d Athletics\n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 Rockies\n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba Cardinals\n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 Cardinals\n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 Cardinals\n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f Pirates\n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 Pirates\n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 Giants\n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f Reds\n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 Reds\n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 Reds\n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 Reds\n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d Padres\n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 Nationals\n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c Brewers\n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 Brewers\n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 Brewers\n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 Astros\n", + "...\n", + "\n", + "[2431 rows x 2 columns]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1 = df[\"gameId\", \"homeTeamName\"]\n", + "df1" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "5331d2c8-7912-4d96-8da1-f64b57374df3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdawayTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfCubs
11f32b347-cbcb-4c31-a145-0e685306d168Cubs
\n", + "
[2 rows x 2 columns in total]" + ], + "text/plain": [ + " gameId awayTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Cubs\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Cubs\n", + "\n", + "[2 rows x 2 columns]" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df2 = df[\"gameId\", \"awayTeamName\"].head(2)\n", + "df2" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "a574ad3e-a219-454c-8bb5-c5ed6627f2c6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlinsCubs
11f32b347-cbcb-4c31-a145-0e685306d168MarlinsCubs
\n", + "
[2 rows x 3 columns in total]" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins Cubs\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins Cubs\n", + "\n", + "[2 rows x 3 columns]" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1.merge(df2, on=\"gameId\", how=\"inner\")" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "288e7a95-a077-46c4-8fe6-802474c01f8b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlinsCubs
11f32b347-cbcb-4c31-a145-0e685306d168MarlinsCubs
20c2292d1-7398-48be-bf8e-b41dad5e1a43Braves<NA>
38fbec734-a15a-42ab-8d51-60790de7750bBraves<NA>
489e514d5-fbf5-4b9d-bdac-6ca45bfd18ddPhillies<NA>
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52Diamondbacks<NA>
676ea8662-c7e6-4c38-8f2a-efe373e428ceAthletics<NA>
766fad23d-6e89-4f99-be29-d49b6e94f95dAthletics<NA>
8d977367c-cf0c-4687-95a0-eb4542efcb01Rockies<NA>
9a87070ff-1084-43ca-a7ba-69278f93ecbaCardinals<NA>
10ea6b350d-3c1d-4737-878d-4465f66999f6Cardinals<NA>
1146463c50-0f5c-4dca-a661-dd194464e791Cardinals<NA>
1259134e6d-9d13-49aa-978e-c3c2300eb90fPirates<NA>
13387630a3-a894-4327-baa1-b24ec1a654d9Pirates<NA>
145d084e13-94fd-4995-b95a-4801ea3ed556Giants<NA>
1534444c94-03ec-4d12-96af-68b8f399a22fReds<NA>
169580bffe-22e1-4975-978b-1b13e7505193Reds<NA>
17645e6a08-afd6-4677-a5c9-01ef446b0cf3Reds<NA>
1808981bd8-d1d7-48e1-8668-9098b8f7fe90Reds<NA>
19303703bb-b55f-476d-8faf-bf582169fb1dPadres<NA>
\n", + "
[2431 rows x 3 columns in total]" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins Cubs\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins Cubs\n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 Braves \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b Braves \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd Phillies \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 Diamondbacks \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce Athletics \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d Athletics \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 Rockies \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba Cardinals \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 Cardinals \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 Cardinals \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f Pirates \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 Pirates \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 Giants \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f Reds \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 Reds \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 Reds \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 Reds \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d Padres \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 Nationals \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c Brewers \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 Brewers \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 Brewers \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 Astros \n", + "...\n", + "\n", + "[2431 rows x 3 columns]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1.merge(df2, on=\"gameId\", how=\"outer\")" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "7ee87a01-2ff5-4021-855d-44b71cf2a225", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlinsCubs
11f32b347-cbcb-4c31-a145-0e685306d168MarlinsCubs
20c2292d1-7398-48be-bf8e-b41dad5e1a43Braves<NA>
38fbec734-a15a-42ab-8d51-60790de7750bBraves<NA>
489e514d5-fbf5-4b9d-bdac-6ca45bfd18ddPhillies<NA>
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52Diamondbacks<NA>
676ea8662-c7e6-4c38-8f2a-efe373e428ceAthletics<NA>
766fad23d-6e89-4f99-be29-d49b6e94f95dAthletics<NA>
8d977367c-cf0c-4687-95a0-eb4542efcb01Rockies<NA>
9a87070ff-1084-43ca-a7ba-69278f93ecbaCardinals<NA>
10ea6b350d-3c1d-4737-878d-4465f66999f6Cardinals<NA>
1146463c50-0f5c-4dca-a661-dd194464e791Cardinals<NA>
1259134e6d-9d13-49aa-978e-c3c2300eb90fPirates<NA>
13387630a3-a894-4327-baa1-b24ec1a654d9Pirates<NA>
145d084e13-94fd-4995-b95a-4801ea3ed556Giants<NA>
1534444c94-03ec-4d12-96af-68b8f399a22fReds<NA>
169580bffe-22e1-4975-978b-1b13e7505193Reds<NA>
17645e6a08-afd6-4677-a5c9-01ef446b0cf3Reds<NA>
1808981bd8-d1d7-48e1-8668-9098b8f7fe90Reds<NA>
19303703bb-b55f-476d-8faf-bf582169fb1dPadres<NA>
\n", + "
[2431 rows x 3 columns in total]" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins Cubs\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins Cubs\n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 Braves \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b Braves \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd Phillies \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 Diamondbacks \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce Athletics \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d Athletics \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 Rockies \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba Cardinals \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 Cardinals \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 Cardinals \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f Pirates \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 Pirates \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 Giants \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f Reds \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 Reds \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 Reds \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 Reds \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d Padres \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 Nationals \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c Brewers \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 Brewers \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 Brewers \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 Astros \n", + "...\n", + "\n", + "[2431 rows x 3 columns]" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1.merge(df2, on=\"gameId\", how=\"left\")" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "330ed69c-f122-4af9-bf5e-96e309d3fa0c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlinsCubs
11f32b347-cbcb-4c31-a145-0e685306d168MarlinsCubs
\n", + "
[2 rows x 3 columns in total]" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins Cubs\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins Cubs\n", + "\n", + "[2 rows x 3 columns]" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1.merge(df2, on=\"gameId\", how=\"right\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "ea42b463-f37b-4ea4-805f-a16a84b656e4", + "metadata": {}, + "source": [ + "### Concat two DFs" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "5181231e-8a2a-4ac5-a379-6aa5ad4fee89", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[4862 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[4862 rows x 5 columns]" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "bpd.concat([df, df])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "dcddbc12-3f17-43f5-a829-84d8093c7760", + "metadata": {}, + "source": [ + "### Access column through property" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "ad1f86f1-890b-462b-b408-b94c073371ff", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0 Marlins\n", + "1 Marlins\n", + "2 Braves\n", + "3 Braves\n", + "4 Phillies\n", + " ... \n", + "2426 Dodgers\n", + "2427 Dodgers\n", + "2428 Mets\n", + "2429 Mets\n", + "2430 American League\n", + "Name: homeTeamName, Length: 2431, dtype: string" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.homeTeamName" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "4ae2e978-ac96-4c6a-8be8-9df9424c3544", + "metadata": {}, + "source": [ + "### Retrieve SQL" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "2ae4924b-37a9-4327-81d4-4f2afb66cb37", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "WITH t0 AS (\n", + " SELECT t17.*, t17.`bigframes_ordering_id` AS `bigframes_index_0`\n", + " FROM (\n", + " SELECT * FROM `7150072e4dc649849bc421d3a492eee1`\n", + " ) t17\n", + "),\n", + "t1 AS (\n", + " SELECT coalesce(`bigframes_ordering_id`, `bigframes_ordering_id`) AS `index_2`,\n", + " CONCAT(CONCAT(t0.`homeTeamName`, ' vs '), t0.`awayTeamName`) AS `col_3`,\n", + " t0.`bigframes_ordering_id` AS `bigframes_ordering_id`\n", + " FROM t0\n", + "),\n", + "t2 AS (\n", + " SELECT `bigframes_ordering_id` AS `bigframes_index_0`, t0.`gameId`, t0.`year`,\n", + " t0.`homeTeamName`, t0.`awayTeamName`, t0.`duration_minutes`,\n", + " t0.`bigframes_ordering_id`\n", + " FROM t0\n", + "),\n", + "t3 AS (\n", + " SELECT t2.`bigframes_index_0`, t2.`gameId`, t2.`year`, t2.`homeTeamName`,\n", + " t2.`awayTeamName`, t2.`duration_minutes`,\n", + " t2.`bigframes_ordering_id` AS `bigframes_ordering_id_x`,\n", + " t1.`index_2`, t1.`col_3`,\n", + " t1.`bigframes_ordering_id` AS `bigframes_ordering_id_y`\n", + " FROM t2\n", + " LEFT OUTER JOIN t1\n", + " ON IFNULL(CAST(t2.`bigframes_index_0` AS STRING), '$NULL_SENTINEL$') = IFNULL(CAST(t1.`index_2` AS STRING), '$NULL_SENTINEL$')\n", + "),\n", + "t4 AS (\n", + " SELECT coalesce(t3.`bigframes_index_0`, t3.`index_2`) AS `index_5`,\n", + " t3.`gameId`, t3.`awayTeamName`,\n", + " ((((t3.`gameId` IS NOT NULL) AND (t3.`year` IS NOT NULL)) AND (t3.`homeTeamName` IS NOT NULL)) AND (t3.`awayTeamName` IS NOT NULL)) AND (t3.`duration_minutes` IS NOT NULL) AS `bigframes_predicate`,\n", + " (row_number() OVER (PARTITION BY ((((t3.`gameId` IS NOT NULL) AND (t3.`year` IS NOT NULL)) AND (t3.`homeTeamName` IS NOT NULL)) AND (t3.`awayTeamName` IS NOT NULL)) AND (t3.`duration_minutes` IS NOT NULL) ORDER BY CONCAT(IFNULL(lpad(CAST(t3.`bigframes_ordering_id_x` AS STRING), 19, '0'), ':::::::::::::::::::'), IFNULL(lpad(CAST(t3.`bigframes_ordering_id_y` AS STRING), 19, '0'), ':::::::::::::::::::')) IS NULL ASC, CONCAT(IFNULL(lpad(CAST(t3.`bigframes_ordering_id_x` AS STRING), 19, '0'), ':::::::::::::::::::'), IFNULL(lpad(CAST(t3.`bigframes_ordering_id_y` AS STRING), 19, '0'), ':::::::::::::::::::')) ASC) - 1) AS `bigframes_ordering_id`\n", + " FROM t3\n", + "),\n", + "t5 AS (\n", + " SELECT t4.*\n", + " FROM t4\n", + " WHERE t4.`bigframes_predicate`\n", + "),\n", + "t6 AS (\n", + " SELECT t5.`index_5`, t5.`gameId`, t5.`awayTeamName`,\n", + " t5.`bigframes_ordering_id`\n", + " FROM t5\n", + "),\n", + "t7 AS (\n", + " SELECT t6.`index_5`, t6.`gameId`, t6.`awayTeamName`,\n", + " (t6.`bigframes_ordering_id` = t6.`bigframes_ordering_id`) AND (t6.`bigframes_ordering_id` < 2) AS `bigframes_predicate`,\n", + " t6.`bigframes_ordering_id`\n", + " FROM t6\n", + "),\n", + "t8 AS (\n", + " SELECT coalesce(t3.`bigframes_index_0`, t3.`index_2`) AS `index_5`,\n", + " t3.`gameId`, t3.`homeTeamName`,\n", + " ((((t3.`gameId` IS NOT NULL) AND (t3.`year` IS NOT NULL)) AND (t3.`homeTeamName` IS NOT NULL)) AND (t3.`awayTeamName` IS NOT NULL)) AND (t3.`duration_minutes` IS NOT NULL) AS `bigframes_predicate`,\n", + " CONCAT(IFNULL(lpad(CAST(t3.`bigframes_ordering_id_x` AS STRING), 19, '0'), ':::::::::::::::::::'), IFNULL(lpad(CAST(t3.`bigframes_ordering_id_y` AS STRING), 19, '0'), ':::::::::::::::::::')) AS `bigframes_ordering_id`\n", + " FROM t3\n", + "),\n", + "t9 AS (\n", + " SELECT t7.*\n", + " FROM t7\n", + " WHERE t7.`bigframes_predicate`\n", + "),\n", + "t10 AS (\n", + " SELECT t8.*\n", + " FROM t8\n", + " WHERE t8.`bigframes_predicate`\n", + "),\n", + "t11 AS (\n", + " SELECT t9.`index_5`, t9.`gameId`, t9.`awayTeamName`,\n", + " t9.`bigframes_ordering_id`\n", + " FROM t9\n", + "),\n", + "t12 AS (\n", + " SELECT t10.`index_5`, t10.`gameId`, t10.`homeTeamName`,\n", + " t10.`bigframes_ordering_id`\n", + " FROM t10\n", + "),\n", + "t13 AS (\n", + " SELECT t12.`index_5` AS `index_5_x`, t12.`gameId` AS `gameId_x`,\n", + " t12.`homeTeamName`,\n", + " t12.`bigframes_ordering_id` AS `bigframes_ordering_id_x`,\n", + " t11.`index_5` AS `index_5_y`, t11.`gameId` AS `gameId_y`,\n", + " t11.`awayTeamName`,\n", + " t11.`bigframes_ordering_id` AS `bigframes_ordering_id_y`\n", + " FROM t12\n", + " INNER JOIN t11\n", + " ON IFNULL(t12.`gameId`, '$NULL_SENTINEL$') = IFNULL(t11.`gameId`, '$NULL_SENTINEL$')\n", + "),\n", + "t14 AS (\n", + " SELECT coalesce(t13.`gameId_x`, t13.`gameId_y`) AS `index_24`,\n", + " t13.`homeTeamName`, t13.`awayTeamName`,\n", + " (row_number() OVER (ORDER BY CONCAT(IFNULL(lpad(t13.`bigframes_ordering_id_x`, 38, '0'), '::::::::::::::::::::::::::::::::::::::'), IFNULL(lpad(CAST(t13.`bigframes_ordering_id_y` AS STRING), 19, '0'), ':::::::::::::::::::')) IS NULL ASC, CONCAT(IFNULL(lpad(t13.`bigframes_ordering_id_x`, 38, '0'), '::::::::::::::::::::::::::::::::::::::'), IFNULL(lpad(CAST(t13.`bigframes_ordering_id_y` AS STRING), 19, '0'), ':::::::::::::::::::')) ASC) - 1) AS `bigframes_ordering_id`\n", + " FROM t13\n", + "),\n", + "t15 AS (\n", + " SELECT t14.`bigframes_ordering_id` AS `col_25`, t14.`index_24`,\n", + " t14.`homeTeamName`, t14.`awayTeamName`\n", + " FROM t14\n", + ")\n", + "SELECT t16.`index_24` AS `gameId`, t16.`homeTeamName`, t16.`awayTeamName`\n", + "FROM (\n", + " SELECT t15.`index_24`, t15.`homeTeamName`, t15.`awayTeamName`\n", + " FROM t15\n", + ") t16\n" + ] + } + ], + "source": [ + "print(df1.merge(df2, on=\"gameId\", how=\"inner\").to_sql_query(always_include_index=False)[0])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "cec555be-0773-45e0-9008-86cd3e645285", + "metadata": {}, + "source": [ + "### Special Column Names" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "5adc7bbf-2c58-4924-964c-ed1b18dc9268", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearHOME TEAMawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year HOME TEAM awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.rename(columns={\"homeTeamName\": \"HOME TEAM\"})" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "ac3ceabe-4317-453c-9418-826de5094454", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeam!@#$%colawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year homeTeam!@#$%col awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.rename(columns={\"homeTeamName\": \"homeTeam!@#$%col\"})" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "e73704c9-7aa9-4f10-b414-3417c3ad9eb8", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearteamteamduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
56a83e76c-dc0d-4924-9d3d-a2e7e0ab5b522016DiamondbacksCubs201
676ea8662-c7e6-4c38-8f2a-efe373e428ce2016AthleticsCubs173
766fad23d-6e89-4f99-be29-d49b6e94f95d2016AthleticsCubs176
8d977367c-cf0c-4687-95a0-eb4542efcb012016RockiesCubs180
9a87070ff-1084-43ca-a7ba-69278f93ecba2016CardinalsCubs157
10ea6b350d-3c1d-4737-878d-4465f66999f62016CardinalsCubs218
1146463c50-0f5c-4dca-a661-dd194464e7912016CardinalsCubs160
1259134e6d-9d13-49aa-978e-c3c2300eb90f2016PiratesCubs178
13387630a3-a894-4327-baa1-b24ec1a654d92016PiratesCubs205
145d084e13-94fd-4995-b95a-4801ea3ed5562016GiantsCubs197
1534444c94-03ec-4d12-96af-68b8f399a22f2016RedsCubs198
169580bffe-22e1-4975-978b-1b13e75051932016RedsCubs188
17645e6a08-afd6-4677-a5c9-01ef446b0cf32016RedsCubs188
1808981bd8-d1d7-48e1-8668-9098b8f7fe902016RedsCubs194
19303703bb-b55f-476d-8faf-bf582169fb1d2016PadresCubs175
\n", + "
[2431 rows x 5 columns in total]" + ], + "text/plain": [ + " gameId year team team \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "5 6a83e76c-dc0d-4924-9d3d-a2e7e0ab5b52 2016 Diamondbacks Cubs \n", + "6 76ea8662-c7e6-4c38-8f2a-efe373e428ce 2016 Athletics Cubs \n", + "7 66fad23d-6e89-4f99-be29-d49b6e94f95d 2016 Athletics Cubs \n", + "8 d977367c-cf0c-4687-95a0-eb4542efcb01 2016 Rockies Cubs \n", + "9 a87070ff-1084-43ca-a7ba-69278f93ecba 2016 Cardinals Cubs \n", + "10 ea6b350d-3c1d-4737-878d-4465f66999f6 2016 Cardinals Cubs \n", + "11 46463c50-0f5c-4dca-a661-dd194464e791 2016 Cardinals Cubs \n", + "12 59134e6d-9d13-49aa-978e-c3c2300eb90f 2016 Pirates Cubs \n", + "13 387630a3-a894-4327-baa1-b24ec1a654d9 2016 Pirates Cubs \n", + "14 5d084e13-94fd-4995-b95a-4801ea3ed556 2016 Giants Cubs \n", + "15 34444c94-03ec-4d12-96af-68b8f399a22f 2016 Reds Cubs \n", + "16 9580bffe-22e1-4975-978b-1b13e7505193 2016 Reds Cubs \n", + "17 645e6a08-afd6-4677-a5c9-01ef446b0cf3 2016 Reds Cubs \n", + "18 08981bd8-d1d7-48e1-8668-9098b8f7fe90 2016 Reds Cubs \n", + "19 303703bb-b55f-476d-8faf-bf582169fb1d 2016 Padres Cubs \n", + "20 71ab82a4-6e07-430a-b695-1af3bc42ea61 2016 Nationals Cubs \n", + "21 d1a110c2-f6c8-4029-bcd8-2f8a01e1561c 2016 Brewers Cubs \n", + "22 6d111b57-fa0b-4f24-82df-ff33a26f0252 2016 Brewers Cubs \n", + "23 a97e9539-bbbd-4e03-bf15-f25ea2c1d923 2016 Brewers Cubs \n", + "24 dc0c9218-505c-4725-8c0c-40b72cca0956 2016 Astros Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 \n", + "5 201 \n", + "6 173 \n", + "7 176 \n", + "8 180 \n", + "9 157 \n", + "10 218 \n", + "11 160 \n", + "12 178 \n", + "13 205 \n", + "14 197 \n", + "15 198 \n", + "16 188 \n", + "17 188 \n", + "18 194 \n", + "19 175 \n", + "20 257 \n", + "21 178 \n", + "22 171 \n", + "23 248 \n", + "24 174 \n", + "...\n", + "\n", + "[2431 rows x 5 columns]" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df3 = df.rename(columns={\"homeTeamName\": \"team\", \"awayTeamName\": \"team\"})\n", + "df3" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "1a80f6f8-a172-4d7d-a2f5-e10871da7224", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
teamteam
0MarlinsCubs
1MarlinsCubs
2BravesCubs
3BravesCubs
4PhilliesCubs
5DiamondbacksCubs
6AthleticsCubs
7AthleticsCubs
8RockiesCubs
9CardinalsCubs
10CardinalsCubs
11CardinalsCubs
12PiratesCubs
13PiratesCubs
14GiantsCubs
15RedsCubs
16RedsCubs
17RedsCubs
18RedsCubs
19PadresCubs
\n", + "
[2431 rows x 2 columns in total]" + ], + "text/plain": [ + " team team\n", + "0 Marlins Cubs\n", + "1 Marlins Cubs\n", + "2 Braves Cubs\n", + "3 Braves Cubs\n", + "4 Phillies Cubs\n", + "5 Diamondbacks Cubs\n", + "6 Athletics Cubs\n", + "7 Athletics Cubs\n", + "8 Rockies Cubs\n", + "9 Cardinals Cubs\n", + "10 Cardinals Cubs\n", + "11 Cardinals Cubs\n", + "12 Pirates Cubs\n", + "13 Pirates Cubs\n", + "14 Giants Cubs\n", + "15 Reds Cubs\n", + "16 Reds Cubs\n", + "17 Reds Cubs\n", + "18 Reds Cubs\n", + "19 Padres Cubs\n", + "20 Nationals Cubs\n", + "21 Brewers Cubs\n", + "22 Brewers Cubs\n", + "23 Brewers Cubs\n", + "24 Astros Cubs\n", + "...\n", + "\n", + "[2431 rows x 2 columns]" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df3[\"team\"]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "8e1a4dc8-de23-466e-8ead-df4af82a8085", + "metadata": {}, + "source": [ + "### Binary Operation" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "2414a095-37df-4755-b86c-2031a6cb9d4a", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
yearduration_minutes
02016187
12016189
22016165
32016222
42016164
52016201
62016173
72016176
82016180
92016157
102016218
112016160
122016178
132016205
142016197
152016198
162016188
172016188
182016194
192016175
\n", + "
[2431 rows x 2 columns in total]" + ], + "text/plain": [ + " year duration_minutes\n", + "0 2016 187\n", + "1 2016 189\n", + "2 2016 165\n", + "3 2016 222\n", + "4 2016 164\n", + "5 2016 201\n", + "6 2016 173\n", + "7 2016 176\n", + "8 2016 180\n", + "9 2016 157\n", + "10 2016 218\n", + "11 2016 160\n", + "12 2016 178\n", + "13 2016 205\n", + "14 2016 197\n", + "15 2016 198\n", + "16 2016 188\n", + "17 2016 188\n", + "18 2016 194\n", + "19 2016 175\n", + "20 2016 257\n", + "21 2016 178\n", + "22 2016 171\n", + "23 2016 248\n", + "24 2016 174\n", + "...\n", + "\n", + "[2431 rows x 2 columns]" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df4 = df[[\"year\", \"duration_minutes\"]]\n", + "df4" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "7d437c7c-ae74-4f0d-a4f8-10a133f4b61e", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
yearduration_minutes
02017188
12017190
22017166
32017223
42017165
52017202
62017174
72017177
82017181
92017158
102017219
112017161
122017179
132017206
142017198
152017199
162017189
172017189
182017195
192017176
\n", + "
[2431 rows x 2 columns in total]" + ], + "text/plain": [ + " year duration_minutes\n", + "0 2017 188\n", + "1 2017 190\n", + "2 2017 166\n", + "3 2017 223\n", + "4 2017 165\n", + "5 2017 202\n", + "6 2017 174\n", + "7 2017 177\n", + "8 2017 181\n", + "9 2017 158\n", + "10 2017 219\n", + "11 2017 161\n", + "12 2017 179\n", + "13 2017 206\n", + "14 2017 198\n", + "15 2017 199\n", + "16 2017 189\n", + "17 2017 189\n", + "18 2017 195\n", + "19 2017 176\n", + "20 2017 258\n", + "21 2017 179\n", + "22 2017 172\n", + "23 2017 249\n", + "24 2017 175\n", + "...\n", + "\n", + "[2431 rows x 2 columns]" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df4 + 1" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/03 - Using ML - ML fundamentals.ipynb b/notebooks/03 - Using ML - ML fundamentals.ipynb new file mode 100644 index 00000000000..c181b67e647 --- /dev/null +++ b/notebooks/03 - Using ML - ML fundamentals.ipynb @@ -0,0 +1,2250 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using ML - ML fundamentals\n", + "\n", + "The `bigframes.ml` module implements Scikit-Learn's machine learning API in\n", + "BigQuery DataFrames. It exposes BigQuery's ML capabilities in a simple, popular\n", + "API that works seamlessly with the rest of the BigQuery DataFrames API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
penguin_id
0Adelie Penguin (Pygoscelis adeliae)Dream36.618.4184.03475.0FEMALE
1Adelie Penguin (Pygoscelis adeliae)Dream39.819.1184.04650.0MALE
2Adelie Penguin (Pygoscelis adeliae)Dream40.918.9184.03900.0MALE
3Chinstrap penguin (Pygoscelis antarctica)Dream46.517.9192.03500.0FEMALE
4Adelie Penguin (Pygoscelis adeliae)Dream37.316.8192.03000.0FEMALE
5Adelie Penguin (Pygoscelis adeliae)Dream43.218.5192.04100.0MALE
6Chinstrap penguin (Pygoscelis antarctica)Dream46.916.6192.02700.0FEMALE
7Chinstrap penguin (Pygoscelis antarctica)Dream50.518.4200.03400.0FEMALE
8Chinstrap penguin (Pygoscelis antarctica)Dream49.519.0200.03800.0MALE
9Adelie Penguin (Pygoscelis adeliae)Dream40.220.1200.03975.0MALE
10Adelie Penguin (Pygoscelis adeliae)Dream40.818.9208.04300.0MALE
11Adelie Penguin (Pygoscelis adeliae)Dream39.018.7185.03650.0MALE
12Adelie Penguin (Pygoscelis adeliae)Dream37.016.9185.03000.0FEMALE
13Chinstrap penguin (Pygoscelis antarctica)Dream47.017.3185.03700.0FEMALE
14Adelie Penguin (Pygoscelis adeliae)Dream34.017.1185.03400.0FEMALE
15Adelie Penguin (Pygoscelis adeliae)Dream37.016.5185.03400.0FEMALE
16Chinstrap penguin (Pygoscelis antarctica)Dream45.717.3193.03600.0FEMALE
17Chinstrap penguin (Pygoscelis antarctica)Dream50.619.4193.03800.0MALE
18Adelie Penguin (Pygoscelis adeliae)Dream39.717.9193.04250.0MALE
19Adelie Penguin (Pygoscelis adeliae)Dream37.818.1193.03750.0MALE
\n", + "
[334 rows x 7 columns in total]" + ], + "text/plain": [ + " species island \\\n", + "penguin_id \n", + "0 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "2 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "4 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "6 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "7 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "8 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "9 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "12 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "13 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "14 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "15 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "16 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "17 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "21 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "22 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "\n", + " culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", + "penguin_id \n", + "0 36.6 18.4 184.0 3475.0 \n", + "1 39.8 19.1 184.0 4650.0 \n", + "2 40.9 18.9 184.0 3900.0 \n", + "3 46.5 17.9 192.0 3500.0 \n", + "4 37.3 16.8 192.0 3000.0 \n", + "5 43.2 18.5 192.0 4100.0 \n", + "6 46.9 16.6 192.0 2700.0 \n", + "7 50.5 18.4 200.0 3400.0 \n", + "8 49.5 19.0 200.0 3800.0 \n", + "9 40.2 20.1 200.0 3975.0 \n", + "10 40.8 18.9 208.0 4300.0 \n", + "11 39.0 18.7 185.0 3650.0 \n", + "12 37.0 16.9 185.0 3000.0 \n", + "13 47.0 17.3 185.0 3700.0 \n", + "14 34.0 17.1 185.0 3400.0 \n", + "15 37.0 16.5 185.0 3400.0 \n", + "16 45.7 17.3 193.0 3600.0 \n", + "17 50.6 19.4 193.0 3800.0 \n", + "18 39.7 17.9 193.0 4250.0 \n", + "19 37.8 18.1 193.0 3750.0 \n", + "20 46.6 17.8 193.0 3800.0 \n", + "21 51.3 19.2 193.0 3650.0 \n", + "22 40.2 17.1 193.0 3400.0 \n", + "23 36.8 18.5 193.0 3500.0 \n", + "24 49.6 18.2 193.0 3775.0 \n", + "\n", + " sex \n", + "penguin_id \n", + "0 FEMALE \n", + "1 MALE \n", + "2 MALE \n", + "3 FEMALE \n", + "4 FEMALE \n", + "5 MALE \n", + "6 FEMALE \n", + "7 FEMALE \n", + "8 MALE \n", + "9 MALE \n", + "10 MALE \n", + "11 MALE \n", + "12 FEMALE \n", + "13 FEMALE \n", + "14 FEMALE \n", + "15 FEMALE \n", + "16 FEMALE \n", + "17 MALE \n", + "18 MALE \n", + "19 MALE \n", + "20 FEMALE \n", + "21 MALE \n", + "22 FEMALE \n", + "23 FEMALE \n", + "24 MALE \n", + "...\n", + "\n", + "[334 rows x 7 columns]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Lets load some test data to use in this tutorial\n", + "import bigframes.pandas\n", + "\n", + "df = bigframes.pandas.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")\n", + "df = df.dropna()\n", + "\n", + "# Temporary workaround: lets name our index so it isn't lost BigQuery DataFrame\n", + "# currently drops unnamed indexes when round-tripping through pandas, which\n", + "# some ML APIs do to route around missing functionality\n", + "df.index.name = \"penguin_id\"\n", + "\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Data split\n", + "\n", + "Part of preparing data for a machine learning task is splitting it into subsets for training and testing, to ensure that the solution is not overfitting. Most commonly this is done with `bigframes.ml.model_selection.train_test_split` like so:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "train_X shape: (267, 6)\n", + "test_X shape: (67, 6)\n", + "train_y shape: (267, 1)\n", + "test_y shape: (67, 1)\n" + ] + } + ], + "source": [ + "# In this example, we're doing supervised learning, where we will learn to predict\n", + "# output variable `y` from input features `X`\n", + "X = df[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex', 'species']]\n", + "y = df[['body_mass_g']] \n", + "\n", + "from bigframes.ml.model_selection import train_test_split\n", + "\n", + "# This will split X and y into test and training sets, with 20% of the rows in the test set,\n", + "# and the rest in the training set\n", + "train_X, test_X, train_y, test_y = train_test_split(\n", + " X, y, test_size=0.2)\n", + "\n", + "# Show the shape of the data after the split\n", + "print(f\"\"\"train_X shape: {train_X.shape}\n", + "test_X shape: {test_X.shape}\n", + "train_y shape: {train_y.shape}\n", + "test_y shape: {test_y.shape}\"\"\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
islandculmen_length_mmculmen_depth_mmflipper_length_mmsexspecies
penguin_id
125Biscoe39.718.9184.0MALEAdelie Penguin (Pygoscelis adeliae)
260Biscoe50.415.7222.0MALEGentoo penguin (Pygoscelis papua)
93Dream41.118.1205.0MALEAdelie Penguin (Pygoscelis adeliae)
255Biscoe44.514.7214.0FEMALEGentoo penguin (Pygoscelis papua)
71Dream32.115.5188.0FEMALEAdelie Penguin (Pygoscelis adeliae)
\n", + "
[5 rows x 6 columns in total]" + ], + "text/plain": [ + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "penguin_id \n", + "125 Biscoe 39.7 18.9 184.0 \n", + "260 Biscoe 50.4 15.7 222.0 \n", + "93 Dream 41.1 18.1 205.0 \n", + "255 Biscoe 44.5 14.7 214.0 \n", + "71 Dream 32.1 15.5 188.0 \n", + "\n", + " sex species \n", + "penguin_id \n", + "125 MALE Adelie Penguin (Pygoscelis adeliae) \n", + "260 MALE Gentoo penguin (Pygoscelis papua) \n", + "93 MALE Adelie Penguin (Pygoscelis adeliae) \n", + "255 FEMALE Gentoo penguin (Pygoscelis papua) \n", + "71 FEMALE Adelie Penguin (Pygoscelis adeliae) \n", + "\n", + "[5 rows x 6 columns]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# If we look at the data, we can see that random rows were selected for\n", + "# each side of the split\n", + "test_X.head(5)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
body_mass_g
penguin_id
1253550.0
2605750.0
934300.0
2554850.0
713050.0
\n", + "
[5 rows x 1 columns in total]" + ], + "text/plain": [ + " body_mass_g\n", + "penguin_id \n", + "125 3550.0\n", + "260 5750.0\n", + "93 4300.0\n", + "255 4850.0\n", + "71 3050.0\n", + "\n", + "[5 rows x 1 columns]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Note that this matches the rows in test_X\n", + "test_y.head(5)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Estimators\n", + "\n", + "Following Scikit-Learn, all learning components are \"estimators\"; objects that can learn from training data and then apply themselves to new data. Estimators share the following patterns:\n", + "\n", + "- a constructor that takes a list of parameters\n", + "- a standard string representation that shows the class name and all non-default parameters, e.g. `LinearRegression(fit_intercept=False)`\n", + "- a `.fit(..)` method to fit the estimator to training data\n", + "\n", + "There estimators can be further broken down into two main subtypes:\n", + "\n", + "### Transformers\n", + "\n", + "Transformers are estimators that are used to prepare data for consumption by other estimators ('preprocessing'). In addition to `.fit(...)`, the transformer implements a `.transform(...)` method, which will apply a transformation based on what was computed during `.fit(..)`. With this pattern dynamic preprocessing steps can be applied to both training and test/production data consistently.\n", + "\n", + "An example of a transformer is `bigframes.ml.preprocessing.StandardScaler`, which rescales a dataset to have a mean of zero and a standard deviation of one:" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
scaled_culmen_length_mmscaled_culmen_depth_mmscaled_flipper_length_mm
penguin_id
23-1.3102310.687253-0.568262
832.5759550.329548-1.432429
1950.816172-1.4078760.151877
2641.109469-0.4369622.096254
2370.816172-1.3056740.872017
2610.614531-0.9479691.520142
2530.321234-1.6633790.944031
1640.101261-0.7435661.160073
293-0.686975-0.079257-1.792499
1760.871166-1.5611780.655975
451.4760911.8114690.655975
1851.347773-1.5100771.232087
1411.036145-1.1012711.088059
1471.109469-0.6413651.66417
320-0.9252790.329548-1.432429
285-0.192036-2.0721851.016045
3050.0279370.431750.655975
371.4210981.607066-0.496248
98-0.5219950.176246-0.784304
10-0.5769890.8916560.511947
\n", + "
[267 rows x 3 columns in total]" + ], + "text/plain": [ + " scaled_culmen_length_mm scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "23 -1.310231 0.687253 \n", + "83 2.575955 0.329548 \n", + "195 0.816172 -1.407876 \n", + "264 1.109469 -0.436962 \n", + "237 0.816172 -1.305674 \n", + "261 0.614531 -0.947969 \n", + "253 0.321234 -1.663379 \n", + "164 0.101261 -0.743566 \n", + "293 -0.686975 -0.079257 \n", + "176 0.871166 -1.561178 \n", + "45 1.476091 1.811469 \n", + "185 1.347773 -1.510077 \n", + "141 1.036145 -1.101271 \n", + "147 1.109469 -0.641365 \n", + "320 -0.925279 0.329548 \n", + "285 -0.192036 -2.072185 \n", + "305 0.027937 0.43175 \n", + "37 1.421098 1.607066 \n", + "98 -0.521995 0.176246 \n", + "10 -0.576989 0.891656 \n", + "307 -1.420218 0.022944 \n", + "337 -0.393678 1.14716 \n", + "26 -0.448671 0.687253 \n", + "97 -1.45688 0.380649 \n", + "27 1.201125 1.249361 \n", + "\n", + " scaled_flipper_length_mm \n", + "penguin_id \n", + "23 -0.568262 \n", + "83 -1.432429 \n", + "195 0.151877 \n", + "264 2.096254 \n", + "237 0.872017 \n", + "261 1.520142 \n", + "253 0.944031 \n", + "164 1.160073 \n", + "293 -1.792499 \n", + "176 0.655975 \n", + "45 0.655975 \n", + "185 1.232087 \n", + "141 1.088059 \n", + "147 1.66417 \n", + "320 -1.432429 \n", + "285 1.016045 \n", + "305 0.655975 \n", + "37 -0.496248 \n", + "98 -0.784304 \n", + "10 0.511947 \n", + "307 -1.000346 \n", + "337 -0.208192 \n", + "26 0.00785 \n", + "97 -0.784304 \n", + "27 0.00785 \n", + "...\n", + "\n", + "[267 rows x 3 columns]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.preprocessing import StandardScaler\n", + "\n", + "# StandardScaler will only work on numeric columns\n", + "numeric_columns = [\"culmen_length_mm\", \"culmen_depth_mm\", \"flipper_length_mm\"]\n", + "\n", + "scaler = StandardScaler()\n", + "scaler.fit(train_X[numeric_columns])\n", + "\n", + "# Now, standardscaler should transform the numbers to have mean of zero\n", + "# and standard deviation of one:\n", + "scaler.transform(train_X[numeric_columns])" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
scaled_culmen_length_mmscaled_culmen_depth_mmscaled_flipper_length_mm
penguin_id
125-0.778630.891656-1.216387
2601.182794-0.7435661.520142
93-0.5219950.4828510.295905
2550.101261-1.2545740.944031
71-2.171791-0.845768-0.928332
327-0.5219950.227347-1.360415
338-0.9802722.066973-0.71229
681.292780.8405560.151877
171.2194561.14716-0.568262
116-0.3020232.066973-0.71229
106-0.943611.453764-0.784304
2271.512753-0.0281561.952226
314-0.2103670.687253-0.424234
246-1.1269210.585052-1.936527
1381.201125-0.999071.088059
126-1.383556-0.028156-1.216387
1210.761179-0.385861-0.136178
231-0.3203541.198261-0.280206
401.1461320.8405560.079863
600.5045440.380649-0.424234
\n", + "
[67 rows x 3 columns in total]" + ], + "text/plain": [ + " scaled_culmen_length_mm scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "125 -0.77863 0.891656 \n", + "260 1.182794 -0.743566 \n", + "93 -0.521995 0.482851 \n", + "255 0.101261 -1.254574 \n", + "71 -2.171791 -0.845768 \n", + "327 -0.521995 0.227347 \n", + "338 -0.980272 2.066973 \n", + "68 1.29278 0.840556 \n", + "17 1.219456 1.14716 \n", + "116 -0.302023 2.066973 \n", + "106 -0.94361 1.453764 \n", + "227 1.512753 -0.028156 \n", + "314 -0.210367 0.687253 \n", + "246 -1.126921 0.585052 \n", + "138 1.201125 -0.99907 \n", + "126 -1.383556 -0.028156 \n", + "121 0.761179 -0.385861 \n", + "231 -0.320354 1.198261 \n", + "40 1.146132 0.840556 \n", + "60 0.504544 0.380649 \n", + "331 -1.548535 0.176246 \n", + "218 0.192916 -0.896869 \n", + "154 -1.10859 0.738354 \n", + "130 -0.357016 1.198261 \n", + "89 1.604408 1.351563 \n", + "\n", + " scaled_flipper_length_mm \n", + "penguin_id \n", + "125 -1.216387 \n", + "260 1.520142 \n", + "93 0.295905 \n", + "255 0.944031 \n", + "71 -0.928332 \n", + "327 -1.360415 \n", + "338 -0.71229 \n", + "68 0.151877 \n", + "17 -0.568262 \n", + "116 -0.71229 \n", + "106 -0.784304 \n", + "227 1.952226 \n", + "314 -0.424234 \n", + "246 -1.936527 \n", + "138 1.088059 \n", + "126 -1.216387 \n", + "121 -0.136178 \n", + "231 -0.280206 \n", + "40 0.079863 \n", + "60 -0.424234 \n", + "331 -0.784304 \n", + "218 1.376114 \n", + "154 -0.568262 \n", + "130 -0.064164 \n", + "89 -0.280206 \n", + "...\n", + "\n", + "[67 rows x 3 columns]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# We can then repeat this transformation on new data\n", + "scaler.transform(test_X[numeric_columns])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Composing transformers\n", + "\n", + "To process data where different columns need different preprocessors, `bigframes.composition.ColumnTransformer` can be employed:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
onehotencoded_islandscaled_culmen_length_mmscaled_culmen_depth_mmscaled_flipper_length_mmonehotencoded_sexonehotencoded_species
penguin_id
23[{'index': 2, 'value': 1.0}]-1.3102310.687253-0.568262[{'index': 1, 'value': 1.0}][{'index': 1, 'value': 1.0}]
83[{'index': 2, 'value': 1.0}]2.5759550.329548-1.432429[{'index': 1, 'value': 1.0}][{'index': 2, 'value': 1.0}]
195[{'index': 1, 'value': 1.0}]0.816172-1.4078760.151877[{'index': 1, 'value': 1.0}][{'index': 3, 'value': 1.0}]
264[{'index': 1, 'value': 1.0}]1.109469-0.4369622.096254[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
237[{'index': 1, 'value': 1.0}]0.816172-1.3056740.872017[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
261[{'index': 1, 'value': 1.0}]0.614531-0.9479691.520142[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
253[{'index': 1, 'value': 1.0}]0.321234-1.6633790.944031[{'index': 1, 'value': 1.0}][{'index': 3, 'value': 1.0}]
164[{'index': 1, 'value': 1.0}]0.101261-0.7435661.160073[{'index': 0, 'value': 1.0}][{'index': 3, 'value': 1.0}]
293[{'index': 3, 'value': 1.0}]-0.686975-0.079257-1.792499[{'index': 1, 'value': 1.0}][{'index': 1, 'value': 1.0}]
176[{'index': 1, 'value': 1.0}]0.871166-1.5611780.655975[{'index': 1, 'value': 1.0}][{'index': 3, 'value': 1.0}]
45[{'index': 2, 'value': 1.0}]1.4760911.8114690.655975[{'index': 2, 'value': 1.0}][{'index': 2, 'value': 1.0}]
185[{'index': 1, 'value': 1.0}]1.347773-1.5100771.232087[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
141[{'index': 1, 'value': 1.0}]1.036145-1.1012711.088059[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
147[{'index': 1, 'value': 1.0}]1.109469-0.6413651.66417[{'index': 2, 'value': 1.0}][{'index': 3, 'value': 1.0}]
320[{'index': 3, 'value': 1.0}]-0.9252790.329548-1.432429[{'index': 1, 'value': 1.0}][{'index': 1, 'value': 1.0}]
285[{'index': 1, 'value': 1.0}]-0.192036-2.0721851.016045[{'index': 1, 'value': 1.0}][{'index': 3, 'value': 1.0}]
305[{'index': 3, 'value': 1.0}]0.0279370.431750.655975[{'index': 2, 'value': 1.0}][{'index': 1, 'value': 1.0}]
37[{'index': 2, 'value': 1.0}]1.4210981.607066-0.496248[{'index': 2, 'value': 1.0}][{'index': 2, 'value': 1.0}]
98[{'index': 2, 'value': 1.0}]-0.5219950.176246-0.784304[{'index': 2, 'value': 1.0}][{'index': 1, 'value': 1.0}]
10[{'index': 2, 'value': 1.0}]-0.5769890.8916560.511947[{'index': 2, 'value': 1.0}][{'index': 1, 'value': 1.0}]
\n", + "
[267 rows x 6 columns in total]" + ], + "text/plain": [ + " onehotencoded_island scaled_culmen_length_mm \\\n", + "penguin_id \n", + "23 [{'index': 2, 'value': 1.0}] -1.310231 \n", + "83 [{'index': 2, 'value': 1.0}] 2.575955 \n", + "195 [{'index': 1, 'value': 1.0}] 0.816172 \n", + "264 [{'index': 1, 'value': 1.0}] 1.109469 \n", + "237 [{'index': 1, 'value': 1.0}] 0.816172 \n", + "261 [{'index': 1, 'value': 1.0}] 0.614531 \n", + "253 [{'index': 1, 'value': 1.0}] 0.321234 \n", + "164 [{'index': 1, 'value': 1.0}] 0.101261 \n", + "293 [{'index': 3, 'value': 1.0}] -0.686975 \n", + "176 [{'index': 1, 'value': 1.0}] 0.871166 \n", + "45 [{'index': 2, 'value': 1.0}] 1.476091 \n", + "185 [{'index': 1, 'value': 1.0}] 1.347773 \n", + "141 [{'index': 1, 'value': 1.0}] 1.036145 \n", + "147 [{'index': 1, 'value': 1.0}] 1.109469 \n", + "320 [{'index': 3, 'value': 1.0}] -0.925279 \n", + "285 [{'index': 1, 'value': 1.0}] -0.192036 \n", + "305 [{'index': 3, 'value': 1.0}] 0.027937 \n", + "37 [{'index': 2, 'value': 1.0}] 1.421098 \n", + "98 [{'index': 2, 'value': 1.0}] -0.521995 \n", + "10 [{'index': 2, 'value': 1.0}] -0.576989 \n", + "307 [{'index': 3, 'value': 1.0}] -1.420218 \n", + "337 [{'index': 3, 'value': 1.0}] -0.393678 \n", + "26 [{'index': 2, 'value': 1.0}] -0.448671 \n", + "97 [{'index': 2, 'value': 1.0}] -1.45688 \n", + "27 [{'index': 2, 'value': 1.0}] 1.201125 \n", + "\n", + " scaled_culmen_depth_mm scaled_flipper_length_mm \\\n", + "penguin_id \n", + "23 0.687253 -0.568262 \n", + "83 0.329548 -1.432429 \n", + "195 -1.407876 0.151877 \n", + "264 -0.436962 2.096254 \n", + "237 -1.305674 0.872017 \n", + "261 -0.947969 1.520142 \n", + "253 -1.663379 0.944031 \n", + "164 -0.743566 1.160073 \n", + "293 -0.079257 -1.792499 \n", + "176 -1.561178 0.655975 \n", + "45 1.811469 0.655975 \n", + "185 -1.510077 1.232087 \n", + "141 -1.101271 1.088059 \n", + "147 -0.641365 1.66417 \n", + "320 0.329548 -1.432429 \n", + "285 -2.072185 1.016045 \n", + "305 0.43175 0.655975 \n", + "37 1.607066 -0.496248 \n", + "98 0.176246 -0.784304 \n", + "10 0.891656 0.511947 \n", + "307 0.022944 -1.000346 \n", + "337 1.14716 -0.208192 \n", + "26 0.687253 0.00785 \n", + "97 0.380649 -0.784304 \n", + "27 1.249361 0.00785 \n", + "\n", + " onehotencoded_sex onehotencoded_species \n", + "penguin_id \n", + "23 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "83 [{'index': 1, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "195 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "264 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "237 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "261 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "253 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "164 [{'index': 0, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "293 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "176 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "45 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "185 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "141 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "147 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "320 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "285 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", + "305 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "37 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "98 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "10 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "307 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "337 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "26 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "97 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "27 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "...\n", + "\n", + "[267 rows x 6 columns]" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.compose import ColumnTransformer\n", + "from bigframes.ml.preprocessing import OneHotEncoder\n", + "\n", + "# Create an aggregate transform that applies StandardScaler to the numeric columns,\n", + "# and OneHotEncoder to the string columns\n", + "preproc = ColumnTransformer([\n", + " (\"scale\", StandardScaler(), [\"culmen_length_mm\", \"culmen_depth_mm\", \"flipper_length_mm\"]),\n", + " (\"encode\", OneHotEncoder(), [\"species\", \"sex\", \"island\"])])\n", + "\n", + "# Now we can fit all columns of the training data\n", + "preproc.fit(train_X)\n", + "\n", + "processed_train_X = preproc.transform(train_X)\n", + "processed_test_X = preproc.transform(test_X)\n", + "\n", + "processed_train_X" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Predictors\n", + "\n", + "Predictors are estimators that learn and make predictions. In addition to `.fit(...)`, the predictor implements a `.predict(...)` method, which will use what was learned during `.fit(...)` to predict some output.\n", + "\n", + "Predictors can be further broken down into two categories:\n", + "\n", + "#### Supervised predictors\n", + "\n", + "Supervised learning is when we train a model on input-output pairs, and then ask it to predict the output for new inputs. An example of such a predictor is `bigframes.ml.linear_models.LinearRegression`." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
penguin_id
2554789.541771
1584642.382347
2303562.484258
2185321.625743
2973926.179647
934158.344927
3384093.191646
2675590.390771
2753623.482531
2005494.402339
1073450.730176
2275707.948707
1994855.771085
2605477.219436
2014965.59121
73670.281516
84072.678667
1263339.624869
393556.890093
123302.203276
\n", + "
[67 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "penguin_id \n", + "255 4789.541771\n", + "158 4642.382347\n", + "230 3562.484258\n", + "218 5321.625743\n", + "297 3926.179647\n", + "93 4158.344927\n", + "338 4093.191646\n", + "267 5590.390771\n", + "275 3623.482531\n", + "200 5494.402339\n", + "107 3450.730176\n", + "227 5707.948707\n", + "199 4855.771085\n", + "260 5477.219436\n", + "201 4965.59121\n", + "7 3670.281516\n", + "8 4072.678667\n", + "126 3339.624869\n", + "39 3556.890093\n", + "12 3302.203276\n", + "34 3227.822355\n", + "206 3864.010614\n", + "53 3329.248666\n", + "282 5297.303727\n", + "225 5770.689468\n", + "...\n", + "\n", + "[67 rows x 1 columns]" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "\n", + "linreg = LinearRegression()\n", + "\n", + "# Learn from the training data how to predict output y\n", + "linreg.fit(processed_train_X, train_y)\n", + "\n", + "# Predict y for the test data\n", + "predicted_test_y = linreg.predict(processed_test_X)\n", + "\n", + "predicted_test_y" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Unsupervised predictors\n", + "\n", + "In unsupervised learning, there are no known outputs in the training data, instead the model learns on input data alone and predicts something else. An example of an unsupervised predictor is `bigframes.ml.cluster.KMeans`, which learns how to fit input data to a target number of clusters." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
CENTROID_ID
penguin_id
171
2302
2914
1773
2062
3312
182
2184
1162
2712
1653
1942
391
2553
2254
401
81
1554
2394
1383
\n", + "
[67 rows x 1 columns in total]" + ], + "text/plain": [ + " CENTROID_ID\n", + "penguin_id \n", + "17 1\n", + "230 2\n", + "291 4\n", + "177 3\n", + "206 2\n", + "331 2\n", + "18 2\n", + "218 4\n", + "116 2\n", + "271 2\n", + "165 3\n", + "194 2\n", + "39 1\n", + "255 3\n", + "225 4\n", + "40 1\n", + "8 1\n", + "155 4\n", + "239 4\n", + "138 3\n", + "89 1\n", + "34 1\n", + "106 2\n", + "12 2\n", + "158 3\n", + "...\n", + "\n", + "[67 rows x 1 columns]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.cluster import KMeans\n", + "\n", + "kmeans = KMeans(n_clusters=4)\n", + "\n", + "kmeans.fit(processed_train_X)\n", + "\n", + "kmeans.predict(processed_test_X)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Pipelines\n", + "\n", + "Transfomers and predictors can be chained into a single estimator component using `bigframes.ml.pipeline.Pipeline`:" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Pipeline(steps=[('preproc',\n", + " ColumnTransformer(transformers=[('scale', StandardScaler(),\n", + " ['culmen_length_mm',\n", + " 'culmen_depth_mm',\n", + " 'flipper_length_mm']),\n", + " ('encode', OneHotEncoder(),\n", + " ['species', 'sex',\n", + " 'island'])])),\n", + " ('linreg', LinearRegression())])" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.pipeline import Pipeline\n", + "\n", + "pipeline = Pipeline([\n", + " ('preproc', preproc),\n", + " ('linreg', linreg)\n", + "])\n", + "\n", + "# Print our pipeline\n", + "pipeline" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pipeline simplifies the workflow by applying each of its component steps automatically:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
penguin_id
1253911.489556
2605477.221344
934158.347237
2554789.543554
713153.676672
3273358.229939
3384093.194446
684133.390538
174013.257221
1164189.275197
1064019.304549
2275707.950809
3144037.311435
2463297.59755
1384971.171877
1263339.627076
1213462.463115
2314206.392245
404102.581912
603483.351408
\n", + "
[67 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "penguin_id \n", + "125 3911.489556\n", + "260 5477.221344\n", + "93 4158.347237\n", + "255 4789.543554\n", + "71 3153.676672\n", + "327 3358.229939\n", + "338 4093.194446\n", + "68 4133.390538\n", + "17 4013.257221\n", + "116 4189.275197\n", + "106 4019.304549\n", + "227 5707.950809\n", + "314 4037.311435\n", + "246 3297.59755\n", + "138 4971.171877\n", + "126 3339.627076\n", + "121 3462.463115\n", + "231 4206.392245\n", + "40 4102.581912\n", + "60 3483.351408\n", + "331 3370.089387\n", + "218 5321.627618\n", + "154 3618.98119\n", + "130 4249.575399\n", + "89 4145.000048\n", + "...\n", + "\n", + "[67 rows x 1 columns]" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pipeline.fit(train_X, train_y)\n", + "\n", + "predicted_test_y = pipeline.predict(test_X)\n", + "predicted_test_y" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the backend, a pipeline will actually be compiled into a single model with an embedded TRANSFORM step." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluating results\n", + "\n", + "Some models include a convenient `.score(X, y)` method for evaulation with a preset accuracy metric:" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.86245773619.1653690.004765209.6072890.8883680.901838
\n", + "
[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 223.862457 73619.165369 0.004765 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 209.607289 0.888368 0.901838 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# In the case of a pipeline, this will be equivalent to calling .score on the contained LinearRegression\n", + "pipeline.score(test_X, test_y)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For a more general approach, the library `bigframes.ml.metrics` is provided:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0.8883676417577451" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.metrics import r2_score\n", + "\n", + "r2_score(test_y, predicted_test_y)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Save/Load to BigQuery\n", + "\n", + "Estimators can be saved to BigQuery as BQML models, and loaded again in future" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Replace with a path where you have permission to save a model\n", + "model_name = \"bigframes-dev.bqml_tutorial.penguins_model\"\n", + "\n", + "linreg.to_gbq(model_name, replace=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# WARNING - until b/281709360 is fixed & pipeline is updated, pipelines will load as models,\n", + "# and details of their transform steps will be lost (the loaded model will behave the same)\n", + "bigframes.pandas.read_gbq_model(model_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "a850322d07d9bdc9ec5f301d307e048bcab2390ae395e1cbce9335f4e081e5e2" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/04 - Using ML - SKLearn linear regression.ipynb b/notebooks/04 - Using ML - SKLearn linear regression.ipynb new file mode 100644 index 00000000000..73cf224bef4 --- /dev/null +++ b/notebooks/04 - Using ML - SKLearn linear regression.ipynb @@ -0,0 +1,977 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using ML - SKLearn linear regression\n", + "\n", + "This demo shows how we can implement a linear regression in BigQuery DataFrames ML, with API that is exactly compatible with Scikit-Learn." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Init & load data" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
tag_numberspeciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
01225Gentoo penguin (Pygoscelis papua)Biscoe<NA><NA><NA><NA><NA>
11278Gentoo penguin (Pygoscelis papua)Biscoe42.013.5210.04150.0FEMALE
21275Gentoo penguin (Pygoscelis papua)Biscoe46.513.5210.04550.0FEMALE
31233Gentoo penguin (Pygoscelis papua)Biscoe43.314.0208.04575.0FEMALE
41311Gentoo penguin (Pygoscelis papua)Biscoe47.514.0212.04875.0FEMALE
51316Gentoo penguin (Pygoscelis papua)Biscoe49.114.5212.04625.0FEMALE
61313Gentoo penguin (Pygoscelis papua)Biscoe45.514.5212.04750.0FEMALE
71381Gentoo penguin (Pygoscelis papua)Biscoe47.614.5215.05400.0MALE
81377Gentoo penguin (Pygoscelis papua)Biscoe45.114.5207.05050.0FEMALE
91380Gentoo penguin (Pygoscelis papua)Biscoe45.114.5215.05000.0FEMALE
101257Gentoo penguin (Pygoscelis papua)Biscoe46.214.5209.04800.0FEMALE
111336Gentoo penguin (Pygoscelis papua)Biscoe46.514.5213.04400.0FEMALE
121237Gentoo penguin (Pygoscelis papua)Biscoe43.214.5208.04450.0FEMALE
131302Gentoo penguin (Pygoscelis papua)Biscoe48.515.0219.04850.0FEMALE
141325Gentoo penguin (Pygoscelis papua)Biscoe49.115.0228.05500.0MALE
151285Gentoo penguin (Pygoscelis papua)Biscoe47.515.0218.04950.0FEMALE
161242Gentoo penguin (Pygoscelis papua)Biscoe49.615.0216.04750.0MALE
171246Gentoo penguin (Pygoscelis papua)Biscoe47.715.0216.04750.0FEMALE
181320Gentoo penguin (Pygoscelis papua)Biscoe45.515.0220.05000.0MALE
191244Gentoo penguin (Pygoscelis papua)Biscoe46.415.0216.04700.0FEMALE
\n", + "
[347 rows x 8 columns in total]" + ], + "text/plain": [ + " tag_number species island culmen_length_mm \\\n", + "0 1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "1 1278 Gentoo penguin (Pygoscelis papua) Biscoe 42.0 \n", + "2 1275 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "3 1233 Gentoo penguin (Pygoscelis papua) Biscoe 43.3 \n", + "4 1311 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "5 1316 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "6 1313 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "7 1381 Gentoo penguin (Pygoscelis papua) Biscoe 47.6 \n", + "8 1377 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "9 1380 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "10 1257 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 \n", + "11 1336 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "12 1237 Gentoo penguin (Pygoscelis papua) Biscoe 43.2 \n", + "13 1302 Gentoo penguin (Pygoscelis papua) Biscoe 48.5 \n", + "14 1325 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "15 1285 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "16 1242 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "17 1246 Gentoo penguin (Pygoscelis papua) Biscoe 47.7 \n", + "18 1320 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "19 1244 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 \n", + "20 1390 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", + "21 1379 Gentoo penguin (Pygoscelis papua) Biscoe 47.8 \n", + "22 1267 Gentoo penguin (Pygoscelis papua) Biscoe 50.1 \n", + "23 1389 Gentoo penguin (Pygoscelis papua) Biscoe 47.2 \n", + "24 1269 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "0 \n", + "1 13.5 210.0 4150.0 FEMALE \n", + "2 13.5 210.0 4550.0 FEMALE \n", + "3 14.0 208.0 4575.0 FEMALE \n", + "4 14.0 212.0 4875.0 FEMALE \n", + "5 14.5 212.0 4625.0 FEMALE \n", + "6 14.5 212.0 4750.0 FEMALE \n", + "7 14.5 215.0 5400.0 MALE \n", + "8 14.5 207.0 5050.0 FEMALE \n", + "9 14.5 215.0 5000.0 FEMALE \n", + "10 14.5 209.0 4800.0 FEMALE \n", + "11 14.5 213.0 4400.0 FEMALE \n", + "12 14.5 208.0 4450.0 FEMALE \n", + "13 15.0 219.0 4850.0 FEMALE \n", + "14 15.0 228.0 5500.0 MALE \n", + "15 15.0 218.0 4950.0 FEMALE \n", + "16 15.0 216.0 4750.0 MALE \n", + "17 15.0 216.0 4750.0 FEMALE \n", + "18 15.0 220.0 5000.0 MALE \n", + "19 15.0 216.0 4700.0 FEMALE \n", + "20 15.0 223.0 5550.0 MALE \n", + "21 15.0 215.0 5650.0 MALE \n", + "22 15.0 225.0 5000.0 MALE \n", + "23 15.5 215.0 4975.0 FEMALE \n", + "24 16.0 225.0 5700.0 MALE \n", + "...\n", + "\n", + "[347 rows x 8 columns]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Initialize BigQuery DataFrame\n", + "import bigframes.pandas\n", + "\n", + "# read a BigQuery table to a BigQuery DataFrame\n", + "df = bigframes.pandas.read_gbq(\"bigframes-dev.bqml_tutorial.penguins\")\n", + "\n", + "# take a peek at the dataframe\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Data cleaning / prep" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
islandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
tag_number
1172Dream32.115.5188.03050.0FEMALE
1371Biscoe37.716.0183.03075.0FEMALE
1417Torgersen38.617.0188.02900.0FEMALE
1204Dream40.717.0190.03725.0MALE
1251Biscoe37.617.0185.03600.0FEMALE
1422Torgersen35.717.0189.03350.0FEMALE
1394Torgersen40.217.0176.03450.0FEMALE
1163Dream36.417.0195.03325.0FEMALE
1329Biscoe38.117.0181.03175.0FEMALE
1406Torgersen44.118.0210.04000.0MALE
1196Dream36.518.0182.03150.0FEMALE
1228Biscoe41.618.0192.03950.0MALE
1412Torgersen40.318.0195.03250.0FEMALE
1142Dream35.718.0202.03550.0FEMALE
1430Torgersen33.519.0190.03600.0FEMALE
1333Biscoe43.219.0197.04775.0MALE
1414Torgersen38.719.0195.03450.0FEMALE
1197Dream41.119.0182.03425.0MALE
1443Torgersen40.619.0199.04000.0MALE
1295Biscoe41.020.0203.04725.0MALE
\n", + "
[146 rows x 6 columns in total]" + ], + "text/plain": [ + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "tag_number \n", + "1172 Dream 32.1 15.5 188.0 \n", + "1371 Biscoe 37.7 16.0 183.0 \n", + "1417 Torgersen 38.6 17.0 188.0 \n", + "1204 Dream 40.7 17.0 190.0 \n", + "1251 Biscoe 37.6 17.0 185.0 \n", + "1422 Torgersen 35.7 17.0 189.0 \n", + "1394 Torgersen 40.2 17.0 176.0 \n", + "1163 Dream 36.4 17.0 195.0 \n", + "1329 Biscoe 38.1 17.0 181.0 \n", + "1406 Torgersen 44.1 18.0 210.0 \n", + "1196 Dream 36.5 18.0 182.0 \n", + "1228 Biscoe 41.6 18.0 192.0 \n", + "1412 Torgersen 40.3 18.0 195.0 \n", + "1142 Dream 35.7 18.0 202.0 \n", + "1430 Torgersen 33.5 19.0 190.0 \n", + "1333 Biscoe 43.2 19.0 197.0 \n", + "1414 Torgersen 38.7 19.0 195.0 \n", + "1197 Dream 41.1 19.0 182.0 \n", + "1443 Torgersen 40.6 19.0 199.0 \n", + "1295 Biscoe 41.0 20.0 203.0 \n", + "1207 Dream 38.8 20.0 190.0 \n", + "1349 Biscoe 38.2 20.0 190.0 \n", + "1350 Biscoe 37.8 20.0 190.0 \n", + "1351 Biscoe 38.1 16.5 198.0 \n", + "1116 Dream 37.0 16.5 185.0 \n", + "\n", + " body_mass_g sex \n", + "tag_number \n", + "1172 3050.0 FEMALE \n", + "1371 3075.0 FEMALE \n", + "1417 2900.0 FEMALE \n", + "1204 3725.0 MALE \n", + "1251 3600.0 FEMALE \n", + "1422 3350.0 FEMALE \n", + "1394 3450.0 FEMALE \n", + "1163 3325.0 FEMALE \n", + "1329 3175.0 FEMALE \n", + "1406 4000.0 MALE \n", + "1196 3150.0 FEMALE \n", + "1228 3950.0 MALE \n", + "1412 3250.0 FEMALE \n", + "1142 3550.0 FEMALE \n", + "1430 3600.0 FEMALE \n", + "1333 4775.0 MALE \n", + "1414 3450.0 FEMALE \n", + "1197 3425.0 MALE \n", + "1443 4000.0 MALE \n", + "1295 4725.0 MALE \n", + "1207 3950.0 MALE \n", + "1349 3900.0 MALE \n", + "1350 4250.0 MALE \n", + "1351 3825.0 FEMALE \n", + "1116 3400.0 FEMALE \n", + "...\n", + "\n", + "[146 rows x 6 columns]" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# set a friendlier index to uniquely identify the rows\n", + "df = df.set_index(\"tag_number\")\n", + "\n", + "# filter down to the data we want to analyze\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# drop the columns we don't care about\n", + "adelie_data = adelie_data.drop(columns=[\"species\"])\n", + "\n", + "# drop rows with nulls to get our training data\n", + "training_data = adelie_data.dropna()\n", + "\n", + "# take a peek at the training data\n", + "training_data" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Use `model_selection.train_test_split` to prepare training data" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "from bigframes.ml.model_selection import train_test_split\n", + "\n", + "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "label_columns = training_data[['body_mass_g']] \n", + "\n", + "train_X, test_X, train_y, test_y = train_test_split(\n", + " feature_columns, label_columns, test_size=0.2)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Configure a linear regression pipeline with preprocessing" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Pipeline(steps=[('preproc',\n", + " ColumnTransformer(transformers=[('onehot', OneHotEncoder(),\n", + " ['island', 'species', 'sex']),\n", + " ('scaler', StandardScaler(),\n", + " ['culmen_depth_mm',\n", + " 'culmen_length_mm',\n", + " 'flipper_length_mm'])])),\n", + " ('linreg', LinearRegression(fit_intercept=False))])" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "from bigframes.ml.pipeline import Pipeline\n", + "from bigframes.ml.compose import ColumnTransformer\n", + "from bigframes.ml.preprocessing import StandardScaler, OneHotEncoder\n", + "\n", + "preprocessing = ColumnTransformer([\n", + " (\"onehot\", OneHotEncoder(), [\"island\", \"species\", \"sex\"]),\n", + " (\"scaler\", StandardScaler(), [\"culmen_depth_mm\", \"culmen_length_mm\", \"flipper_length_mm\"]),\n", + "])\n", + "\n", + "model = LinearRegression(fit_intercept=False)\n", + "\n", + "pipeline = Pipeline([\n", + " ('preproc', preprocessing),\n", + " ('linreg', model)\n", + "])\n", + "\n", + "# TODO(bmil): pretty printing for pipelines\n", + "pipeline" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Fit the pipeline to the training data\n", + "\n", + "This will create a temporary BQML model in BigQuery" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "pipeline.fit(train_X, train_y)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Score the pipeline on the test data with `metrics.r2_score`" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0.4597467714807377" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.metrics import r2_score\n", + "\n", + "pred_y = pipeline.predict(test_X)\n", + "\n", + "r2_score(test_y, pred_y)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Inference the model on new data" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas\n", + "\n", + "new_penguins = bigframes.pandas.read_pandas(\n", + " pandas.DataFrame(\n", + " {\n", + " \"tag_number\": [1633, 1672, 1690],\n", + " \"species\": [\n", + " \"Adelie Penguin (Pygoscelis adeliae)\",\n", + " \"Adelie Penguin (Pygoscelis adeliae)\",\n", + " \"Adelie Penguin (Pygoscelis adeliae)\",\n", + " ],\n", + " \"island\": [\"Torgersen\", \"Torgersen\", \"Dream\"],\n", + " \"culmen_length_mm\": [39.5, 38.5, 37.9],\n", + " \"culmen_depth_mm\": [18.8, 17.2, 18.1],\n", + " \"flipper_length_mm\": [196.0, 181.0, 188.0],\n", + " \"sex\": [\"MALE\", \"FEMALE\", \"FEMALE\"],\n", + " }\n", + " ).set_index(\"tag_number\")\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
tag_number
16334034.682043
16723267.196667
16903445.920252
\n", + "
[3 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "tag_number \n", + "1633 4034.682043\n", + "1672 3267.196667\n", + "1690 3445.920252\n", + "\n", + "[3 rows x 1 columns]" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pipeline.predict(new_penguins)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Save in BigQuery" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Pipeline(steps=[('preproc',\n", + " ColumnTransformer(transformers=[('onehot', OneHotEncoder(),\n", + " ['island', 'species', 'sex']),\n", + " ('scaler', StandardScaler(),\n", + " ['culmen_depth_mm',\n", + " 'culmen_length_mm',\n", + " 'flipper_length_mm'])])),\n", + " ('linreg', LinearRegression(fit_intercept=False))])" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pipeline.to_gbq(\"bigframes-dev.bigframes_demo_us.penguin_model\", replace=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "a850322d07d9bdc9ec5f301d307e048bcab2390ae395e1cbce9335f4e081e5e2" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/05 - Using ML - Easy linear regression.ipynb b/notebooks/05 - Using ML - Easy linear regression.ipynb new file mode 100644 index 00000000000..90d6f256a7b --- /dev/null +++ b/notebooks/05 - Using ML - Easy linear regression.ipynb @@ -0,0 +1,920 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using ML - Easy linear regression\n", + "\n", + "This demo shows BigQuery DataFrames ML providing an SKLearn-like experience for\n", + "training a linear regression model.\n", + "\n", + "In this \"easy\" version of linear regression, we use a couple of BQML features to simplify our code:\n", + "\n", + "- We rely on automatic preprocessing to encode string values and scale numeric values\n", + "- We rely on automatic data split & evaluation to test the model\n", + "\n", + "This example is adapted from the [BQML linear regression tutorial](https://cloud.google.com/bigquery-ml/docs/linear-regression-tutorial)." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Init & load data" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
tag_numberspeciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
01225Gentoo penguin (Pygoscelis papua)Biscoe<NA><NA><NA><NA><NA>
11278Gentoo penguin (Pygoscelis papua)Biscoe42.013.5210.04150.0FEMALE
21275Gentoo penguin (Pygoscelis papua)Biscoe46.513.5210.04550.0FEMALE
31233Gentoo penguin (Pygoscelis papua)Biscoe43.314.0208.04575.0FEMALE
41311Gentoo penguin (Pygoscelis papua)Biscoe47.514.0212.04875.0FEMALE
51316Gentoo penguin (Pygoscelis papua)Biscoe49.114.5212.04625.0FEMALE
61313Gentoo penguin (Pygoscelis papua)Biscoe45.514.5212.04750.0FEMALE
71381Gentoo penguin (Pygoscelis papua)Biscoe47.614.5215.05400.0MALE
81377Gentoo penguin (Pygoscelis papua)Biscoe45.114.5207.05050.0FEMALE
91380Gentoo penguin (Pygoscelis papua)Biscoe45.114.5215.05000.0FEMALE
101257Gentoo penguin (Pygoscelis papua)Biscoe46.214.5209.04800.0FEMALE
111336Gentoo penguin (Pygoscelis papua)Biscoe46.514.5213.04400.0FEMALE
121237Gentoo penguin (Pygoscelis papua)Biscoe43.214.5208.04450.0FEMALE
131302Gentoo penguin (Pygoscelis papua)Biscoe48.515.0219.04850.0FEMALE
141325Gentoo penguin (Pygoscelis papua)Biscoe49.115.0228.05500.0MALE
151285Gentoo penguin (Pygoscelis papua)Biscoe47.515.0218.04950.0FEMALE
161242Gentoo penguin (Pygoscelis papua)Biscoe49.615.0216.04750.0MALE
171246Gentoo penguin (Pygoscelis papua)Biscoe47.715.0216.04750.0FEMALE
181320Gentoo penguin (Pygoscelis papua)Biscoe45.515.0220.05000.0MALE
191244Gentoo penguin (Pygoscelis papua)Biscoe46.415.0216.04700.0FEMALE
\n", + "
[347 rows x 8 columns in total]" + ], + "text/plain": [ + " tag_number species island culmen_length_mm \\\n", + "0 1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "1 1278 Gentoo penguin (Pygoscelis papua) Biscoe 42.0 \n", + "2 1275 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "3 1233 Gentoo penguin (Pygoscelis papua) Biscoe 43.3 \n", + "4 1311 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "5 1316 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "6 1313 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "7 1381 Gentoo penguin (Pygoscelis papua) Biscoe 47.6 \n", + "8 1377 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "9 1380 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "10 1257 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 \n", + "11 1336 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "12 1237 Gentoo penguin (Pygoscelis papua) Biscoe 43.2 \n", + "13 1302 Gentoo penguin (Pygoscelis papua) Biscoe 48.5 \n", + "14 1325 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "15 1285 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "16 1242 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "17 1246 Gentoo penguin (Pygoscelis papua) Biscoe 47.7 \n", + "18 1320 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "19 1244 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 \n", + "20 1390 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", + "21 1379 Gentoo penguin (Pygoscelis papua) Biscoe 47.8 \n", + "22 1267 Gentoo penguin (Pygoscelis papua) Biscoe 50.1 \n", + "23 1389 Gentoo penguin (Pygoscelis papua) Biscoe 47.2 \n", + "24 1269 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "0 \n", + "1 13.5 210.0 4150.0 FEMALE \n", + "2 13.5 210.0 4550.0 FEMALE \n", + "3 14.0 208.0 4575.0 FEMALE \n", + "4 14.0 212.0 4875.0 FEMALE \n", + "5 14.5 212.0 4625.0 FEMALE \n", + "6 14.5 212.0 4750.0 FEMALE \n", + "7 14.5 215.0 5400.0 MALE \n", + "8 14.5 207.0 5050.0 FEMALE \n", + "9 14.5 215.0 5000.0 FEMALE \n", + "10 14.5 209.0 4800.0 FEMALE \n", + "11 14.5 213.0 4400.0 FEMALE \n", + "12 14.5 208.0 4450.0 FEMALE \n", + "13 15.0 219.0 4850.0 FEMALE \n", + "14 15.0 228.0 5500.0 MALE \n", + "15 15.0 218.0 4950.0 FEMALE \n", + "16 15.0 216.0 4750.0 MALE \n", + "17 15.0 216.0 4750.0 FEMALE \n", + "18 15.0 220.0 5000.0 MALE \n", + "19 15.0 216.0 4700.0 FEMALE \n", + "20 15.0 223.0 5550.0 MALE \n", + "21 15.0 215.0 5650.0 MALE \n", + "22 15.0 225.0 5000.0 MALE \n", + "23 15.5 215.0 4975.0 FEMALE \n", + "24 16.0 225.0 5700.0 MALE \n", + "...\n", + "\n", + "[347 rows x 8 columns]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import bigframes.pandas\n", + "\n", + "# read a BigQuery table to a BigQuery DataFrame\n", + "df = bigframes.pandas.read_gbq(\"bigframes-dev.bqml_tutorial.penguins\")\n", + "\n", + "# take a peek at the dataframe\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Data cleaning / prep" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
islandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
tag_number
1172Dream32.115.5188.03050.0FEMALE
1371Biscoe37.716.0183.03075.0FEMALE
1417Torgersen38.617.0188.02900.0FEMALE
1204Dream40.717.0190.03725.0MALE
1251Biscoe37.617.0185.03600.0FEMALE
1422Torgersen35.717.0189.03350.0FEMALE
1394Torgersen40.217.0176.03450.0FEMALE
1163Dream36.417.0195.03325.0FEMALE
1329Biscoe38.117.0181.03175.0FEMALE
1406Torgersen44.118.0210.04000.0MALE
1196Dream36.518.0182.03150.0FEMALE
1228Biscoe41.618.0192.03950.0MALE
1412Torgersen40.318.0195.03250.0FEMALE
1142Dream35.718.0202.03550.0FEMALE
1430Torgersen33.519.0190.03600.0FEMALE
1333Biscoe43.219.0197.04775.0MALE
1414Torgersen38.719.0195.03450.0FEMALE
1197Dream41.119.0182.03425.0MALE
1443Torgersen40.619.0199.04000.0MALE
1295Biscoe41.020.0203.04725.0MALE
\n", + "
[146 rows x 6 columns in total]" + ], + "text/plain": [ + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "tag_number \n", + "1172 Dream 32.1 15.5 188.0 \n", + "1371 Biscoe 37.7 16.0 183.0 \n", + "1417 Torgersen 38.6 17.0 188.0 \n", + "1204 Dream 40.7 17.0 190.0 \n", + "1251 Biscoe 37.6 17.0 185.0 \n", + "1422 Torgersen 35.7 17.0 189.0 \n", + "1394 Torgersen 40.2 17.0 176.0 \n", + "1163 Dream 36.4 17.0 195.0 \n", + "1329 Biscoe 38.1 17.0 181.0 \n", + "1406 Torgersen 44.1 18.0 210.0 \n", + "1196 Dream 36.5 18.0 182.0 \n", + "1228 Biscoe 41.6 18.0 192.0 \n", + "1412 Torgersen 40.3 18.0 195.0 \n", + "1142 Dream 35.7 18.0 202.0 \n", + "1430 Torgersen 33.5 19.0 190.0 \n", + "1333 Biscoe 43.2 19.0 197.0 \n", + "1414 Torgersen 38.7 19.0 195.0 \n", + "1197 Dream 41.1 19.0 182.0 \n", + "1443 Torgersen 40.6 19.0 199.0 \n", + "1295 Biscoe 41.0 20.0 203.0 \n", + "1207 Dream 38.8 20.0 190.0 \n", + "1349 Biscoe 38.2 20.0 190.0 \n", + "1350 Biscoe 37.8 20.0 190.0 \n", + "1351 Biscoe 38.1 16.5 198.0 \n", + "1116 Dream 37.0 16.5 185.0 \n", + "\n", + " body_mass_g sex \n", + "tag_number \n", + "1172 3050.0 FEMALE \n", + "1371 3075.0 FEMALE \n", + "1417 2900.0 FEMALE \n", + "1204 3725.0 MALE \n", + "1251 3600.0 FEMALE \n", + "1422 3350.0 FEMALE \n", + "1394 3450.0 FEMALE \n", + "1163 3325.0 FEMALE \n", + "1329 3175.0 FEMALE \n", + "1406 4000.0 MALE \n", + "1196 3150.0 FEMALE \n", + "1228 3950.0 MALE \n", + "1412 3250.0 FEMALE \n", + "1142 3550.0 FEMALE \n", + "1430 3600.0 FEMALE \n", + "1333 4775.0 MALE \n", + "1414 3450.0 FEMALE \n", + "1197 3425.0 MALE \n", + "1443 4000.0 MALE \n", + "1295 4725.0 MALE \n", + "1207 3950.0 MALE \n", + "1349 3900.0 MALE \n", + "1350 4250.0 MALE \n", + "1351 3825.0 FEMALE \n", + "1116 3400.0 FEMALE \n", + "...\n", + "\n", + "[146 rows x 6 columns]" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# set a friendlier index to uniquely identify the rows\n", + "df = df.set_index(\"tag_number\")\n", + "\n", + "# filter down to the data we want to analyze\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# drop the columns we don't care about\n", + "adelie_data = adelie_data.drop(columns=[\"species\"])\n", + "\n", + "# drop rows with nulls to get our training data\n", + "training_data = adelie_data.dropna()\n", + "\n", + "# take a peek at the training data\n", + "training_data" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "# pick feature columns and label column\n", + "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "label_columns = training_data[['body_mass_g']] \n", + "\n", + "# also get the rows that we want to make predictions for (i.e. where the feature column is null)\n", + "missing_body_mass = adelie_data[adelie_data.body_mass_g.isnull()]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Create, score, fit, predict" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "\n", + "model = LinearRegression()\n", + "\n", + "# Here we pass the feature columns without transforms - BQML will then use\n", + "# automatic preprocessing to encode these columns\n", + "model.fit(feature_columns, label_columns)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.87876378553.6016340.005614181.3309110.6239510.623951
\n", + "
[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 223.878763 78553.601634 0.005614 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 181.330911 0.623951 0.623951 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# check how the model performed\n", + "model.score(feature_columns, label_columns)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
tag_number
13933459.735118
15253947.881639
15244304.175638
15233471.668379
\n", + "
[4 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "tag_number \n", + "1393 3459.735118\n", + "1525 3947.881639\n", + "1524 4304.175638\n", + "1523 3471.668379\n", + "\n", + "[4 rows x 1 columns]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# use the model to predict the missing labels\n", + "model.predict(missing_body_mass)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Save in BigQuery" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# save the model to a permanent location in BigQuery, so we can use it in future sessions (and elsewhere in BQ)\n", + "model.to_gbq(\"bqml_tutorial.penguins_model\", replace=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "a850322d07d9bdc9ec5f301d307e048bcab2390ae395e1cbce9335f4e081e5e2" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/06 - Using ML - Large Language Models.ipynb b/notebooks/06 - Using ML - Large Language Models.ipynb new file mode 100644 index 00000000000..c2c0b83bef2 --- /dev/null +++ b/notebooks/06 - Using ML - Large Language Models.ipynb @@ -0,0 +1,201 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import bigframes.pandas\n", + "import pandas as pd\n", + "from bigframes.ml.llm import PaLM2TextGenerator" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "Create session and define a BQ connection which we already created and allowlisted. " + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "session = bigframes.pandas.get_global_session()\n", + "connection = \"bigframes-dev.us.bigframes-ml\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Define the model" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "model = PaLM2TextGenerator(session=session, connection_name=connection)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create Sample Data\n", + "\n", + "Read as a BigQuery DataFrames." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "df = pd.DataFrame(\n", + " {\n", + " \"prompt\": [\"What is BigQuery?\", \"What is BQML?\", \"What is BigQuery DataFrame?\"],\n", + " })\n", + "bf_df = bigframes.pandas.read_pandas(df)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Make Predictions" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
ml_generate_text_llm_result
0BigQuery is a fully managed, petabyte-scale an...
1BQML stands for BigQuery Machine Learning. It ...
2A BigQuery DataFrames is a distributed collecti...
\n", + "
" + ], + "text/plain": [ + " ml_generate_text_llm_result\n", + "0 BigQuery is a fully managed, petabyte-scale an...\n", + "1 BQML stands for BigQuery Machine Learning. It ...\n", + "2 A BigQuery DataFrames is a distributed collecti..." + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pred = model.predict(bf_df).compute()\n", + "pred" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Fetch Predictions" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'BigQuery is a fully managed, petabyte-scale analytics data warehouse that enables businesses to analyze all their data very quickly. It is a cloud-based service that offers a pay-as-you-go pricing model. BigQuery is designed to handle large amounts of data and provide fast performance. It is a good choice for businesses that need to analyze large amounts of data quickly and easily.'" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pred.iloc[0, 0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/notebooks/50 - Remote Function.ipynb b/notebooks/50 - Remote Function.ipynb new file mode 100644 index 00000000000..76c87f8629c --- /dev/null +++ b/notebooks/50 - Remote Function.ipynb @@ -0,0 +1,839 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "3613b1cd", + "metadata": {}, + "outputs": [], + "source": [ + "TABLE='bigquery-public-data.stackoverflow.comments'\n", + "MAX_ROWS=1000000 # 1 Million" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f1175247", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 30.5 s, sys: 4.25 s, total: 34.8 s\n", + "Wall time: 1min 49s\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtextscore
010It will help if you give some details of which...6
125infact it does. Look a the first lines of your...10
227\"Currently + is implemented using StringBuffer...7
341I don't think that's the magic number he was r...18
459It's still very useful to know that magic numb...12
596This implementation is also nice if you wish t...9
6108That's not full text searching, it's searching...6
7109That's not full text searching, it's searching...6
8137In vim you can open > 1 buffer. :e filename. T...9
9154Sure, but what about a solution using O(1) mem...8
\n", + "
" + ], + "text/plain": [ + " id text score\n", + "0 10 It will help if you give some details of which... 6\n", + "1 25 infact it does. Look a the first lines of your... 10\n", + "2 27 \"Currently + is implemented using StringBuffer... 7\n", + "3 41 I don't think that's the magic number he was r... 18\n", + "4 59 It's still very useful to know that magic numb... 12\n", + "5 96 This implementation is also nice if you wish t... 9\n", + "6 108 That's not full text searching, it's searching... 6\n", + "7 109 That's not full text searching, it's searching... 6\n", + "8 137 In vim you can open > 1 buffer. :e filename. T... 9\n", + "9 154 Sure, but what about a solution using O(1) mem... 8" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "\n", + "##############################\n", + "# Pandas World #\n", + "##############################\n", + "\n", + "import pandas as pd\n", + "df = pd.read_gbq(TABLE, max_results=MAX_ROWS)[['id', 'text', 'score']]\n", + "df.head(10)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "fd8a04a3", + "metadata": {}, + "outputs": [], + "source": [ + "# User defined function\n", + "# https://www.codespeedy.com/find-nth-prime-number-in-python/\n", + "def nth_prime(n):\n", + " prime_numbers = [2,3]\n", + " i=3\n", + " if(02):\n", + " while (True):\n", + " i+=1\n", + " status = True\n", + " for j in range(2,int(i/2)+1):\n", + " if(i%j==0):\n", + " status = False\n", + " break\n", + " if(status==True):\n", + " prime_numbers.append(i)\n", + " if(len(prime_numbers)==n):\n", + " break\n", + " return prime_numbers[n-1]\n", + " else:\n", + " return -1" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "2b5e4568", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 5.48 s, sys: 27.6 ms, total: 5.51 s\n", + "Wall time: 5.49 s\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtextscoren_prime
010It will help if you give some details of which...613
125infact it does. Look a the first lines of your...1029
227\"Currently + is implemented using StringBuffer...717
341I don't think that's the magic number he was r...1861
459It's still very useful to know that magic numb...1237
596This implementation is also nice if you wish t...923
6108That's not full text searching, it's searching...613
7109That's not full text searching, it's searching...613
8137In vim you can open > 1 buffer. :e filename. T...923
9154Sure, but what about a solution using O(1) mem...819
\n", + "
" + ], + "text/plain": [ + " id text score n_prime\n", + "0 10 It will help if you give some details of which... 6 13\n", + "1 25 infact it does. Look a the first lines of your... 10 29\n", + "2 27 \"Currently + is implemented using StringBuffer... 7 17\n", + "3 41 I don't think that's the magic number he was r... 18 61\n", + "4 59 It's still very useful to know that magic numb... 12 37\n", + "5 96 This implementation is also nice if you wish t... 9 23\n", + "6 108 That's not full text searching, it's searching... 6 13\n", + "7 109 That's not full text searching, it's searching... 6 13\n", + "8 137 In vim you can open > 1 buffer. :e filename. T... 9 23\n", + "9 154 Sure, but what about a solution using O(1) mem... 8 19" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "\n", + "df = df.assign(n_prime=df['score'].apply(nth_prime))\n", + "df.head(10)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b81feaef", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 835 ms, sys: 62.5 ms, total: 897 ms\n", + "Wall time: 24.1 s\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtextscore
035153602Are you trying to access a nested array?0
135156124Sorry I didn't include my timeout method befor...0
235157401As soon as I defined some sort of primary Key ...0
335158649@user3355243 I've edited it to give $values an...0
435162039I pasted my exes @Matt.0
535162396@Gene - I do have separate fields, but I also ...0
635162907could you please provide any kind of sketch ho...0
735166498We use PhoneGap. Using Angular.0
835170138`decltype((int(Foo::*)(int))(&Foo::foo))` That...0
935172348Yes, I found that one already. And I understan...0
\n", + "
[10 rows x 3 columns in total]" + ], + "text/plain": [ + " id text score\n", + "0 35153602 Are you trying to access a nested array? 0\n", + "1 35156124 Sorry I didn't include my timeout method befor... 0\n", + "2 35157401 As soon as I defined some sort of primary Key ... 0\n", + "3 35158649 @user3355243 I've edited it to give $values an... 0\n", + "4 35162039 I pasted my exes @Matt. 0\n", + "5 35162396 @Gene - I do have separate fields, but I also ... 0\n", + "6 35162907 could you please provide any kind of sketch ho... 0\n", + "7 35166498 We use PhoneGap. Using Angular. 0\n", + "8 35170138 `decltype((int(Foo::*)(int))(&Foo::foo))` That... 0\n", + "9 35172348 Yes, I found that one already. And I understan... 0\n", + "\n", + "[10 rows x 3 columns]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "\n", + "##############################\n", + "# BigQuery DataFrames #\n", + "##############################\n", + "\n", + "import bigframes.pandas as pd\n", + "\n", + "df = pd.read_gbq(TABLE).head(MAX_ROWS)[['id', 'text', 'score']]\n", + "df.head(10)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "55ed241e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Help on function remote_function in module bigframes.pandas:\n", + "\n", + "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True)\n", + " Decorator to turn a user defined function into a BigQuery remote function.\n", + " \n", + " Args:\n", + " input_types : list(type)\n", + " List of input data types in the user defined function.\n", + " output_type : type\n", + " Data type of the output in the user defined function.\n", + " dataset : str, Optional\n", + " Dataset to use to create a BigQuery function. It should be in\n", + " `.` or `` format. If this\n", + " param is not provided then session dataset id would be used.\n", + " bigquery_connection : str, Optional\n", + " Name of the BigQuery connection. If it is pre created in the same\n", + " location as the `bigquery_client.location` then it would be used,\n", + " otherwise it would be created dynamically assuming the user has\n", + " necessary priviliges. If this param is not provided then the\n", + " bigquery connection from the session would be used.\n", + " reuse : bool, Optional\n", + " Reuse the remote function if already exists.\n", + " `True` by default, which will result in reusing an existing remote\n", + " function (if any) that was previously created for the same udf.\n", + " Setting it to false would force creating a unique remote function.\n", + " If the required remote function does not exist then it would be\n", + " created irrespective of this param.\n", + " \n", + " Notes:\n", + " Please make sure following is setup before using this API:\n", + " \n", + " 1. Have the below APIs enabled for your project:\n", + " a. BigQuery Connection API\n", + " b. Cloud Functions API\n", + " c. Cloud Run API\n", + " d. Cloud Build API\n", + " e. Artifact Registry API\n", + " f. Cloud Resource Manager API\n", + " \n", + " This can be done from the cloud console (change PROJECT_ID to yours):\n", + " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", + " Or from the gcloud CLI:\n", + " $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com\n", + " \n", + " 2. Have following IAM roles enabled for you:\n", + " a. BigQuery Data Editor (roles/bigquery.dataEditor)\n", + " b. BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", + " c. Cloud Functions Developer (roles/cloudfunctions.developer)\n", + " d. Service Account User (roles/iam.serviceAccountUser)\n", + " e. Storage Object Viewer (roles/storage.objectViewer)\n", + " f. Project IAM Admin (roles/resourcemanager.projectIamAdmin)\n", + " (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", + " \n", + " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", + " a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", + " b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", + " Alternatively, the IAM could also be setup via the gcloud CLI:\n", + " $ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"\n", + "\n" + ] + } + ], + "source": [ + "# Tell the user what needs to be done offline before using BigQuery DataFrame\n", + "# remote functions\n", + "help(pd.remote_function)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "c9a8d03d", + "metadata": {}, + "outputs": [], + "source": [ + "# BigQuery DataFrames user is a data scientist and may not have privileges to\n", + "# create a BQ connector and set it up for invoking a cloud function. They\n", + "# should get such a connector created from their cloud admin and use it with\n", + "# BigQuery DataFrames remote functions. If the provided connection name does not\n", + "# exist, BigQuery DataFrames will try to create it on the fly assuming the user\n", + "# has sufficient privileges.\n", + "bq_connection_name = 'bigframes-rf-conn'" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "fbc27f81", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[INFO][2023-06-29 01:03:30,557][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-ac72c931423f68dddb4f84f6754e2b28 --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmposspvw3v --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", + "Preparing function...\n", + ".done.\n", + "Deploying function...\n", + "[Build]......................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................done\n", + "[Service]....................................................................................................done\n", + "Done.\n", + "You can view your function in the Cloud Console here: https://console.cloud.google.com/functions/details/us-central1/bigframes-ac72c931423f68dddb4f84f6754e2b28?project=bigframes-dev\n", + "\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "buildConfig:\n", + " build: projects/1084210331973/locations/us-central1/builds/57859622-4287-47bc-ab86-3cd3f4c47548\n", + " entryPoint: udf_http\n", + " runtime: python310\n", + " source:\n", + " storageSource:\n", + " bucket: gcf-v2-sources-1084210331973-us-central1\n", + " object: bigframes-ac72c931423f68dddb4f84f6754e2b28/function-source.zip\n", + " sourceProvenance:\n", + " resolvedStorageSource:\n", + " bucket: gcf-v2-sources-1084210331973-us-central1\n", + " generation: '1688000614202977'\n", + " object: bigframes-ac72c931423f68dddb4f84f6754e2b28/function-source.zip\n", + "environment: GEN_2\n", + "labels:\n", + " deployment-tool: cli-gcloud\n", + "name: projects/bigframes-dev/locations/us-central1/functions/bigframes-ac72c931423f68dddb4f84f6754e2b28\n", + "serviceConfig:\n", + " allTrafficOnLatestRevision: true\n", + " availableCpu: '0.1666'\n", + " availableMemory: 256M\n", + " ingressSettings: ALLOW_ALL\n", + " maxInstanceCount: 100\n", + " maxInstanceRequestConcurrency: 1\n", + " revision: bigframes-ac72c931423f68dddb4f84f6754e2b28-00001-faf\n", + " service: projects/bigframes-dev/locations/us-central1/services/bigframes-ac72c931423f68dddb4f84f6754e2b28\n", + " serviceAccountEmail: 1084210331973-compute@developer.gserviceaccount.com\n", + " timeoutSeconds: 60\n", + " uri: https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app\n", + "state: ACTIVE\n", + "updateTime: '2023-06-29T01:04:34.306070995Z'\n", + "url: https://us-central1-bigframes-dev.cloudfunctions.net/bigframes-ac72c931423f68dddb4f84f6754e2b28\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[INFO][2023-06-29 01:04:37,239][bigframes.remote_function] Successfully created cloud function bigframes-ac72c931423f68dddb4f84f6754e2b28 with uri (https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app)\n", + "[INFO][2023-06-29 01:04:42,402][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", + "[INFO][2023-06-29 01:04:42,404][bigframes.remote_function] Creating BQ remote function: \n", + " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_ac72c931423f68dddb4f84f6754e2b28(n INT64)\n", + " RETURNS INT64\n", + " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", + " OPTIONS (\n", + " endpoint = \"https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app\"\n", + " )\n", + "[INFO][2023-06-29 01:04:43,456][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_ac72c931423f68dddb4f84f6754e2b28\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Wall time: 76.2659 s\n" + ] + } + ], + "source": [ + "from time import perf_counter\n", + "start = perf_counter()\n", + "#-------------------------------------------------------------------------------------\n", + "\n", + "# User defined function\n", + "# https://www.codespeedy.com/find-nth-prime-number-in-python/\n", + "@pd.remote_function([int], int, bigquery_connection=bq_connection_name)\n", + "def nth_prime(n):\n", + " prime_numbers = [2,3]\n", + " i=3\n", + " if(02):\n", + " while (True):\n", + " i+=1\n", + " status = True\n", + " for j in range(2,int(i/2)+1):\n", + " if(i%j==0):\n", + " status = False\n", + " break\n", + " if(status==True):\n", + " prime_numbers.append(i)\n", + " if(len(prime_numbers)==n):\n", + " break\n", + " return prime_numbers[n-1]\n", + " else:\n", + " return -1\n", + "\n", + "#-------------------------------------------------------------------------------------\n", + "print(f\"\\nWall time: {(perf_counter()-start):.4f} s\")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "c1c9355f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 22.2 ms, sys: 0 ns, total: 22.2 ms\n", + "Wall time: 22.4 ms\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtextscoren_prime
035153602Are you trying to access a nested array?0-1
135156124Sorry I didn't include my timeout method befor...0-1
235157401As soon as I defined some sort of primary Key ...0-1
335158649@user3355243 I've edited it to give $values an...0-1
435162039I pasted my exes @Matt.0-1
535162396@Gene - I do have separate fields, but I also ...0-1
635162907could you please provide any kind of sketch ho...0-1
735166498We use PhoneGap. Using Angular.0-1
835170138`decltype((int(Foo::*)(int))(&Foo::foo))` That...0-1
935172348Yes, I found that one already. And I understan...0-1
\n", + "
[10 rows x 4 columns in total]" + ], + "text/plain": [ + " id text score n_prime\n", + "0 35153602 Are you trying to access a nested array? 0 -1\n", + "1 35156124 Sorry I didn't include my timeout method befor... 0 -1\n", + "2 35157401 As soon as I defined some sort of primary Key ... 0 -1\n", + "3 35158649 @user3355243 I've edited it to give $values an... 0 -1\n", + "4 35162039 I pasted my exes @Matt. 0 -1\n", + "5 35162396 @Gene - I do have separate fields, but I also ... 0 -1\n", + "6 35162907 could you please provide any kind of sketch ho... 0 -1\n", + "7 35166498 We use PhoneGap. Using Angular. 0 -1\n", + "8 35170138 `decltype((int(Foo::*)(int))(&Foo::foo))` That... 0 -1\n", + "9 35172348 Yes, I found that one already. And I understan... 0 -1\n", + "\n", + "[10 rows x 4 columns]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "\n", + "df = df.assign(n_prime=df['score'].apply(nth_prime))\n", + "df.head(10)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/99 - Longer ML demo.ipynb b/notebooks/99 - Longer ML demo.ipynb new file mode 100644 index 00000000000..c4d133421f6 --- /dev/null +++ b/notebooks/99 - Longer ML demo.ipynb @@ -0,0 +1,1925 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "71fbfc47", + "metadata": {}, + "source": [ + "**Note: this notebook requires changes not yet checked in**\n", + "\n", + "# Introduction\n", + "\n", + "This is a prototype for how a minimal SKLearn-like wrapper for BQML might work in BigQuery DataFrames.\n", + "\n", + "Disclaimer - this is not a polished design or a robust implementation, this is a quick prototype to workshop some ideas. Design will be next.\n", + "\n", + "What is BigQuery DataFrame?\n", + "- Pandas API for BigQuery\n", + "- Lets data scientists quickly iterate and prepare their data as they do in Pandas, but executed by BigQuery\n", + "\n", + "What is meant by SKLearn-like?\n", + "- Follow the API design practices from the SKLearn project\n", + " - [API design for machine learning software: experiences from the scikit-learn project](https://arxiv.org/pdf/1309.0238.pdf)\n", + "- Not a copy of, or compatible with, SKLearn\n", + "\n", + "Briefly, patterns taken from SKLearn are:\n", + "- Models and transforms are 'Estimators'\n", + " - A bundle of parameters with a consistent way to initialize/get/set\n", + " - And a .fit(..) method to fit to training data\n", + "- Models additionally have a .predict(..)\n", + "- By default, these objects are transient, making them easy to play around with. No need to give them names or decide how to persist them.\n", + "\n", + "\n", + "Design goals:\n", + "- Zero friction ML capabilities for BigQuery DataFrames users (no extra auth, configuration, etc)\n", + "- Offers first class integration with the Pandas-like BigQuery DataFrames API\n", + "- Uses SKLearn-like design patterns that feel familiar to data scientists\n", + "- Also a first class BigQuery experience\n", + " - Offers BigQuery's scalability and storage / compute management\n", + " - Works naturally with BigQuery's other interfaces, e.g. GUI and SQL\n", + " - BQML features" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "345c2163", + "metadata": {}, + "source": [ + "# Linear regression tutorial\n", + "\n", + "Adapted from the \"Penguin weight\" Linear Regression tutorial for BQML: https://cloud.google.com/bigquery-ml/docs/linear-regression-tutorial\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "03c9e168", + "metadata": {}, + "source": [ + "## Setting the scene\n", + "\n", + "Our conservationists have sent us some measurements of penguins found in the Antarctic islands. They say that some of the body mass measurements for the Adelie penguins are missing, and ask if we can use some data science magic to estimate them. Sounds like a job for a linear regression!\n", + "\n", + "Lets take a look at the data..." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d7a03de2-c0ef-4f80-9cd5-f96e87cf2d54", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
tag_numberspeciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
01225Gentoo penguin (Pygoscelis papua)Biscoe<NA><NA><NA><NA><NA>
11278Gentoo penguin (Pygoscelis papua)Biscoe42.013.5210.04150.0FEMALE
21275Gentoo penguin (Pygoscelis papua)Biscoe46.513.5210.04550.0FEMALE
31233Gentoo penguin (Pygoscelis papua)Biscoe43.314.0208.04575.0FEMALE
41311Gentoo penguin (Pygoscelis papua)Biscoe47.514.0212.04875.0FEMALE
51316Gentoo penguin (Pygoscelis papua)Biscoe49.114.5212.04625.0FEMALE
61313Gentoo penguin (Pygoscelis papua)Biscoe45.514.5212.04750.0FEMALE
71381Gentoo penguin (Pygoscelis papua)Biscoe47.614.5215.05400.0MALE
81377Gentoo penguin (Pygoscelis papua)Biscoe45.114.5207.05050.0FEMALE
91380Gentoo penguin (Pygoscelis papua)Biscoe45.114.5215.05000.0FEMALE
101257Gentoo penguin (Pygoscelis papua)Biscoe46.214.5209.04800.0FEMALE
111336Gentoo penguin (Pygoscelis papua)Biscoe46.514.5213.04400.0FEMALE
121237Gentoo penguin (Pygoscelis papua)Biscoe43.214.5208.04450.0FEMALE
131302Gentoo penguin (Pygoscelis papua)Biscoe48.515.0219.04850.0FEMALE
141325Gentoo penguin (Pygoscelis papua)Biscoe49.115.0228.05500.0MALE
151285Gentoo penguin (Pygoscelis papua)Biscoe47.515.0218.04950.0FEMALE
161242Gentoo penguin (Pygoscelis papua)Biscoe49.615.0216.04750.0MALE
171246Gentoo penguin (Pygoscelis papua)Biscoe47.715.0216.04750.0FEMALE
181320Gentoo penguin (Pygoscelis papua)Biscoe45.515.0220.05000.0MALE
191244Gentoo penguin (Pygoscelis papua)Biscoe46.415.0216.04700.0FEMALE
\n", + "
[347 rows x 8 columns in total]" + ], + "text/plain": [ + " tag_number species island culmen_length_mm \\\n", + "0 1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "1 1278 Gentoo penguin (Pygoscelis papua) Biscoe 42.0 \n", + "2 1275 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "3 1233 Gentoo penguin (Pygoscelis papua) Biscoe 43.3 \n", + "4 1311 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "5 1316 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "6 1313 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "7 1381 Gentoo penguin (Pygoscelis papua) Biscoe 47.6 \n", + "8 1377 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "9 1380 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "10 1257 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 \n", + "11 1336 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "12 1237 Gentoo penguin (Pygoscelis papua) Biscoe 43.2 \n", + "13 1302 Gentoo penguin (Pygoscelis papua) Biscoe 48.5 \n", + "14 1325 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "15 1285 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "16 1242 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "17 1246 Gentoo penguin (Pygoscelis papua) Biscoe 47.7 \n", + "18 1320 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "19 1244 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 \n", + "20 1390 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", + "21 1379 Gentoo penguin (Pygoscelis papua) Biscoe 47.8 \n", + "22 1267 Gentoo penguin (Pygoscelis papua) Biscoe 50.1 \n", + "23 1389 Gentoo penguin (Pygoscelis papua) Biscoe 47.2 \n", + "24 1269 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "0 \n", + "1 13.5 210.0 4150.0 FEMALE \n", + "2 13.5 210.0 4550.0 FEMALE \n", + "3 14.0 208.0 4575.0 FEMALE \n", + "4 14.0 212.0 4875.0 FEMALE \n", + "5 14.5 212.0 4625.0 FEMALE \n", + "6 14.5 212.0 4750.0 FEMALE \n", + "7 14.5 215.0 5400.0 MALE \n", + "8 14.5 207.0 5050.0 FEMALE \n", + "9 14.5 215.0 5000.0 FEMALE \n", + "10 14.5 209.0 4800.0 FEMALE \n", + "11 14.5 213.0 4400.0 FEMALE \n", + "12 14.5 208.0 4450.0 FEMALE \n", + "13 15.0 219.0 4850.0 FEMALE \n", + "14 15.0 228.0 5500.0 MALE \n", + "15 15.0 218.0 4950.0 FEMALE \n", + "16 15.0 216.0 4750.0 MALE \n", + "17 15.0 216.0 4750.0 FEMALE \n", + "18 15.0 220.0 5000.0 MALE \n", + "19 15.0 216.0 4700.0 FEMALE \n", + "20 15.0 223.0 5550.0 MALE \n", + "21 15.0 215.0 5650.0 MALE \n", + "22 15.0 225.0 5000.0 MALE \n", + "23 15.5 215.0 4975.0 FEMALE \n", + "24 16.0 225.0 5700.0 MALE \n", + "...\n", + "\n", + "[347 rows x 8 columns]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import bigframes.pandas\n", + "\n", + "df = bigframes.pandas.read_gbq(\"bigframes-dev.bqml_tutorial.penguins\")\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "359524c4", + "metadata": {}, + "source": [ + "First we note that while we have a default numbered index generated by BigQuery, actually the penguins are uniquely identified by their tags.\n", + "\n", + "Lets make the data a bit friendlier to work with by setting the tag number column as the index." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "93d01411", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
tag_number
1225Gentoo penguin (Pygoscelis papua)Biscoe<NA><NA><NA><NA><NA>
1278Gentoo penguin (Pygoscelis papua)Biscoe42.013.5210.04150.0FEMALE
1275Gentoo penguin (Pygoscelis papua)Biscoe46.513.5210.04550.0FEMALE
1233Gentoo penguin (Pygoscelis papua)Biscoe43.314.0208.04575.0FEMALE
1311Gentoo penguin (Pygoscelis papua)Biscoe47.514.0212.04875.0FEMALE
1316Gentoo penguin (Pygoscelis papua)Biscoe49.114.5212.04625.0FEMALE
1313Gentoo penguin (Pygoscelis papua)Biscoe45.514.5212.04750.0FEMALE
1381Gentoo penguin (Pygoscelis papua)Biscoe47.614.5215.05400.0MALE
1377Gentoo penguin (Pygoscelis papua)Biscoe45.114.5207.05050.0FEMALE
1380Gentoo penguin (Pygoscelis papua)Biscoe45.114.5215.05000.0FEMALE
1257Gentoo penguin (Pygoscelis papua)Biscoe46.214.5209.04800.0FEMALE
1336Gentoo penguin (Pygoscelis papua)Biscoe46.514.5213.04400.0FEMALE
1237Gentoo penguin (Pygoscelis papua)Biscoe43.214.5208.04450.0FEMALE
1302Gentoo penguin (Pygoscelis papua)Biscoe48.515.0219.04850.0FEMALE
1325Gentoo penguin (Pygoscelis papua)Biscoe49.115.0228.05500.0MALE
1285Gentoo penguin (Pygoscelis papua)Biscoe47.515.0218.04950.0FEMALE
1242Gentoo penguin (Pygoscelis papua)Biscoe49.615.0216.04750.0MALE
1246Gentoo penguin (Pygoscelis papua)Biscoe47.715.0216.04750.0FEMALE
1320Gentoo penguin (Pygoscelis papua)Biscoe45.515.0220.05000.0MALE
1244Gentoo penguin (Pygoscelis papua)Biscoe46.415.0216.04700.0FEMALE
\n", + "
[347 rows x 7 columns in total]" + ], + "text/plain": [ + " species island culmen_length_mm \\\n", + "tag_number \n", + "1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "1278 Gentoo penguin (Pygoscelis papua) Biscoe 42.0 \n", + "1275 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "1233 Gentoo penguin (Pygoscelis papua) Biscoe 43.3 \n", + "1311 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "1316 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "1313 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "1381 Gentoo penguin (Pygoscelis papua) Biscoe 47.6 \n", + "1377 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "1380 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", + "1257 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 \n", + "1336 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "1237 Gentoo penguin (Pygoscelis papua) Biscoe 43.2 \n", + "1302 Gentoo penguin (Pygoscelis papua) Biscoe 48.5 \n", + "1325 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", + "1285 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", + "1242 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "1246 Gentoo penguin (Pygoscelis papua) Biscoe 47.7 \n", + "1320 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "1244 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 \n", + "1390 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", + "1379 Gentoo penguin (Pygoscelis papua) Biscoe 47.8 \n", + "1267 Gentoo penguin (Pygoscelis papua) Biscoe 50.1 \n", + "1389 Gentoo penguin (Pygoscelis papua) Biscoe 47.2 \n", + "1269 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "tag_number \n", + "1225 \n", + "1278 13.5 210.0 4150.0 FEMALE \n", + "1275 13.5 210.0 4550.0 FEMALE \n", + "1233 14.0 208.0 4575.0 FEMALE \n", + "1311 14.0 212.0 4875.0 FEMALE \n", + "1316 14.5 212.0 4625.0 FEMALE \n", + "1313 14.5 212.0 4750.0 FEMALE \n", + "1381 14.5 215.0 5400.0 MALE \n", + "1377 14.5 207.0 5050.0 FEMALE \n", + "1380 14.5 215.0 5000.0 FEMALE \n", + "1257 14.5 209.0 4800.0 FEMALE \n", + "1336 14.5 213.0 4400.0 FEMALE \n", + "1237 14.5 208.0 4450.0 FEMALE \n", + "1302 15.0 219.0 4850.0 FEMALE \n", + "1325 15.0 228.0 5500.0 MALE \n", + "1285 15.0 218.0 4950.0 FEMALE \n", + "1242 15.0 216.0 4750.0 MALE \n", + "1246 15.0 216.0 4750.0 FEMALE \n", + "1320 15.0 220.0 5000.0 MALE \n", + "1244 15.0 216.0 4700.0 FEMALE \n", + "1390 15.0 223.0 5550.0 MALE \n", + "1379 15.0 215.0 5650.0 MALE \n", + "1267 15.0 225.0 5000.0 MALE \n", + "1389 15.5 215.0 4975.0 FEMALE \n", + "1269 16.0 225.0 5700.0 MALE \n", + "...\n", + "\n", + "[347 rows x 7 columns]" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.set_index(\"tag_number\")\n", + "df" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "f95fda6a", + "metadata": {}, + "source": [ + "We saw in the first view that there were some missing values. We're especially interested in observations that are missing just the body_mass_g, so lets look at those:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "941cb6c3-8c54-42ce-a945-4fa604176b2e", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
tag_number
1225Gentoo penguin (Pygoscelis papua)Biscoe<NA><NA><NA><NA><NA>
1393Adelie Penguin (Pygoscelis adeliae)Torgersen<NA><NA><NA><NA><NA>
1524Adelie Penguin (Pygoscelis adeliae)Dream41.620.0204.0<NA>MALE
1523Adelie Penguin (Pygoscelis adeliae)Dream38.017.5194.0<NA>FEMALE
1525Adelie Penguin (Pygoscelis adeliae)Dream36.318.5194.0<NA>MALE
\n", + "
[5 rows x 7 columns in total]" + ], + "text/plain": [ + " species island culmen_length_mm \\\n", + "tag_number \n", + "1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "1393 Adelie Penguin (Pygoscelis adeliae) Torgersen \n", + "1524 Adelie Penguin (Pygoscelis adeliae) Dream 41.6 \n", + "1523 Adelie Penguin (Pygoscelis adeliae) Dream 38.0 \n", + "1525 Adelie Penguin (Pygoscelis adeliae) Dream 36.3 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "tag_number \n", + "1225 \n", + "1393 \n", + "1524 20.0 204.0 MALE \n", + "1523 17.5 194.0 FEMALE \n", + "1525 18.5 194.0 MALE \n", + "\n", + "[5 rows x 7 columns]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df[df.body_mass_g.isnull()]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a70c2027", + "metadata": {}, + "source": [ + "Here we see three Adelie penguins with tag numbers 1523, 1524, 1525 are missing their body_mass_g but have the other measurements. These are the ones we need to guess. We can do this by training a statistical model on the measurements that we do have, and then using it to predict the missing values.\n", + "\n", + "Our conservationists warned us that trying to generalize across species is a bad idea, so for now lets just try building a model for Adelie penguins. We can revisit it later and see if including the other observations improves the model performance." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "93ff013a", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
tag_number
1172Adelie Penguin (Pygoscelis adeliae)Dream32.115.5188.03050.0FEMALE
1371Adelie Penguin (Pygoscelis adeliae)Biscoe37.716.0183.03075.0FEMALE
1417Adelie Penguin (Pygoscelis adeliae)Torgersen38.617.0188.02900.0FEMALE
1204Adelie Penguin (Pygoscelis adeliae)Dream40.717.0190.03725.0MALE
1251Adelie Penguin (Pygoscelis adeliae)Biscoe37.617.0185.03600.0FEMALE
1422Adelie Penguin (Pygoscelis adeliae)Torgersen35.717.0189.03350.0FEMALE
1394Adelie Penguin (Pygoscelis adeliae)Torgersen40.217.0176.03450.0FEMALE
1163Adelie Penguin (Pygoscelis adeliae)Dream36.417.0195.03325.0FEMALE
1329Adelie Penguin (Pygoscelis adeliae)Biscoe38.117.0181.03175.0FEMALE
1406Adelie Penguin (Pygoscelis adeliae)Torgersen44.118.0210.04000.0MALE
1196Adelie Penguin (Pygoscelis adeliae)Dream36.518.0182.03150.0FEMALE
1228Adelie Penguin (Pygoscelis adeliae)Biscoe41.618.0192.03950.0MALE
1412Adelie Penguin (Pygoscelis adeliae)Torgersen40.318.0195.03250.0FEMALE
1142Adelie Penguin (Pygoscelis adeliae)Dream35.718.0202.03550.0FEMALE
1430Adelie Penguin (Pygoscelis adeliae)Torgersen33.519.0190.03600.0FEMALE
1333Adelie Penguin (Pygoscelis adeliae)Biscoe43.219.0197.04775.0MALE
1414Adelie Penguin (Pygoscelis adeliae)Torgersen38.719.0195.03450.0FEMALE
1197Adelie Penguin (Pygoscelis adeliae)Dream41.119.0182.03425.0MALE
1443Adelie Penguin (Pygoscelis adeliae)Torgersen40.619.0199.04000.0MALE
1295Adelie Penguin (Pygoscelis adeliae)Biscoe41.020.0203.04725.0MALE
\n", + "
[146 rows x 7 columns in total]" + ], + "text/plain": [ + " species island culmen_length_mm \\\n", + "tag_number \n", + "1172 Adelie Penguin (Pygoscelis adeliae) Dream 32.1 \n", + "1371 Adelie Penguin (Pygoscelis adeliae) Biscoe 37.7 \n", + "1417 Adelie Penguin (Pygoscelis adeliae) Torgersen 38.6 \n", + "1204 Adelie Penguin (Pygoscelis adeliae) Dream 40.7 \n", + "1251 Adelie Penguin (Pygoscelis adeliae) Biscoe 37.6 \n", + "1422 Adelie Penguin (Pygoscelis adeliae) Torgersen 35.7 \n", + "1394 Adelie Penguin (Pygoscelis adeliae) Torgersen 40.2 \n", + "1163 Adelie Penguin (Pygoscelis adeliae) Dream 36.4 \n", + "1329 Adelie Penguin (Pygoscelis adeliae) Biscoe 38.1 \n", + "1406 Adelie Penguin (Pygoscelis adeliae) Torgersen 44.1 \n", + "1196 Adelie Penguin (Pygoscelis adeliae) Dream 36.5 \n", + "1228 Adelie Penguin (Pygoscelis adeliae) Biscoe 41.6 \n", + "1412 Adelie Penguin (Pygoscelis adeliae) Torgersen 40.3 \n", + "1142 Adelie Penguin (Pygoscelis adeliae) Dream 35.7 \n", + "1430 Adelie Penguin (Pygoscelis adeliae) Torgersen 33.5 \n", + "1333 Adelie Penguin (Pygoscelis adeliae) Biscoe 43.2 \n", + "1414 Adelie Penguin (Pygoscelis adeliae) Torgersen 38.7 \n", + "1197 Adelie Penguin (Pygoscelis adeliae) Dream 41.1 \n", + "1443 Adelie Penguin (Pygoscelis adeliae) Torgersen 40.6 \n", + "1295 Adelie Penguin (Pygoscelis adeliae) Biscoe 41.0 \n", + "1207 Adelie Penguin (Pygoscelis adeliae) Dream 38.8 \n", + "1349 Adelie Penguin (Pygoscelis adeliae) Biscoe 38.2 \n", + "1350 Adelie Penguin (Pygoscelis adeliae) Biscoe 37.8 \n", + "1351 Adelie Penguin (Pygoscelis adeliae) Biscoe 38.1 \n", + "1116 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "\n", + " culmen_depth_mm flipper_length_mm body_mass_g sex \n", + "tag_number \n", + "1172 15.5 188.0 3050.0 FEMALE \n", + "1371 16.0 183.0 3075.0 FEMALE \n", + "1417 17.0 188.0 2900.0 FEMALE \n", + "1204 17.0 190.0 3725.0 MALE \n", + "1251 17.0 185.0 3600.0 FEMALE \n", + "1422 17.0 189.0 3350.0 FEMALE \n", + "1394 17.0 176.0 3450.0 FEMALE \n", + "1163 17.0 195.0 3325.0 FEMALE \n", + "1329 17.0 181.0 3175.0 FEMALE \n", + "1406 18.0 210.0 4000.0 MALE \n", + "1196 18.0 182.0 3150.0 FEMALE \n", + "1228 18.0 192.0 3950.0 MALE \n", + "1412 18.0 195.0 3250.0 FEMALE \n", + "1142 18.0 202.0 3550.0 FEMALE \n", + "1430 19.0 190.0 3600.0 FEMALE \n", + "1333 19.0 197.0 4775.0 MALE \n", + "1414 19.0 195.0 3450.0 FEMALE \n", + "1197 19.0 182.0 3425.0 MALE \n", + "1443 19.0 199.0 4000.0 MALE \n", + "1295 20.0 203.0 4725.0 MALE \n", + "1207 20.0 190.0 3950.0 MALE \n", + "1349 20.0 190.0 3900.0 MALE \n", + "1350 20.0 190.0 4250.0 MALE \n", + "1351 16.5 198.0 3825.0 FEMALE \n", + "1116 16.5 185.0 3400.0 FEMALE \n", + "...\n", + "\n", + "[146 rows x 7 columns]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# get all the rows with adelie penguins\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# separate out the rows that have a body mass measurement\n", + "training_data = adelie_data[adelie_data.body_mass_g.notnull()]\n", + "\n", + "# we noticed there were also some rows that were missing other values,\n", + "# lets remove these so they don't affect our results\n", + "training_data = training_data.dropna()\n", + "\n", + "# lets take a quick peek and make sure things look right:\n", + "training_data" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "d55a39f9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "species string[pyarrow]\n", + "island string[pyarrow]\n", + "culmen_length_mm Float64\n", + "culmen_depth_mm Float64\n", + "flipper_length_mm Float64\n", + "body_mass_g Float64\n", + "sex string[pyarrow]\n", + "dtype: object" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# we'll look at the schema too:\n", + "training_data.dtypes" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "59d374b5", + "metadata": {}, + "source": [ + "Great! Now lets configure a linear regression model to predict body mass from the other columns" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "18c4cecf", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import bigframes.ml.linear_model as ml\n", + "\n", + "model = ml.LinearRegression()\n", + "model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "6e54a1a2", + "metadata": {}, + "source": [ + "As in SKLearn, an unfitted model object is just a bundle of parameters." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a2060cf1", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'fit_intercept': True}" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# lets view the parameters\n", + "model.get_params()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "8e25fe41", + "metadata": {}, + "source": [ + "For this task, really all the default options are fine. But just so we can see how configuration works, lets specify that we want to use gradient descent to find the solution:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "327e2232", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.optimize_strategy = \"BATCH_GRADIENT_DESCENT\"\n", + "model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "2c2e0835", + "metadata": {}, + "source": [ + "BigQuery models provide a couple of extra conveniences:\n", + "\n", + "1. By default, they will automatically perform feature engineering on the inputs - encoding our string columns and scaling our numeric columns.\n", + "2. By default, they will also automatically manage the test/training data split for us.\n", + "\n", + "So all we need to do is hook our chosen feature and label columns into the model and call .fit()!" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "085c9a99", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "train_x = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "train_y = training_data[['body_mass_g']]\n", + "model.fit(train_x, train_y)\n", + "model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "9e76e10c", + "metadata": {}, + "source": [ + "...and there, we've successfully trained a linear regressor model. Lets see how it performs, using the automatic data split:" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "c9458c02", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.87876378553.6016340.005614181.3309110.6239510.623951
\n", + "
[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 223.878763 78553.601634 0.005614 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 181.330911 0.623951 0.623951 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.score(train_x, train_y)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "f0b39603", + "metadata": {}, + "source": [ + "Great! The model seems useful, predicting 62% of the variance.\n", + "\n", + "We realize we made a mistake though - we're trying to predict mass using a linear model, mass will increase with the cube of the penguin's size, whereas our inputs are linear with size. Can we improve our model by cubing them?" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "b94eddc7", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'\\ndef cubify(penguin_df):\\n penguin_df.culmen_length_mm = train_x.culmen_length_mm.pow(3)\\n penguin_df.culmen_depth_mm = train_x.culmen_depth_mm.pow(3)\\n penguin_df.flipper_length_mm = train_x.flipper_length_mm.pow(3)\\n\\ncubify(train_x)\\ntrain_x\\n'" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# SKIP THIS STEP (not yet work working in BigQuery DataFrame)\n", + "\n", + "# lets define a preprocessing step that adjust the linear measurements to use the cube\n", + "'''\n", + "def cubify(penguin_df):\n", + " penguin_df.culmen_length_mm = train_x.culmen_length_mm.pow(3)\n", + " penguin_df.culmen_depth_mm = train_x.culmen_depth_mm.pow(3)\n", + " penguin_df.flipper_length_mm = train_x.flipper_length_mm.pow(3)\n", + "\n", + "cubify(train_x)\n", + "train_x\n", + "'''" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "1b0e3f02", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'\\nmodel.fit(train_x, train_y)\\nmodel.evaluate()\\n'" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# AS ABOVE, SKIP FOR NOW\n", + "'''\n", + "model.fit(train_x, train_y)\n", + "model.evaluate()\n", + "'''" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "45c5e755", + "metadata": {}, + "source": [ + "Now that we're satisfied with our model, lets see what it predicts for those Adelie penguins with no body mass measurement:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "f21ebc1f", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
tag_number
13933459.735118
15244304.175638
15233471.668379
15253947.881639
\n", + "
[4 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "tag_number \n", + "1393 3459.735118\n", + "1524 4304.175638\n", + "1523 3471.668379\n", + "1525 3947.881639\n", + "\n", + "[4 rows x 1 columns]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Lets predict the missing observations\n", + "missing_body_mass = adelie_data[adelie_data.body_mass_g.isnull()]\n", + "\n", + "model.predict(missing_body_mass)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "e66bd0b0", + "metadata": {}, + "source": [ + "Because we created it without a name, it was just a temporary model that will disappear after 24 hours. \n", + "\n", + "We decide that this approach is promising, so lets tell BigQuery to save it." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "c508691b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.to_gbq(\"bqml_tutorial.penguins_model\", replace=True)\n", + "model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "46abef08", + "metadata": {}, + "source": [ + "We can now use this model anywhere in BigQuery with this name. We can also load\n", + "it again in our BigQuery DataFrames session and evaluate or inference it without\n", + "needing to retrain it:" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "0c87e972", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model = bigframes.pandas.read_gbq_model(\"bqml_tutorial.penguins_model\")\n", + "model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "d6ab8def", + "metadata": {}, + "source": [ + "And of course we can retrain it if we like. Lets make another version that is based on all the penguins, so we can test that assumption we made at the beginning that it would be best to separate them:" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "f4960452", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0224.71743379527.8796230.005693169.2358690.6192870.619287
\n", + "
[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 224.717433 79527.879623 0.005693 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 169.235869 0.619287 0.619287 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# This time we'll take all the training data, for all species\n", + "training_data = df[df.body_mass_g.notnull()]\n", + "training_data = training_data.dropna()\n", + "\n", + "# And we'll include species in our features\n", + "train_x = training_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "train_y = training_data[['body_mass_g']]\n", + "model.fit(train_x, train_y)\n", + "\n", + "# And we'll evaluate it on the Adelie penguins only\n", + "adelie_data = training_data[training_data.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "test_x = adelie_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "test_y = adelie_data[['body_mass_g']]\n", + "model.score(test_x, test_y)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7d101140", + "metadata": {}, + "source": [ + "It looks like the conservationists were right! Including other species, even though it gave us more training data, worsened prediction on the Adelie penguins." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7f3fe50d", + "metadata": {}, + "source": [ + "===============================================\n", + "\n", + "**Everything below this line not yet implemented**" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "62577c72", + "metadata": {}, + "source": [ + "We want to productionalize this model, so lets start publishing it to the vertex model registry ([prerequisites](https://cloud.google.com/bigquery-ml/docs/managing-models-vertex#prerequisites))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b82e79ee", + "metadata": {}, + "outputs": [], + "source": [ + "model.publish(\n", + " registry=\"vertex_ai\",\n", + " vertex_ai_model_version_aliases=[\"experimental\"])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "69d2482c", + "metadata": {}, + "source": [ + "Now when we fit the model, we can see it published here: https://console.cloud.google.com/vertex-ai/models" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "b97d9b64", + "metadata": {}, + "source": [ + "# Custom feature engineering" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "c837ace9", + "metadata": {}, + "source": [ + "So far, we've relied on BigQuery to do our feature engineering for us. What if we want to do it manually?\n", + "\n", + "BigQuery DataFrames provides a way to do this using Pipelines." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "480cb12f", + "metadata": {}, + "outputs": [], + "source": [ + "from bigframes.ml.pipeline import Pipeline\n", + "from bigframes.ml.preprocessing import StandardScaler\n", + "\n", + "pipe = Pipeline([\n", + " ('scaler', StandardScaler()),\n", + " ('linreg', LinearRegression())\n", + "])\n", + "\n", + "pipe.fit(train_x, train_y)\n", + "pipe.evaluate()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "9a0e7d19", + "metadata": {}, + "source": [ + "We then can then save the entire pipeline to BigQuery, BigQuery will save this as a single model, with the pre-processing steps embedded in the TRANSFORM property:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d1831ed", + "metadata": {}, + "outputs": [], + "source": [ + "pipe.to_gbq(\"bqml_tutorial.penguins_pipeline\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "f6b60898", + "metadata": {}, + "source": [ + "# Custom data split" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "60ac0174", + "metadata": {}, + "source": [ + "BigQuery has also managed splitting out our training data. What if we want to do this manually?\n", + "\n", + "*TODO: Write this section*" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + }, + "vscode": { + "interpreter": { + "hash": "a850322d07d9bdc9ec5f301d307e048bcab2390ae395e1cbce9335f4e081e5e2" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 00000000000..5e54c12fb2f --- /dev/null +++ b/noxfile.py @@ -0,0 +1,636 @@ +# -*- coding: utf-8 -*- +# +# Copyright 2018 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. + +from __future__ import absolute_import + +import os +import pathlib +import re +import shutil +from typing import Dict, List +import warnings + +import nox + +BLACK_VERSION = "black==22.3.0" +ISORT_VERSION = "isort==5.12.0" +SPHINX_VERSION = "sphinx==4.5.0" +LINT_PATHS = ["docs", "bigframes", "tests", "noxfile.py", "setup.py"] + +DEFAULT_PYTHON_VERSION = "3.10" + +UNIT_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11"] +UNIT_TEST_STANDARD_DEPENDENCIES = [ + "mock", + "asyncmock", + "pytest", + "pytest-cov", + "pytest-asyncio", +] +UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] +UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = [] +UNIT_TEST_DEPENDENCIES: List[str] = [] +UNIT_TEST_EXTRAS: List[str] = [] +UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {} + +SYSTEM_TEST_PYTHON_VERSIONS = ["3.9", "3.11"] +SYSTEM_TEST_STANDARD_DEPENDENCIES = [ + "jinja2", + "mock", + "openpyxl", + "pytest", + "pytest-cov", + "pytest-retry", + "pytest-timeout", + "pytest-xdist", + "google-cloud-testutils", + "tabulate", + "xarray", +] +SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [ + "google-cloud-bigquery", +] +SYSTEM_TEST_LOCAL_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_EXTRAS: List[str] = ["tests"] +SYSTEM_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {} + +CURRENT_DIRECTORY = pathlib.Path(__file__).parent.absolute() + +# Sessions are executed in the order so putting the smaller sessions +# ahead to fail fast at presubmit running. +# 'docfx' is excluded since it only needs to run in 'docs-presubmit' +nox.options.sessions = [ + "lint", + "lint_setup_py", + "mypy", + "format", + "docs", + "docfx", + "unit", + "unit_noextras", + "unit_prerelease", + "system", + "doctest", + "cover", + "release_dry_run", +] + +# Error if a python version is missing +nox.options.error_on_missing_interpreters = True + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def lint(session): + """Run linters. + + Returns a failure if the linters find linting errors or sufficiently + serious code quality issues. + """ + session.install("flake8", BLACK_VERSION) + session.run( + "black", + "--check", + *LINT_PATHS, + ) + session.run("flake8", "bigframes", "tests") + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def blacken(session): + """Run black. Format code to uniform standard.""" + session.install(BLACK_VERSION) + session.run( + "black", + *LINT_PATHS, + ) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def format(session): + """ + Run isort to sort imports. Then run black + to format code to uniform standard. + """ + session.install(BLACK_VERSION, ISORT_VERSION) + # 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", + *LINT_PATHS, + ) + session.run( + "black", + *LINT_PATHS, + ) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def lint_setup_py(session): + """Verify that setup.py is valid (including RST check).""" + session.install("docutils", "pygments") + session.run("python", "setup.py", "check", "--restructuredtext", "--strict") + + +def install_unittest_dependencies(session, install_test_extra, *constraints): + standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES + session.install(*standard_deps, *constraints) + + if UNIT_TEST_EXTERNAL_DEPENDENCIES: + warnings.warn( + "'unit_test_external_dependencies' is deprecated. Instead, please " + "use 'unit_test_dependencies' or 'unit_test_local_dependencies'.", + DeprecationWarning, + ) + session.install(*UNIT_TEST_EXTERNAL_DEPENDENCIES, *constraints) + + if UNIT_TEST_LOCAL_DEPENDENCIES: + session.install(*UNIT_TEST_LOCAL_DEPENDENCIES, *constraints) + + if install_test_extra and UNIT_TEST_EXTRAS_BY_PYTHON: + extras = UNIT_TEST_EXTRAS_BY_PYTHON.get(session.python, []) + elif install_test_extra and UNIT_TEST_EXTRAS: + extras = UNIT_TEST_EXTRAS + else: + extras = [] + + if extras: + session.install("-e", f".[{','.join(extras)}]", *constraints) + else: + session.install("-e", ".", *constraints) + + +def run_unit(session, install_test_extra): + """Run the unit test suite.""" + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + install_unittest_dependencies(session, install_test_extra, "-c", constraints_path) + + # Run py.test against the unit tests. + tests_path = os.path.join("tests", "unit") + session.run( + "py.test", + "--quiet", + f"--junitxml=unit_{session.python}_sponge_log.xml", + "--cov=bigframes", + f"--cov={tests_path}", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=term-missing", + "--cov-fail-under=0", + tests_path, + *session.posargs, + ) + + +@nox.session(python=UNIT_TEST_PYTHON_VERSIONS) +def unit(session): + run_unit(session, install_test_extra=True) + + +@nox.session(python=UNIT_TEST_PYTHON_VERSIONS[-1]) +def unit_noextras(session): + run_unit(session, install_test_extra=False) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def mypy(session): + """Run type checks with mypy.""" + session.install("-e", ".") + + # Just install the dependencies' type info directly, since "mypy --install-types" + # might require an additional pass. + deps = ( + set( + [ + "mypy", + "pandas-stubs", + "types-protobuf", + "types-python-dateutil", + "types-requests", + "types-setuptools", + ] + ) + | set(SYSTEM_TEST_STANDARD_DEPENDENCIES) + | set(UNIT_TEST_STANDARD_DEPENDENCIES) + ) + + session.install(*deps) + shutil.rmtree(".mypy_cache", ignore_errors=True) + session.run( + "mypy", + "bigframes", + os.path.join("tests", "system"), + os.path.join("tests", "unit"), + "--explicit-package-bases", + '--exclude="^third_party"', + ) + + +def install_systemtest_dependencies(session, install_test_extra, *constraints): + # Use pre-release gRPC for system tests. + # Exclude version 1.49.0rc1 which has a known issue. + # See https://github.com/grpc/grpc/pull/30642 + session.install("--pre", "grpcio!=1.49.0rc1") + + session.install(*SYSTEM_TEST_STANDARD_DEPENDENCIES, *constraints) + + if SYSTEM_TEST_EXTERNAL_DEPENDENCIES: + session.install(*SYSTEM_TEST_EXTERNAL_DEPENDENCIES, *constraints) + + if SYSTEM_TEST_LOCAL_DEPENDENCIES: + session.install("-e", *SYSTEM_TEST_LOCAL_DEPENDENCIES, *constraints) + + if SYSTEM_TEST_DEPENDENCIES: + session.install("-e", *SYSTEM_TEST_DEPENDENCIES, *constraints) + + if install_test_extra and SYSTEM_TEST_EXTRAS_BY_PYTHON: + extras = SYSTEM_TEST_EXTRAS_BY_PYTHON.get(session.python, []) + elif install_test_extra and SYSTEM_TEST_EXTRAS: + extras = SYSTEM_TEST_EXTRAS + else: + extras = [] + + if extras: + session.install("-e", f".[{','.join(extras)}]", *constraints) + else: + session.install("-e", ".", *constraints) + + +def run_system( + session, + prefix_name, + test_folder, + *, + check_cov=False, + install_test_extra=True, + print_duration=False, + extra_pytest_options=(), +): + """Run the system test suite.""" + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + + # Check the value of `RUN_SYSTEM_TESTS` env var. It defaults to true. + if os.environ.get("RUN_SYSTEM_TESTS", "true") == "false": + session.skip("RUN_SYSTEM_TESTS is set to false, skipping") + # Install pyopenssl for mTLS testing. + if os.environ.get("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") == "true": + session.install("pyopenssl") + + install_systemtest_dependencies(session, install_test_extra, "-c", constraints_path) + + # Run py.test against the system tests. + pytest_cmd = [ + "py.test", + "--quiet", + "-n 20", + # Any indivisual test taking longer than 10 mins will be terminated. + "--timeout=900", + f"--junitxml={prefix_name}_{session.python}_sponge_log.xml", + ] + if print_duration: + pytest_cmd.extend( + [ + "--durations=0", + ] + ) + if check_cov: + pytest_cmd.extend( + [ + "--cov=bigframes", + f"--cov={test_folder}", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=term-missing", + "--cov-fail-under=0", + ] + ) + + pytest_cmd.extend(extra_pytest_options) + session.run( + *pytest_cmd, + test_folder, + *session.posargs, + ) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +def system(session): + """Run the system test suite.""" + run_system( + session=session, + prefix_name="system", + test_folder=os.path.join("tests", "system", "small"), + check_cov=True, + ) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS[-1]) +def system_noextras(session): + """Run the system test suite.""" + run_system( + session=session, + prefix_name="system_noextras", + test_folder=os.path.join("tests", "system", "small"), + install_test_extra=False, + ) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS[-1]) +def doctest(session): + """Run the system test suite.""" + run_system( + session=session, + prefix_name="doctest", + extra_pytest_options=("--doctest-modules",), + test_folder="bigframes", + check_cov=True, + ) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS[-1]) +def e2e(session): + """Run the large tests in system test suite.""" + run_system( + session=session, + prefix_name="e2e", + test_folder=os.path.join("tests", "system", "large"), + print_duration=True, + ) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +def samples(session): + """Run the samples test suite.""" + + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + + # TODO(swast): Use `requirements.txt` files from the samples directories to + # test samples. + install_test_extra = True + install_systemtest_dependencies(session, install_test_extra, "-c", constraints_path) + + session.run( + "py.test", + "samples", + *session.posargs, + ) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def cover(session): + """Run the final coverage report. + + This outputs the coverage report aggregating coverage from the test runs + (including system test runs), and then erases coverage data. + """ + session.install("coverage", "pytest-cov") + session.run("coverage", "report", "--show-missing", "--fail-under=90") + + # Make sure there is no dead code in our test directories. + # TODO(swast): Cleanup dead code in the system tests directory. + session.run( + "coverage", + "report", + "--show-missing", + "--include=tests/unit/*", + "--fail-under=100", + ) + + session.run("coverage", "erase") + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def docs(session): + """Build the docs for this library.""" + + session.install("-e", ".") + session.install( + SPHINX_VERSION, + "alabaster", + "recommonmark", + ) + + shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True) + session.run( + "sphinx-build", + "-W", # warnings as errors + "-T", # show full traceback on exception + "-N", # no colors + "-b", + "html", + "-d", + os.path.join("docs", "_build", "doctrees", ""), + os.path.join("docs", ""), + os.path.join("docs", "_build", "html", ""), + ) + + +# docfx doesn't yet support Python 3.10. +# https://github.com/googleapis/sphinx-docfx-yaml/issues/305 +@nox.session(python="3.9") +def docfx(session): + """Build the docfx yaml files for this library.""" + + session.install("-e", ".") + session.install( + SPHINX_VERSION, + "alabaster", + "recommonmark", + "gcp-sphinx-docfx-yaml", + ) + + shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True) + session.run( + "sphinx-build", + "-T", # show full traceback on exception + "-N", # no colors + "-D", + ( + "extensions=sphinx.ext.autodoc," + "sphinx.ext.autosummary," + "docfx_yaml.extension," + "sphinx.ext.intersphinx," + "sphinx.ext.coverage," + "sphinx.ext.napoleon," + "sphinx.ext.todo," + "sphinx.ext.viewcode," + "recommonmark" + ), + "-b", + "html", + "-d", + os.path.join("docs", "_build", "doctrees", ""), + os.path.join("docs", ""), + os.path.join("docs", "_build", "html", ""), + ) + + +def prerelease(session, tests_path): + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + + # PyArrow prerelease packages are published to an alternative PyPI host. + # https://arrow.apache.org/docs/python/install.html#installing-nightly-packages + session.install( + "--extra-index-url", + "https://pypi.fury.io/arrow-nightlies/", + "--prefer-binary", + "--pre", + "--upgrade", + "pyarrow", + ) + session.install( + "--extra-index-url", + "https://pypi.anaconda.org/scipy-wheels-nightly/simple", + "--prefer-binary", + "--pre", + "--upgrade", + "pandas", + ) + session.install( + "--upgrade", + "-e", # Use -e so that py.typed file is included. + "git+https://github.com/ibis-project/ibis.git#egg=ibis-framework", + ) + # Workaround https://github.com/googleapis/python-db-dtypes-pandas/issues/178 + session.install("--no-deps", "db-dtypes") + + # Workaround to install pandas-gbq >=0.15.0, which is required by test only. + session.install("--no-deps", "pandas-gbq") + + session.install( + *set(UNIT_TEST_STANDARD_DEPENDENCIES + SYSTEM_TEST_STANDARD_DEPENDENCIES), + "-c", + constraints_path, + ) + + # Because we test minimum dependency versions on the minimum Python + # version, the first version we test with in the unit tests sessions has a + # constraints file containing all dependencies and extras. + with open( + CURRENT_DIRECTORY + / "testing" + / f"constraints-{UNIT_TEST_PYTHON_VERSIONS[0]}.txt", + encoding="utf-8", + ) as constraints_file: + constraints_text = constraints_file.read() + + # Ignore leading whitespace and comment lines. + already_installed = frozenset( + ("db-dtypes", "pandas", "pyarrow", "ibis-framework", "pandas-gbq") + ) + deps = [ + match.group(1) + for match in re.finditer( + r"^\s*(\S+)(?===\S+)", constraints_text, flags=re.MULTILINE + ) + if match.group(1) not in already_installed + ] + + # We use --no-deps to ensure that pre-release versions aren't overwritten + # by the version ranges in setup.py. + session.install(*deps) + session.install("--no-deps", "-e", ".") + + # Print out prerelease package versions. + session.run("python", "-m", "pip", "freeze") + + # Run py.test against the tests. + session.run( + "py.test", + "--quiet", + "-n 20", + # Any indivisual test taking longer than 10 mins will be terminated. + "--timeout=600", + f"--junitxml={os.path.split(tests_path)[-1]}_prerelease_{session.python}_sponge_log.xml", + "--cov=bigframes", + f"--cov={tests_path}", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=term-missing", + "--cov-fail-under=0", + tests_path, + *session.posargs, + ) + + +@nox.session(python=UNIT_TEST_PYTHON_VERSIONS[-1]) +def unit_prerelease(session): + """Run the unit test suite with prerelease dependencies.""" + prerelease(session, os.path.join("tests", "unit")) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS[-1]) +def system_prerelease(session): + """Run the system test suite with prerelease dependencies.""" + prerelease(session, os.path.join("tests", "system", "small")) + + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +def notebook(session): + session.install("-e", ".[all]") + session.install("pytest", "pytest-xdist", "pytest-retry", "nbmake") + + notebooks = [ + "00 - Summary.ipynb", + "01 - Getting Started.ipynb", + "02 - DataFrame.ipynb", + "03 - Using ML - ML fundamentals.ipynb", + "04 - Using ML - SKLearn linear regression.ipynb", + "05 - Using ML - Easy linear regression.ipynb", + "06 - Using ML - Large Language Models.ipynb", + "50 - Remote Function.ipynb", + ] + notebooks = [os.path.join("notebooks", nb) for nb in notebooks] + + # For some reason nbmake exits silently with "no tests ran" message if + # one of the notebook paths supplied does not exist. Let's make sure that + # each path exists + for nb in notebooks: + assert os.path.exists(nb), nb + + # Use retries because sometimes parallel runs of the same notebook can try + # to create the same artifacts and may run into resoure conflict at the GCP + # level. + session.run( + "py.test", + "-nauto", + "--nbmake", + "--nbmake-timeout=600", + "--retries=3", + *notebooks, + ) + + +@nox.session(python="3.10") +def release_dry_run(session): + env = {} + + # If the project root is not set, then take current directory as the project + # root. See the release script for how the project root is set/used. This is + # specially useful when the developer runs the nox session on local machine. + if not os.environ.get("PROJECT_ROOT") and not os.environ.get( + "KOKORO_ARTIFACTS_DIR" + ): + env["PROJECT_ROOT"] = "." + session.run(".kokoro/release-nightly.sh", "--dry-run", env=env) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000000..693439f47ca --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +filterwarnings = + ignore::pandas.errors.SettingWithCopyWarning 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/quickstart.py b/samples/snippets/quickstart.py new file mode 100644 index 00000000000..dc8ab7ba2d0 --- /dev/null +++ b/samples/snippets/quickstart.py @@ -0,0 +1,89 @@ +# 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): + # [START bigquery_bigframes_quickstart] + import bigframes.pandas as pd + + # TODO: (Optional) Setup your session with the configuration. Some of these + # settings cannot be changed once a session has started. + pd.options.bigquery.project = "your-gcp-project-id" + pd.options.bigquery.location = "us" + + # [START_EXCLUDE silent] + # Ignore the lines between "EXCLUDE" comments. They replace values so this + # sample runs in our test suite. + pd.options.bigquery.project = project_id + # [END_EXCLUDE] + + query_or_table = "bigquery-public-data.ml_datasets.penguins" + df = pd.read_gbq(query_or_table) + + # 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}") + + # IMPORTANT: The `bigframes.pandas` package creates a BigQuery session for + # queries and temporary tables. A BigQuery session has a limited lifetime + # (https://cloud.google.com/bigquery/docs/sessions-intro#limitations) and + # does not support concurrent queries. For long lived applications, create + # session objects as needed, instead. + + import bigframes + + session_options = bigframes.BigQueryOptions() + session_options.project = "your-gcp-project-id" + session_options.location = "us" + + # [START_EXCLUDE silent] + # Ignore the lines between "EXCLUDE" comments. They replace values so this + # sample runs in our test suite. + session_options.project = project_id + # [END_EXCLUDE] + + session = bigframes.connect(session_options) + df_session = session.read_gbq(query_or_table) + average_body_mass = df_session["body_mass_g"].mean() + print(f"average_body_mass (df_session): {average_body_mass}") + + 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] diff --git a/samples/snippets/quickstart_test.py b/samples/snippets/quickstart_test.py new file mode 100644 index 00000000000..6b0c69de999 --- /dev/null +++ b/samples/snippets/quickstart_test.py @@ -0,0 +1,31 @@ +# 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.reset_session() + + # TODO(swast): Get project from environment so contributors can run tests. + quickstart.run_quickstart("bigframes-dev") + out, _ = capsys.readouterr() + assert "average_body_mass (df_session):" in out diff --git a/scripts/generate_third_party_notices.py b/scripts/generate_third_party_notices.py new file mode 100644 index 00000000000..7040bb2e5f8 --- /dev/null +++ b/scripts/generate_third_party_notices.py @@ -0,0 +1,332 @@ +# 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 argparse +import glob +import importlib.metadata +import json +import os.path +import re +import sys + +import piplicenses +import requests + +DEPENDENCY_INFO_SEPARATOR = "*" * 80 + "\n" +PACKAGE_NAME_EXTRACTOR = re.compile("^[a-zA-Z0-9._-]+") + +# These packages don't have LICENSE files distributed in their packages, +# but we have manually confirmed they have a compatible license and +# included it manually in our `third_party` directory. +# +# TODO(swast): We can remove this workaround once these packages bundle the +# license file. +# +# ipython-genutils and recommonmark are both in an archived state with no likely updates in the future +# +# Tracking issues: +# * https://github.com/grpc/grpc/issues/33557 +# * https://github.com/gsnedders/python-webencodings/issues/33 +# * https://github.com/pickleshare/pickleshare/issues/34 +DIRECT_LICENSE_MAPPINGS = { + "grpcio-status": "https://raw.githubusercontent.com/grpc/grpc/master/LICENSE", + "webencodings": "https://raw.githubusercontent.com/gsnedders/python-webencodings/master/LICENSE", + "ipython-genutils": "https://raw.githubusercontent.com/ipython/ipython_genutils/master/COPYING.md", + "pickleshare": "https://raw.githubusercontent.com/pickleshare/pickleshare/master/LICENSE", + "recommonmark": "https://raw.githubusercontent.com/readthedocs/recommonmark/master/license.md", +} + + +def get_package_dependencies(pkg_name): + """Get all package dependencies for a given package, both required and optional.""" + packages = set() + requirements = importlib.metadata.requires(pkg_name) + if requirements: + for req in requirements: + match = PACKAGE_NAME_EXTRACTOR.match(req) + assert match, f"Could not parse {req} for package name" + packages.add(match.group(0)) + return packages + + +# Inspired by third_party/colab/cleanup_filesets.py +def find_dependencies( + roots: set[str], ignore_missing_metadata=False +) -> dict[str, dict[str, set[str]]]: + """Return the transitive dependencies of a set of packages. + Args: + roots: List of package names, e.g. ["pkg1", "pkg2"] + Returns: + A dictionary of dependencies, e.g. + { + "pkg3" : { + "Requires" : set(["pkg4", "pkg5", "pkg6"]), + "RequiredBy": set(["pkg1"]) + }, + "pkg4" : { + "Requires" : set([]), + "RequiredBy": set(["pkg3"]) + }, + ... + } + """ + hops = set() + visited = set() + deps: dict[str, dict[str, set[str]]] = dict() + + # Initialize the start of the graph walk + for root in roots: + # Get the normalized package name + try: + pkg = importlib.metadata.metadata(root) + except importlib.metadata.PackageNotFoundError: + if not ignore_missing_metadata: + raise + continue + hops.add(pkg["Name"]) + + # Start the graph walk + while True: + if not hops: + break + hop = hops.pop() + if hop in visited: + continue + visited.add(hop) + + for dep in get_package_dependencies(hop): + # Get the normalized package name + try: + req_pkg = importlib.metadata.metadata(dep) + except importlib.metadata.PackageNotFoundError: + if not ignore_missing_metadata: + raise + continue + dep = req_pkg["Name"] + + # Create outgoing edge only for non root packages, for which an + # entry must have been created in the deps dictionary when we + # saw the package for the first time during the graph walk + if hop in deps: + deps[hop]["Requires"].add(dep) + + if dep in deps: + # We have already seen this requirement in the graph walk. + # Just update the incoming dependency and carry on. + deps[dep]["RequiredBy"].add(hop) + else: + # This is the first time we came across this requirement. + # Create a new entry with the incoming dependency. + deps[dep] = {"RequiredBy": {hop}, "Requires": set()} + + # Put it in the next hops for further graph traversal + hops.add(dep) + + return deps + + +def get_metadata_and_filename( + package_name: str, + metadata_name: str, + metadata_file: str, + metadata_text: str, + ignore_missing=True, +) -> tuple[str, str] | None: + """Get package metadata and corresponsing file name.""" + + # Check metadata file + metadata_filepath_known = metadata_file != piplicenses.LICENSE_UNKNOWN + if not metadata_filepath_known and not ignore_missing: + raise ValueError(f"No {metadata_name} file found for {package_name}") + + # Check metadata text + if metadata_text != piplicenses.LICENSE_UNKNOWN: + output_filename = metadata_name + if metadata_filepath_known: + output_filename = os.path.basename(metadata_file) + if not output_filename: + raise ValueError( + f"Need a file name to write {metadata_name} text for {package_name}." + ) + return metadata_text, output_filename + elif not ignore_missing: + raise ValueError(f"No {metadata_name} text found for {package_name}") + + return None + + +def fetch_license_and_notice_metadata(packages: list[str]): + """Fetch metadata including license and notice for given packages. + Returns a json object. + """ + parser = piplicenses.create_parser() + args = parser.parse_args( + [ + "--format", + "json", + "--with-license-file", + "--with-notice-file", + "--with-urls", + "--with-description", + "--packages", + *packages, + ] + ) + output_str = piplicenses.create_output_string(args) + metadatas = json.loads(output_str) + return metadatas + + +def write_lines_without_trailing_spaces(file, text: str, key: str): + """Write text lines to a file without the trailing spaces. + This will stop complaints by the trailing-whitespace pre-commit hook.""" + text = "\n".join([line.rstrip() for line in text.split("\n")]) + file.write(f"{key}:\n{text}\n") + + +def write_metadata_to_file( + file, metadata, with_version=False, requires_packages=[], packages_required_by=[] +): + """Write package metadata to a file object.""" + file.write(DEPENDENCY_INFO_SEPARATOR) + + info_keys = ["Name"] + if with_version: + info_keys.append("Version") + info_keys.extend(["License", "URL"]) + file.writelines([f"{key}: {metadata[key]}\n" for key in info_keys]) + + if requires_packages: + file.write(f"Requires: {', '.join(sorted(requires_packages))}\n") + + if packages_required_by: + file.write(f"Required By: {', '.join(sorted(packages_required_by))}\n") + + # Try to generate third party license + + license_info = get_metadata_and_filename( + metadata["Name"], + "LICENSE", + metadata["LicenseFile"], + metadata["LicenseText"], + ignore_missing=metadata["Name"] in DIRECT_LICENSE_MAPPINGS, + ) + + license_text = "" + if license_info: + license_text = license_info[0] + else: + license_text_response = requests.get(DIRECT_LICENSE_MAPPINGS[metadata["Name"]]) + license_text = license_text_response.text + + write_lines_without_trailing_spaces(file, license_text, "License") + + # Try to generate third party notice + notice_info = get_metadata_and_filename( + metadata["Name"], + "NOTICE", + metadata["NoticeFile"], + metadata["NoticeText"], + ignore_missing=True, + ) + + if notice_info: + write_lines_without_trailing_spaces(file, notice_info[0], "Notice") + + file.write(DEPENDENCY_INFO_SEPARATOR) + + +def write_third_party_vendored_license(file, path): + """Write license of a vendored third party library to notices file.""" + file.write(DEPENDENCY_INFO_SEPARATOR) + file.write(f"Vendored Code: {os.path.dirname(path)}\n") + notice_key = f"Notice ({os.path.basename(path)})" + write_lines_without_trailing_spaces(file, open(path).read(), notice_key) + file.write(DEPENDENCY_INFO_SEPARATOR) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate third party notices for bigframes dependencies." + ) + parser.add_argument( + "--with-version", + action="store_true", + default=False, + help="Include the version information for each package.", + ) + parser.add_argument( + "--with-requires", + action="store_true", + default=False, + help="Include for each package the packages it requires.", + ) + parser.add_argument( + "--with-required-by", + action="store_true", + default=False, + help="Include for each package the packages that require it.", + ) + parser.add_argument( + "--output-file", + action="store", + default="THIRD_PARTY_NOTICES", + help="The output file to write third party notices in.", + ) + args = parser.parse_args(sys.argv[1:]) + + # Initialize the root package + roots = {"bigframes"} + + # Find dependencies + # Let's ignore the packages that are not installed assuming they are + # just the optional dependencies that bigframes does not require. + # One example is the dependency path bigframes -> SQLAlchemy -> pg8000, + # where pg8000 is only an optional dependency for SQLAlchemy which bigframes + # is not depending on + # https://github.com/sqlalchemy/sqlalchemy/blob/7bc81947e22dc32368b0c49a41c398cd251d94af/setup.cfg#LL62C21-L62C27 + deps = find_dependencies(roots, ignore_missing_metadata=True) + + # Use third party solution to fetch dependency metadata + deps_metadata = fetch_license_and_notice_metadata(list(deps)) + deps_metadata = sorted(deps_metadata, key=lambda m: m["Name"]) + + # Write the file + with open(args.output_file, "w") as f: + # Generate third party metadata for each dependency + for metadata in deps_metadata: + dep = deps[metadata["Name"]] + write_metadata_to_file( + f, + metadata, + args.with_version, + dep["Requires"] if args.with_requires else [], + dep["RequiredBy"] if args.with_required_by else [], + ) + + # Generate third party vendored notices + notices = set() + for filename in [ + "LICENCE", + "LICENCE.txt", + "LICENSE", + "LICENSE.txt", + "NOTICE", + "NOTICE.txt", + "COPYING", + "COPYING.txt", + ]: + notices.update(glob.glob(f"third_party/bigframes_vendored/*/{filename}")) + for path in sorted(notices): + write_third_party_vendored_license(f, path) diff --git a/scripts/publish_api_coverage.py b/scripts/publish_api_coverage.py new file mode 100644 index 00000000000..856307e4407 --- /dev/null +++ b/scripts/publish_api_coverage.py @@ -0,0 +1,179 @@ +# 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. + +"""This script runs with each release and publishes information about our pandas +and sklearn API coverage to BigQuery, where it can be used in our dashboards.""" + +import argparse +import inspect + +import pandas as pd + +import bigframes.pandas as bpd + + +def generate_pandas_api_coverage(): + """Inspect all our pandas objects, and compare with the real pandas objects, to see + which methods we implement. For each, generate a regex that can be used to check if + its present in a notebook""" + header = ["api", "pattern", "kind", "is_in_bigframes"] + api_patterns = [] + targets = [ + ("pandas", pd, bpd), + ("dataframe", pd.DataFrame, bpd.DataFrame), + ("series", pd.Series, bpd.Series), + ("index", pd.Index, bpd.Index), + ] + indexers = ["loc", "iloc", "iat", "ix", "at"] + for name, pandas_obj, bigframes_obj in targets: + for member in dir(pandas_obj): + # skip private functions and properties + if member[0] == "_" and member[1] != "_": + continue + + # skip members that are also common python methods + if member in {"append", "tolist", "format", "items", "keys"}: + continue + + if inspect.isfunction(getattr(pandas_obj, member)): + # Function, match .member( + token = f"\\.{member}\\(" + token_type = "function" + elif member in indexers: + # Indexer, match .indexer[ + token = f"\\.{member}\\[" + token_type = "indexer" + else: + # Property + token = f"\\.{member}\\b" + token_type = "property" + + is_in_bigframes = hasattr(bigframes_obj, member) + + api_patterns.append( + [f"{name}.{member}", token, token_type, is_in_bigframes] + ) + + return pd.DataFrame(api_patterns, columns=header) + + +def generate_sklearn_api_coverage(): + """Explore all SKLearn modules, and for each item contained generate a + regex to detect it being imported, and record whether we implement it""" + sklearn_modules = [ + "sklearn", + "sklearn.model_selection", + "sklearn.preprocessing", + "sklearn.metrics", + "sklearn.linear_model", + "sklearn.ensemble", + "sklearn.tree", + "sklearn.neighbors", + "sklearn.svm", + "sklearn.naive_bayes", + "sklearn.pipeline", + "sklearn.decomposition", + "sklearn.impute", + "sklearn.cluster", + "sklearn.feature_selection", + "sklearn.utils", + "sklearn.compose", + "sklearn.neural_network", + "sklearn.datasets", + "sklearn.base", + "sklearn.manifold", + "sklearn.discriminant_analysis", + "sklearn.experimental", + "sklearn.multiclass", + "sklearn.kernel_ridge", + "sklearn.feature_extraction", + "sklearn.dummy", + "sklearn.mixture", + "sklearn.gaussian_process", + "sklearn.calibration", + "sklearn.multioutput", + "sklearn.inspection", + "sklearn.exceptions", + "sklearn.cross_decomposition", + "sklearn.random_projection", + "sklearn.covariance", + "sklearn.semi_supervised", + "sklearn.isotonic", + "sklearn.kernel_approximation", + ] + + header = ["api", "pattern", "kind", "is_in_bigframes"] + api_patterns = [] + for module in sklearn_modules: + exec(f"import {module}") + members = eval(f"dir({module})") + bigframes_has_module = False + bigframes_members = [] + try: + bigframes_module = module.replace("sklearn", "bigframes.ml") + exec(f"import {bigframes_module}") + bigframes_has_module = True + bigframes_members = eval(f"dir({bigframes_module})") + except ImportError: + pass + + api_patterns.append( + [ + module, + f"from {module} import ", + "module", + bigframes_has_module, + ] + ) + for member in members: + # skip private functions and properties + if member[0] == "_": + continue + + api_patterns.append( + [ + f"{module}.{member}", + rf"from {module} import [^\n]*\b{member}\b", + "api", + member in bigframes_members, + ] + ) + + return pd.DataFrame(api_patterns, columns=header) + + +def build_api_coverage_table(bigframes_version: str, release_version: str): + pandas_cov_df = generate_pandas_api_coverage() + pandas_cov_df["module"] = "bigframes" + sklearn_cov_df = generate_sklearn_api_coverage() + sklearn_cov_df["module"] = "bigframes.ml" + combined_df = pd.concat([pandas_cov_df, sklearn_cov_df]) + combined_df["timestamp"] = pd.Timestamp.now() + combined_df["bigframes_version"] = bigframes_version + combined_df["release_version"] = release_version + return combined_df.infer_objects().convert_dtypes() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--bigframes_version") + parser.add_argument("--release_version") + parser.add_argument("--bigquery_table_name") + args = parser.parse_args() + df = build_api_coverage_table(args.bigframes_version, args.release_version) + df.to_gbq(args.bigquery_table_name, if_exists="append") + + +if __name__ == "__main__": + main() diff --git a/scripts/test_publish_api_coverage.py b/scripts/test_publish_api_coverage.py new file mode 100644 index 00000000000..96b2d1bb48d --- /dev/null +++ b/scripts/test_publish_api_coverage.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. + +import pandas +import publish_api_coverage + + +def test_api_coverage_produces_expected_schema(): + df = publish_api_coverage.build_api_coverage_table("my_bf_ver", "my_release_ver") + pandas.testing.assert_series_equal( + df.dtypes, + pandas.Series( + data=[ + "string", + "string", + "string", + "boolean", + "string", + "datetime64[ns]", + "string", + "string", + ], + index=[ + "api", + "pattern", + "kind", + "is_in_bigframes", + "module", + "timestamp", + "bigframes_version", + "release_version", + ], + ), + ) diff --git a/scripts/update_firebase_docs_site.sh b/scripts/update_firebase_docs_site.sh new file mode 100644 index 00000000000..f0ef866c90b --- /dev/null +++ b/scripts/update_firebase_docs_site.sh @@ -0,0 +1,105 @@ +#!/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 +# +# 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 -e -o pipefail +set -x + +# Replace the docs version by the desired commit +BIGFRAMES_DOCS_VERSION=8ce2a3e + +BIGFRAMES_FIREBASE_PROJECT=bigframes-dev-d4d9a +BIGFRAMES_DOCS_GCS_BUCKET=bigframes-docs + +BIGFRAMES_DOCS_DIR=`mktemp -d` + + +# Install firebase if not already installed +if ! which firebase; then + npm install -g firebase-tools +fi + +# Prepare a working directory for firebase +mkdir -p ${BIGFRAMES_DOCS_DIR} +pushd ${BIGFRAMES_DOCS_DIR} + +# Copy the bigframes version +if [ ! -d ${BIGFRAMES_DOCS_VERSION} ]; then + gsutil -m cp -r gs://${BIGFRAMES_DOCS_GCS_BUCKET}/${BIGFRAMES_DOCS_VERSION} . +fi + +rm -f latest +ln -s ${BIGFRAMES_DOCS_VERSION} latest + +# Set up firebase +firebase login --no-localhost +firebase init hosting + +versions="${BIGFRAMES_DOCS_VERSION} latest" +for version in ${versions}; do + site_name=bigframes-docs-${version} + if ! firebase hosting:sites:list | grep ${site_name}; then + firebase hosting:sites:create ${site_name} + fi + + firebase target:apply hosting ${version} ${site_name} +done + +# Make sure the firebase json config is consistent with ${versions} +# TODO(shobs): Come up with a better way of updating the config than +# a hard overwrite +cat > firebase.json << EOF +{ + "hosting": [ + { + "target": "latest", + "public": "latest", + "ignore": [ + "firebase.json", + "**/.*", + "**/node_modules/**" + ] + } + , + { + "target": "${BIGFRAMES_DOCS_VERSION}", + "public": "${BIGFRAMES_DOCS_VERSION}", + "ignore": [ + "firebase.json", + "**/.*", + "**/node_modules/**" + ] + } + ] +} +EOF + +# Verify that the intended sites look good +for version in ${versions}; do + echo "Preview the local hosting of the docs site \"${version}\" before actually deploying (Press Ctrl+C to stop)" + firebase serve --only hosting:${version} +done + +echo -n "Go ahead and deploy? [y/N]: " +read deploy_consent + +# Deploy the sites +if [ "$deploy_consent" = y ]; then + echo "Deploying ..." + firebase deploy --only hosting +else + echo "Not Deploying anything." +fi + +popd diff --git a/scripts/update_x20_docs_site.sh b/scripts/update_x20_docs_site.sh new file mode 100644 index 00000000000..31da116bdd4 --- /dev/null +++ b/scripts/update_x20_docs_site.sh @@ -0,0 +1,106 @@ +#!/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. + +################################################################################ +# This script uses corp credentials to sync the files to x20. +# Make sure you: +# 1. Have write permission to /x20/teams/bigframes-swe +# 2. Have run `gcert` before running the script +################################################################################ + +set -e -o pipefail +set -x + +GIT_DOCS_DIR=docs/_build/html/ +X20_BIGFRAMES_DIR=/x20/teams/bigframes-swe/bigframes/docs +MAX_BACKFILL=10 +REQ_PYTHON_VERSION=3.9 + +# Create a temporary directory +tempdir=`mktemp --directory --tmpdir bigframes.XXXXXXXXXX` + +# Clone the repository +git clone sso://team/bigquery-query-swe/bigframes ${tempdir} + +# Enter the temporary bigframes directory +pushd ${tempdir} + +# Python version 3.9 is required to build bigframes docs, install if not present +if ! python3 --version | grep ${REQ_PYTHON_VERSION}; then + # Install pyenv to install the required python version + ## https://github.com/pyenv/pyenv#basic-github-checkout + git clone https://github.com/pyenv/pyenv.git .pyenv + pushd .pyenv && src/configure && make -C src && popd + + ## https://github.com/pyenv/pyenv#set-up-your-shell-environment-for-pyenv + export PYENV_ROOT=${PWD}/.pyenv + PATH=${PYENV_ROOT}/bin:${PATH} + eval "$(pyenv init -)" + + ## Install the required python version + pyenv install ${REQ_PYTHON_VERSION} + + ## Make the required python version available + pyenv global ${REQ_PYTHON_VERSION} +fi + +# Create a virtual environment with nox installed +python3 -m venv venv +source venv/bin/activate +pip install nox + +# i = 0 means docs for the latest version, and i = 1 onwards means backfill +for i in `seq 0 ${MAX_BACKFILL}`; do + # If it is backfill turn, back off the version by 1 + if [ ${i} -ne 0 ]; then + git reset --hard HEAD~1 + + # Clean up any old docs + rm -rf ${GIT_DOCS_DIR} + fi + + # Construct a docs path in x20 + commit_hash=`git rev-parse --short HEAD` + x20_docs_dir_commit=${X20_BIGFRAMES_DIR}/${commit_hash} + + # If the x20 docs path already exists, let's assume that it was created + # properly in the previous attempt + if fileutil test -d ${x20_docs_dir_commit}; then + echo ${x20_docs_dir_commit} exists, skipping rebuilding it.. + continue + fi + + # Build the docs + echo Building docs for commit ${commit_hash}.. + nox -s docs + + # TODO(shobs): Check if a symlink can be created instead of another copy of + # the latest commit's docs, using fileutil CLI or otherwise + x20_docs_dirs=${x20_docs_dir_commit} + if [ ${i} -eq 0 ]; then + x20_docs_dirs="${x20_docs_dirs} ${X20_BIGFRAMES_DIR}/latest" + fi + + for x20_docs_dir in ${x20_docs_dirs}; do + fileutil mirror -parallelism=4 -force ${GIT_DOCS_DIR} ${x20_docs_dir} + x20_own request_change --recursive --path=${x20_docs_dir} --new_mode=a+r + done +done + +# Exit the temporary bigframes directory +popd + +# Clean up the temporary bigframes directory +rm -rf ${tempdir} diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000000..8bd749387ed --- /dev/null +++ b/setup.cfg @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +# +# 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. + +# Generated by synthtool. DO NOT EDIT! +[bdist_wheel] +universal = 1 + +[pytype] +python_version = 3.9 +inputs = + google/cloud/ +exclude = + tests/ + google/cloud/bigquery_v2/ # Legacy proto-based types. +output = .pytype/ +disable = + # There's some issue with finding some pyi files, thus disabling. + # The issue https://github.com/google/pytype/issues/150 is closed, but the + # error still occurs for some reason. + pyi-error diff --git a/setup.py b/setup.py new file mode 100644 index 00000000000..311471f2627 --- /dev/null +++ b/setup.py @@ -0,0 +1,109 @@ +# Copyright 2022 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 io +import itertools +import os +from typing import Dict + +import setuptools + +# Package metadata. + +name = "bigframes" +description = "Scalable DataFrames with BigQuery" + +# Should be one of: +# 'Development Status :: 3 - Alpha' +# 'Development Status :: 4 - Beta' +# 'Development Status :: 5 - Production/Stable' +release_status = "Development Status :: 3 - Alpha" +dependencies = [ + "cloudpickle >= 2.2.1", + "fsspec >=2023.3.0", + "gcsfs >=2023.3.0", + "geopandas >=0.12.2", + "google-auth >2.14.1,<3.0dev", + "google-cloud-bigquery[bqstorage,pandas] >=3.10.0", + "google-cloud-functions >=1.10.1", + "google-cloud-bigquery-connection >=1.12.0", + "google-cloud-storage >=2.0.0", + "ibis-framework[bigquery] >=6.0.0", + "pandas >=1.5.0", + "pydata-google-auth >=1.5.0", + "scikit-learn >=1.2.2", + "sqlalchemy >=1.4,<3.0", + "ipywidgets >=7.7.1", + "humanize >= 4.6.0", +] +extras = { + # Optional test dependencies packages. If they're missed, may skip some tests. + "tests": [ + "pandas-gbq >=0.19.0", + ], + # Packages required for basic development flow. + "dev": ["pytest", "pre-commit", "nox", "google-cloud-testutils"], +} +extras["all"] = list(sorted(frozenset(itertools.chain.from_iterable(extras.values())))) + +# Setup boilerplate below this line. + +package_root = os.path.abspath(os.path.dirname(__file__)) + +readme_filename = os.path.join(package_root, "README.rst") +with io.open(readme_filename, encoding="utf-8") as readme_file: + readme = readme_file.read() + +version: Dict[str, str] = {} +with open(os.path.join(package_root, "bigframes/version.py")) as fp: + exec(fp.read(), version) +version_id = version["__version__"] + +# Only include packages under the 'bigframes' namespace. Do not include tests, +# benchmarks, etc. +packages = [ + package + for package in setuptools.PEP420PackageFinder.find() + if package.startswith("bigframes") or package.startswith("third_party") +] + +setuptools.setup( + name=name, + version=version_id, + description=description, + long_description=readme, + author="Google LLC", + author_email="googleapis-packages@google.com", + license="Apache 2.0", + url="https://github.com/googleapis/python-bigquery", + classifiers=[ + release_status, + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Operating System :: OS Independent", + "Topic :: Internet", + ], + install_requires=dependencies, + extras_require=extras, + platforms="Posix; MacOS X; Windows", + packages=packages, + python_requires=">=3.9", + include_package_data=True, + zip_safe=False, +) diff --git a/testing/constraints-3.10.txt b/testing/constraints-3.10.txt new file mode 100644 index 00000000000..9f0786f47e6 --- /dev/null +++ b/testing/constraints-3.10.txt @@ -0,0 +1,15 @@ +# Keep in sync with colab/containers/requirements.core.in image +google-auth==2.17.3 +ipykernel==5.5.6 +ipython==7.34.0 +notebook==6.4.8 +pandas==1.5.3 +portpicker==1.3.9 +requests==2.27.1 +tornado==6.3.1 +absl-py==1.4.0 +debugpy==1.6.6 +ipywidgets==7.7.1 +matplotlib==3.7.1 +psutil==5.9.5 +traitlets==5.7.1 diff --git a/testing/constraints-3.11.txt b/testing/constraints-3.11.txt new file mode 100644 index 00000000000..e69de29bb2d diff --git a/testing/constraints-3.9.txt b/testing/constraints-3.9.txt new file mode 100644 index 00000000000..aa72e5426b7 --- /dev/null +++ b/testing/constraints-3.9.txt @@ -0,0 +1,118 @@ +argcomplete==2.1.2 +asyncmock==0.4.2 +atpublic==3.1.1 +attrs==22.2.0 +bidict==0.22.1 +black==23.3.0 +bleach==6.0.0 +cachetools==5.3.0 +certifi==2022.12.7 +cffi==1.15.1 +cfgv==3.3.1 +charset-normalizer==3.1.0 +click==8.1.3 +cloudpickle==2.2.1 +colorlog==6.7.0 +coverage==7.2.2 +cryptography==40.0.1 +distlib==0.3.6 +docstring-inheritance==2.0.0 +docutils==0.19 +exceptiongroup==1.1.1 +execnet==1.9.0 +filelock==3.10.7 +fsspec==2023.3.0 +gcp-docuploader==0.6.5 +gcp-releasetool==1.11.0 +gcsfs==2023.3.0 +geopandas==0.12.2 +google-api-core==2.11.0 +google-auth==2.17.0 +google-auth-oauthlib==1.0.0 +google-cloud-bigquery==3.10.0 +google-cloud-bigquery-connection==1.12.0 +google-cloud-bigquery-storage==2.19.1 +google-cloud-core==2.3.2 +google-cloud-functions==1.10.1 +google-cloud-storage==2.0.0 +google-cloud-testutils==1.3.3 +google-crc32c==1.5.0 +google-resumable-media==2.4.1 +googleapis-common-protos==1.59.0 +greenlet==2.0.2 +grpc-google-iam-v1==0.12.6 +grpcio==1.53.0 +grpcio-status==1.48.2 +ibis-framework==6.0.0 +humanize==4.6.0 +identify==2.5.22 +idna==3.4 +importlib-metadata==6.1.0 +iniconfig==2.0.0 +ipywidgets==7.7.1 +jaraco.classes==3.2.3 +jeepney==0.8.0 +Jinja2==3.1.2 +keyring==23.13.1 +markdown-it-py==2.2.0 +MarkupSafe==2.1.2 +mdurl==0.1.2 +mock==5.0.1 +more-itertools==9.1.0 +multipledispatch==0.6.0 +mypy-extensions==1.0.0 +nodeenv==1.7.0 +nox==2022.11.21 +numpy==1.24.2 +oauthlib==3.2.2 +packaging==23.0 +pandas==1.5.0 +pandas-gbq==0.19.0 +parsy==2.1 +pathspec==0.11.1 +pkginfo==1.9.6 +platformdirs==3.2.0 +pluggy==1.0.0 +pooch==1.7.0 +pre-commit==3.2.1 +proto-plus==1.22.2 +protobuf==3.20.3 +pyarrow==11.0.0 +pyasn1==0.4.8 +pyasn1-modules==0.2.8 +pycparser==2.21 +pydata-google-auth==1.5.0 +Pygments==2.14.0 +PyJWT==2.6.0 +pyperclip==1.8.2 +pytest==7.2.2 +pytest-asyncio==0.21.0 +pytest-cov==4.0.0 +pytest-retry==1.1.0 +pytest-xdist==3.2.1 +python-dateutil==2.8.2 +pytz==2023.3 +PyYAML==6.0 +readme-renderer==37.3 +requests==2.28.2 +requests-oauthlib==1.3.1 +requests-toolbelt==0.10.1 +rfc3986==2.0.0 +rich==13.3.3 +rsa==4.9 +scikit-learn==1.2.2 +SecretStorage==3.3.3 +six==1.16.0 +SQLAlchemy==1.4.0 +sqlglot==10.6.4 +tomli==2.0.1 +toolz==0.12.0 +tqdm==4.65.0 +twine==4.0.2 +typing_extensions==4.5.0 +tzdata==2023.3 +urllib3==1.26.15 +virtualenv==20.21.0 +webencodings==0.5.1 +xxhash==3.2.0 +zipp==3.15.0 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/__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/tests/data/penguins.jsonl b/tests/data/penguins.jsonl new file mode 100644 index 00000000000..d3aa06e3339 --- /dev/null +++ b/tests/data/penguins.jsonl @@ -0,0 +1,344 @@ +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.6,"culmen_depth_mm":18.4,"flipper_length_mm":184,"body_mass_g":3475,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.8,"culmen_depth_mm":19.1,"flipper_length_mm":184,"body_mass_g":4650,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.9,"culmen_depth_mm":18.9,"flipper_length_mm":184,"body_mass_g":3900,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.5,"culmen_depth_mm":17.9,"flipper_length_mm":192,"body_mass_g":3500,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.3,"culmen_depth_mm":16.8,"flipper_length_mm":192,"body_mass_g":3000,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":43.2,"culmen_depth_mm":18.5,"flipper_length_mm":192,"body_mass_g":4100,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.9,"culmen_depth_mm":16.6,"flipper_length_mm":192,"body_mass_g":2700,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.5,"culmen_depth_mm":18.4,"flipper_length_mm":200,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.5,"culmen_depth_mm":19,"flipper_length_mm":200,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.2,"culmen_depth_mm":20.1,"flipper_length_mm":200,"body_mass_g":3975,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.8,"culmen_depth_mm":18.9,"flipper_length_mm":208,"body_mass_g":4300,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39,"culmen_depth_mm":18.7,"flipper_length_mm":185,"body_mass_g":3650,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37,"culmen_depth_mm":16.9,"flipper_length_mm":185,"body_mass_g":3000,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":47,"culmen_depth_mm":17.3,"flipper_length_mm":185,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":34,"culmen_depth_mm":17.1,"flipper_length_mm":185,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37,"culmen_depth_mm":16.5,"flipper_length_mm":185,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.7,"culmen_depth_mm":17.3,"flipper_length_mm":193,"body_mass_g":3600,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.6,"culmen_depth_mm":19.4,"flipper_length_mm":193,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.7,"culmen_depth_mm":17.9,"flipper_length_mm":193,"body_mass_g":4250,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.8,"culmen_depth_mm":18.1,"flipper_length_mm":193,"body_mass_g":3750,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.6,"culmen_depth_mm":17.8,"flipper_length_mm":193,"body_mass_g":3800,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.3,"culmen_depth_mm":19.2,"flipper_length_mm":193,"body_mass_g":3650,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.2,"culmen_depth_mm":17.1,"flipper_length_mm":193,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.8,"culmen_depth_mm":18.5,"flipper_length_mm":193,"body_mass_g":3500,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.6,"culmen_depth_mm":18.2,"flipper_length_mm":193,"body_mass_g":3775,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52,"culmen_depth_mm":18.1,"flipper_length_mm":201,"body_mass_g":4050,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":41.5,"culmen_depth_mm":18.5,"flipper_length_mm":201,"body_mass_g":4000,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.5,"culmen_depth_mm":19.6,"flipper_length_mm":201,"body_mass_g":4050,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":54.2,"culmen_depth_mm":20.8,"flipper_length_mm":201,"body_mass_g":4300,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.4,"culmen_depth_mm":19,"flipper_length_mm":201,"body_mass_g":3950,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.8,"culmen_depth_mm":18.5,"flipper_length_mm":201,"body_mass_g":4450,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":33.1,"culmen_depth_mm":16.1,"flipper_length_mm":178,"body_mass_g":2900,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.2,"culmen_depth_mm":18.1,"flipper_length_mm":178,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.5,"culmen_depth_mm":16.7,"flipper_length_mm":178,"body_mass_g":3250,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.1,"culmen_depth_mm":18.2,"flipper_length_mm":178,"body_mass_g":3250,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36,"culmen_depth_mm":18.5,"flipper_length_mm":186,"body_mass_g":3100,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.6,"culmen_depth_mm":18.1,"flipper_length_mm":186,"body_mass_g":4450,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.7,"culmen_depth_mm":20.3,"flipper_length_mm":194,"body_mass_g":3775,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":41.3,"culmen_depth_mm":20.3,"flipper_length_mm":194,"body_mass_g":3550,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.6,"culmen_depth_mm":19.4,"flipper_length_mm":194,"body_mass_g":3525,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.2,"culmen_depth_mm":18.8,"flipper_length_mm":202,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":35.7,"culmen_depth_mm":18,"flipper_length_mm":202,"body_mass_g":3550,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":43.5,"culmen_depth_mm":18.1,"flipper_length_mm":202,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49,"culmen_depth_mm":19.5,"flipper_length_mm":210,"body_mass_g":3950,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.8,"culmen_depth_mm":19,"flipper_length_mm":210,"body_mass_g":4100,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52,"culmen_depth_mm":20.7,"flipper_length_mm":210,"body_mass_g":4800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.5,"culmen_depth_mm":18.9,"flipper_length_mm":179,"body_mass_g":2975} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.2,"culmen_depth_mm":17.5,"flipper_length_mm":187,"body_mass_g":3650,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.5,"culmen_depth_mm":18.7,"flipper_length_mm":187,"body_mass_g":3250,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":40.9,"culmen_depth_mm":16.6,"flipper_length_mm":187,"body_mass_g":3200,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":42.5,"culmen_depth_mm":16.7,"flipper_length_mm":187,"body_mass_g":3350,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":38.1,"culmen_depth_mm":17.6,"flipper_length_mm":187,"body_mass_g":3425,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":42.5,"culmen_depth_mm":17.3,"flipper_length_mm":187,"body_mass_g":3350,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36,"culmen_depth_mm":17.1,"flipper_length_mm":187,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.6,"culmen_depth_mm":17.2,"flipper_length_mm":187,"body_mass_g":3475,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.2,"culmen_depth_mm":17.3,"flipper_length_mm":187,"body_mass_g":3300,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":43.2,"culmen_depth_mm":16.6,"flipper_length_mm":187,"body_mass_g":2900,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.2,"culmen_depth_mm":18.2,"flipper_length_mm":195,"body_mass_g":4400,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":47.6,"culmen_depth_mm":18.3,"flipper_length_mm":195,"body_mass_g":3850,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.7,"culmen_depth_mm":17,"flipper_length_mm":195,"body_mass_g":3650,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.7,"culmen_depth_mm":17.9,"flipper_length_mm":195,"body_mass_g":3300,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46,"culmen_depth_mm":18.9,"flipper_length_mm":195,"body_mass_g":4150,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.4,"culmen_depth_mm":17,"flipper_length_mm":195,"body_mass_g":3325,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.7,"culmen_depth_mm":18.6,"flipper_length_mm":195,"body_mass_g":3600,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36,"culmen_depth_mm":17.8,"flipper_length_mm":195,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.8,"culmen_depth_mm":18.4,"flipper_length_mm":195,"body_mass_g":3900,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.3,"culmen_depth_mm":19.9,"flipper_length_mm":203,"body_mass_g":4050,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.7,"culmen_depth_mm":19.7,"flipper_length_mm":203,"body_mass_g":4050,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51,"culmen_depth_mm":18.8,"flipper_length_mm":203,"body_mass_g":4100,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":42.2,"culmen_depth_mm":18.5,"flipper_length_mm":180,"body_mass_g":3550,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.5,"culmen_depth_mm":17.8,"flipper_length_mm":188,"body_mass_g":3300,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":32.1,"culmen_depth_mm":15.5,"flipper_length_mm":188,"body_mass_g":3050,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.4,"culmen_depth_mm":18.7,"flipper_length_mm":188,"body_mass_g":3525,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.9,"culmen_depth_mm":19.1,"flipper_length_mm":196,"body_mass_g":3550,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.3,"culmen_depth_mm":18.5,"flipper_length_mm":196,"body_mass_g":4350,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.5,"culmen_depth_mm":17,"flipper_length_mm":196,"body_mass_g":3500,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50,"culmen_depth_mm":19.5,"flipper_length_mm":196,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.2,"culmen_depth_mm":21.1,"flipper_length_mm":196,"body_mass_g":4150,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":44.1,"culmen_depth_mm":19.7,"flipper_length_mm":196,"body_mass_g":4400,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.9,"culmen_depth_mm":17.9,"flipper_length_mm":196,"body_mass_g":3675,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49,"culmen_depth_mm":19.6,"flipper_length_mm":212,"body_mass_g":4300,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.6,"culmen_depth_mm":19.3,"flipper_length_mm":181,"body_mass_g":3300,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":42.4,"culmen_depth_mm":17.3,"flipper_length_mm":181,"body_mass_g":3600,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":58,"culmen_depth_mm":17.8,"flipper_length_mm":181,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":38.3,"culmen_depth_mm":19.2,"flipper_length_mm":189,"body_mass_g":3950,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.9,"culmen_depth_mm":18.6,"flipper_length_mm":189,"body_mass_g":3500,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.8,"culmen_depth_mm":16.5,"flipper_length_mm":189,"body_mass_g":3650,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52.2,"culmen_depth_mm":18.8,"flipper_length_mm":197,"body_mass_g":3450,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.3,"culmen_depth_mm":18.2,"flipper_length_mm":197,"body_mass_g":3750,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52.7,"culmen_depth_mm":19.8,"flipper_length_mm":197,"body_mass_g":3725,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52,"culmen_depth_mm":19,"flipper_length_mm":197,"body_mass_g":4150,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.3,"culmen_depth_mm":20,"flipper_length_mm":197,"body_mass_g":3300,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":53.5,"culmen_depth_mm":19.9,"flipper_length_mm":205,"body_mass_g":4500,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":41.1,"culmen_depth_mm":18.1,"flipper_length_mm":205,"body_mass_g":4300,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":52.8,"culmen_depth_mm":20,"flipper_length_mm":205,"body_mass_g":4550,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.5,"culmen_depth_mm":18,"flipper_length_mm":182,"body_mass_g":3150,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":41.1,"culmen_depth_mm":19,"flipper_length_mm":182,"body_mass_g":3425,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36,"culmen_depth_mm":17.9,"flipper_length_mm":190,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":41.1,"culmen_depth_mm":17.5,"flipper_length_mm":190,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":36.3,"culmen_depth_mm":19.5,"flipper_length_mm":190,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.2,"culmen_depth_mm":18.6,"flipper_length_mm":190,"body_mass_g":4250,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":38.9,"culmen_depth_mm":18.8,"flipper_length_mm":190,"body_mass_g":3600,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":39.6,"culmen_depth_mm":18.8,"flipper_length_mm":190,"body_mass_g":4600,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":40.7,"culmen_depth_mm":17,"flipper_length_mm":190,"body_mass_g":3725,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.9,"culmen_depth_mm":17.1,"flipper_length_mm":190,"body_mass_g":3575,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":38.1,"culmen_depth_mm":18.6,"flipper_length_mm":190,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":38.8,"culmen_depth_mm":20,"flipper_length_mm":190,"body_mass_g":3950,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.4,"culmen_depth_mm":18.6,"flipper_length_mm":190,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.1,"culmen_depth_mm":17.9,"flipper_length_mm":190,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.2,"culmen_depth_mm":17.8,"flipper_length_mm":198,"body_mass_g":3950,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.3,"culmen_depth_mm":19.9,"flipper_length_mm":198,"body_mass_g":3700,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":49.8,"culmen_depth_mm":17.3,"flipper_length_mm":198,"body_mass_g":3675,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":50.2,"culmen_depth_mm":18.7,"flipper_length_mm":198,"body_mass_g":3775,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":51.9,"culmen_depth_mm":19.5,"flipper_length_mm":206,"body_mass_g":3950,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":45.2,"culmen_depth_mm":16.6,"flipper_length_mm":191,"body_mass_g":3250,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":48.5,"culmen_depth_mm":17.5,"flipper_length_mm":191,"body_mass_g":3400,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":42.3,"culmen_depth_mm":21.2,"flipper_length_mm":191,"body_mass_g":4150,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.3,"culmen_depth_mm":17.8,"flipper_length_mm":191,"body_mass_g":3350,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":35.6,"culmen_depth_mm":17.5,"flipper_length_mm":191,"body_mass_g":3175,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":46.4,"culmen_depth_mm":17.8,"flipper_length_mm":191,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":47.5,"culmen_depth_mm":16.8,"flipper_length_mm":199,"body_mass_g":3900,"sex":"FEMALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":48.1,"culmen_depth_mm":16.4,"flipper_length_mm":199,"body_mass_g":3325,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Dream","culmen_length_mm":37.5,"culmen_depth_mm":18.5,"flipper_length_mm":199,"body_mass_g":4475,"sex":"MALE"} +{"species":"Chinstrap penguin (Pygoscelis antarctica)","island":"Dream","culmen_length_mm":55.8,"culmen_depth_mm":19.8,"flipper_length_mm":207,"body_mass_g":4000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":39.7,"culmen_depth_mm":18.9,"flipper_length_mm":184,"body_mass_g":3550,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":36.4,"culmen_depth_mm":17.1,"flipper_length_mm":184,"body_mass_g":2850,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41.6,"culmen_depth_mm":18,"flipper_length_mm":192,"body_mass_g":3950,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35,"culmen_depth_mm":17.9,"flipper_length_mm":192,"body_mass_g":3725,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41.1,"culmen_depth_mm":18.2,"flipper_length_mm":192,"body_mass_g":4050,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":42,"culmen_depth_mm":19.5,"flipper_length_mm":200,"body_mass_g":4050,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.8,"culmen_depth_mm":13.9,"flipper_length_mm":208,"body_mass_g":4300,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.3,"culmen_depth_mm":14,"flipper_length_mm":208,"body_mass_g":4575,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44,"culmen_depth_mm":13.6,"flipper_length_mm":208,"body_mass_g":4350,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":42.7,"culmen_depth_mm":13.7,"flipper_length_mm":208,"body_mass_g":3950,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.3,"culmen_depth_mm":13.8,"flipper_length_mm":208,"body_mass_g":4200,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.2,"culmen_depth_mm":14.5,"flipper_length_mm":208,"body_mass_g":4450,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.7,"culmen_depth_mm":15.7,"flipper_length_mm":208,"body_mass_g":5350,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.5,"culmen_depth_mm":15.2,"flipper_length_mm":216,"body_mass_g":5000,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.4,"culmen_depth_mm":15.8,"flipper_length_mm":216,"body_mass_g":4925,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.3,"culmen_depth_mm":13.8,"flipper_length_mm":216,"body_mass_g":4725} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.6,"culmen_depth_mm":15,"flipper_length_mm":216,"body_mass_g":4750,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.5,"culmen_depth_mm":14.3,"flipper_length_mm":216,"body_mass_g":4100} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.4,"culmen_depth_mm":15,"flipper_length_mm":216,"body_mass_g":4700,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49,"culmen_depth_mm":16.1,"flipper_length_mm":216,"body_mass_g":5550,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.7,"culmen_depth_mm":15,"flipper_length_mm":216,"body_mass_g":4750,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.5,"culmen_depth_mm":16.1,"flipper_length_mm":224,"body_mass_g":5650,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50,"culmen_depth_mm":15.9,"flipper_length_mm":224,"body_mass_g":5350,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.4,"culmen_depth_mm":15.3,"flipper_length_mm":224,"body_mass_g":5550,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35.7,"culmen_depth_mm":16.9,"flipper_length_mm":185,"body_mass_g":3150,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.6,"culmen_depth_mm":17,"flipper_length_mm":185,"body_mass_g":3600,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.2,"culmen_depth_mm":18.1,"flipper_length_mm":185,"body_mass_g":3950,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":40.6,"culmen_depth_mm":18.8,"flipper_length_mm":193,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":39.7,"culmen_depth_mm":17.7,"flipper_length_mm":193,"body_mass_g":3200,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.9,"culmen_depth_mm":18.6,"flipper_length_mm":193,"body_mass_g":2925,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.1,"culmen_depth_mm":15.1,"flipper_length_mm":209,"body_mass_g":5500,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.2,"culmen_depth_mm":14.5,"flipper_length_mm":209,"body_mass_g":4800,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.3,"culmen_depth_mm":13.4,"flipper_length_mm":209,"body_mass_g":4400,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":42.8,"culmen_depth_mm":14.2,"flipper_length_mm":209,"body_mass_g":4700,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.5,"culmen_depth_mm":14.2,"flipper_length_mm":209,"body_mass_g":4600,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.6,"culmen_depth_mm":13.9,"flipper_length_mm":217,"body_mass_g":4900,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.3,"culmen_depth_mm":15.7,"flipper_length_mm":217,"body_mass_g":5850,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.5,"culmen_depth_mm":14.8,"flipper_length_mm":217,"body_mass_g":5200,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.5,"culmen_depth_mm":14.4,"flipper_length_mm":217,"body_mass_g":4900,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.5,"culmen_depth_mm":15.7,"flipper_length_mm":217,"body_mass_g":4875,"sex":"."} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.2,"culmen_depth_mm":14.1,"flipper_length_mm":217,"body_mass_g":4375,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.1,"culmen_depth_mm":15,"flipper_length_mm":225,"body_mass_g":5000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":51.1,"culmen_depth_mm":16.5,"flipper_length_mm":225,"body_mass_g":5250,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.6,"culmen_depth_mm":16,"flipper_length_mm":225,"body_mass_g":5700,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.5,"culmen_depth_mm":15.9,"flipper_length_mm":225,"body_mass_g":5400,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":39,"culmen_depth_mm":17.5,"flipper_length_mm":186,"body_mass_g":3550,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":39.6,"culmen_depth_mm":17.7,"flipper_length_mm":186,"body_mass_g":3500,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.6,"culmen_depth_mm":19.1,"flipper_length_mm":194,"body_mass_g":3750,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.8,"culmen_depth_mm":14.6,"flipper_length_mm":210,"body_mass_g":4200,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.5,"culmen_depth_mm":13.5,"flipper_length_mm":210,"body_mass_g":4550,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.3,"culmen_depth_mm":13.7,"flipper_length_mm":210,"body_mass_g":4300,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.7,"culmen_depth_mm":14.1,"flipper_length_mm":210,"body_mass_g":4450,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":42,"culmen_depth_mm":13.5,"flipper_length_mm":210,"body_mass_g":4150,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.1,"culmen_depth_mm":14.4,"flipper_length_mm":210,"body_mass_g":4400,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.5,"culmen_depth_mm":13.9,"flipper_length_mm":210,"body_mass_g":4200,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.2,"culmen_depth_mm":14.3,"flipper_length_mm":210,"body_mass_g":4600,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.6,"culmen_depth_mm":14.2,"flipper_length_mm":210,"body_mass_g":4850,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":41.7,"culmen_depth_mm":14.7,"flipper_length_mm":210,"body_mass_g":4700,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.2,"culmen_depth_mm":14.3,"flipper_length_mm":218,"body_mass_g":5700,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.5,"culmen_depth_mm":15,"flipper_length_mm":218,"body_mass_g":4950,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":51.3,"culmen_depth_mm":14.2,"flipper_length_mm":218,"body_mass_g":5300,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50,"culmen_depth_mm":15.2,"flipper_length_mm":218,"body_mass_g":5700,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.4,"culmen_depth_mm":14.4,"flipper_length_mm":218,"body_mass_g":4600,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.8,"culmen_depth_mm":15.7,"flipper_length_mm":226,"body_mass_g":5200,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35.3,"culmen_depth_mm":18.9,"flipper_length_mm":187,"body_mass_g":3800,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":40.5,"culmen_depth_mm":17.9,"flipper_length_mm":187,"body_mass_g":3200,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":34.5,"culmen_depth_mm":18.1,"flipper_length_mm":187,"body_mass_g":2900,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41.3,"culmen_depth_mm":21.1,"flipper_length_mm":195,"body_mass_g":4400,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35.5,"culmen_depth_mm":16.2,"flipper_length_mm":195,"body_mass_g":3350,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41,"culmen_depth_mm":20,"flipper_length_mm":203,"body_mass_g":4725,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.4,"culmen_depth_mm":14.4,"flipper_length_mm":203,"body_mass_g":4625,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.4,"culmen_depth_mm":14.6,"flipper_length_mm":211,"body_mass_g":4800,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.1,"culmen_depth_mm":13.2,"flipper_length_mm":211,"body_mass_g":4500,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.4,"culmen_depth_mm":17.3,"flipper_length_mm":219,"body_mass_g":5250,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.8,"culmen_depth_mm":14.2,"flipper_length_mm":219,"body_mass_g":4700,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":53.4,"culmen_depth_mm":15.8,"flipper_length_mm":219,"body_mass_g":5500,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.5,"culmen_depth_mm":15,"flipper_length_mm":219,"body_mass_g":4850,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.7,"culmen_depth_mm":15.3,"flipper_length_mm":219,"body_mass_g":5200,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.9,"culmen_depth_mm":18.6,"flipper_length_mm":172,"body_mass_g":3150,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.8,"culmen_depth_mm":17.2,"flipper_length_mm":180,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.7,"culmen_depth_mm":18.7,"flipper_length_mm":180,"body_mass_g":3600,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":40.5,"culmen_depth_mm":18.9,"flipper_length_mm":180,"body_mass_g":3950,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41.1,"culmen_depth_mm":19.1,"flipper_length_mm":188,"body_mass_g":4100,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":40.1,"culmen_depth_mm":18.9,"flipper_length_mm":188,"body_mass_g":4300,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":42.7,"culmen_depth_mm":18.3,"flipper_length_mm":196,"body_mass_g":4075,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.5,"culmen_depth_mm":14,"flipper_length_mm":212,"body_mass_g":4875,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.4,"culmen_depth_mm":14.6,"flipper_length_mm":212,"body_mass_g":4725,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.5,"culmen_depth_mm":14.5,"flipper_length_mm":212,"body_mass_g":4750,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.2,"culmen_depth_mm":14.8,"flipper_length_mm":212,"body_mass_g":5200,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.9,"culmen_depth_mm":13.8,"flipper_length_mm":212,"body_mass_g":4750,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.1,"culmen_depth_mm":14.5,"flipper_length_mm":212,"body_mass_g":4625,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.4,"culmen_depth_mm":16.3,"flipper_length_mm":220,"body_mass_g":5400,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50,"culmen_depth_mm":15.3,"flipper_length_mm":220,"body_mass_g":5550,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45,"culmen_depth_mm":15.4,"flipper_length_mm":220,"body_mass_g":5050,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.5,"culmen_depth_mm":15,"flipper_length_mm":220,"body_mass_g":5000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.5,"culmen_depth_mm":14.1,"flipper_length_mm":220,"body_mass_g":5300,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":51.1,"culmen_depth_mm":16.3,"flipper_length_mm":220,"body_mass_g":6000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.1,"culmen_depth_mm":14.8,"flipper_length_mm":220,"body_mass_g":5150,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.5,"culmen_depth_mm":14.2,"flipper_length_mm":220,"body_mass_g":4700,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.1,"culmen_depth_mm":15,"flipper_length_mm":228,"body_mass_g":5500,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":55.9,"culmen_depth_mm":17,"flipper_length_mm":228,"body_mass_g":5600,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.8,"culmen_depth_mm":17.3,"flipper_length_mm":228,"body_mass_g":5600,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":52.2,"culmen_depth_mm":17.1,"flipper_length_mm":228,"body_mass_g":5400,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.1,"culmen_depth_mm":17,"flipper_length_mm":181,"body_mass_g":3175,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":36.5,"culmen_depth_mm":16.6,"flipper_length_mm":181,"body_mass_g":2850,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35.9,"culmen_depth_mm":19.2,"flipper_length_mm":189,"body_mass_g":3800,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":42.2,"culmen_depth_mm":19.5,"flipper_length_mm":197,"body_mass_g":4275,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":43.2,"culmen_depth_mm":19,"flipper_length_mm":197,"body_mass_g":4775,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":42.6,"culmen_depth_mm":13.7,"flipper_length_mm":213,"body_mass_g":4950,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.9,"culmen_depth_mm":13.3,"flipper_length_mm":213,"body_mass_g":5100,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.5,"culmen_depth_mm":14.5,"flipper_length_mm":213,"body_mass_g":4400,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":43.5,"culmen_depth_mm":15.2,"flipper_length_mm":213,"body_mass_g":4650,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.4,"culmen_depth_mm":14.6,"flipper_length_mm":213,"body_mass_g":5850,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.9,"culmen_depth_mm":16.1,"flipper_length_mm":213,"body_mass_g":5400,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.4,"culmen_depth_mm":15.6,"flipper_length_mm":221,"body_mass_g":5000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.2,"culmen_depth_mm":15.2,"flipper_length_mm":221,"body_mass_g":6300,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.2,"culmen_depth_mm":14.9,"flipper_length_mm":221,"body_mass_g":5300,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":52.5,"culmen_depth_mm":15.6,"flipper_length_mm":221,"body_mass_g":5450,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.2,"culmen_depth_mm":15.6,"flipper_length_mm":221,"body_mass_g":5100,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.8,"culmen_depth_mm":15.9,"flipper_length_mm":229,"body_mass_g":5950,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.5,"culmen_depth_mm":16.2,"flipper_length_mm":229,"body_mass_g":5800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.8,"culmen_depth_mm":18.3,"flipper_length_mm":174,"body_mass_g":3400,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":35,"culmen_depth_mm":17.9,"flipper_length_mm":190,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.2,"culmen_depth_mm":20,"flipper_length_mm":190,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.8,"culmen_depth_mm":20,"flipper_length_mm":190,"body_mass_g":4250,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.1,"culmen_depth_mm":16.5,"flipper_length_mm":198,"body_mass_g":3825,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.5,"culmen_depth_mm":13.7,"flipper_length_mm":214,"body_mass_g":4650,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.2,"culmen_depth_mm":14.4,"flipper_length_mm":214,"body_mass_g":4650} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.7,"culmen_depth_mm":13.9,"flipper_length_mm":214,"body_mass_g":4400,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":40.9,"culmen_depth_mm":13.7,"flipper_length_mm":214,"body_mass_g":4650,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":44.5,"culmen_depth_mm":14.7,"flipper_length_mm":214,"body_mass_g":4850,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.2,"culmen_depth_mm":13.7,"flipper_length_mm":214,"body_mass_g":4925,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.8,"culmen_depth_mm":16.2,"flipper_length_mm":222,"body_mass_g":6000,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.9,"culmen_depth_mm":14.6,"flipper_length_mm":222,"body_mass_g":4875,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.7,"culmen_depth_mm":15.1,"flipper_length_mm":222,"body_mass_g":5350,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.4,"culmen_depth_mm":15.7,"flipper_length_mm":222,"body_mass_g":5750,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.3,"culmen_depth_mm":15.3,"flipper_length_mm":222,"body_mass_g":5250,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.5,"culmen_depth_mm":15.9,"flipper_length_mm":222,"body_mass_g":5550,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":52.1,"culmen_depth_mm":17,"flipper_length_mm":230,"body_mass_g":5550,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50,"culmen_depth_mm":16.3,"flipper_length_mm":230,"body_mass_g":5700,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":51.5,"culmen_depth_mm":16.3,"flipper_length_mm":230,"body_mass_g":5500,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":59.6,"culmen_depth_mm":17,"flipper_length_mm":230,"body_mass_g":6050,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":48.6,"culmen_depth_mm":16,"flipper_length_mm":230,"body_mass_g":5800,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":49.8,"culmen_depth_mm":16.8,"flipper_length_mm":230,"body_mass_g":5700,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":55.1,"culmen_depth_mm":16,"flipper_length_mm":230,"body_mass_g":5850,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":37.7,"culmen_depth_mm":16,"flipper_length_mm":183,"body_mass_g":3075,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":40.6,"culmen_depth_mm":18.6,"flipper_length_mm":183,"body_mass_g":3550,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":41.4,"culmen_depth_mm":18.6,"flipper_length_mm":191,"body_mass_g":3700,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":45.6,"culmen_depth_mm":20.3,"flipper_length_mm":191,"body_mass_g":4600,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":39.6,"culmen_depth_mm":20.7,"flipper_length_mm":191,"body_mass_g":3900,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Biscoe","culmen_length_mm":38.6,"culmen_depth_mm":17.2,"flipper_length_mm":199,"body_mass_g":3750,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.1,"culmen_depth_mm":14.5,"flipper_length_mm":207,"body_mass_g":5050,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.8,"culmen_depth_mm":15.4,"flipper_length_mm":215,"body_mass_g":5150,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.8,"culmen_depth_mm":15,"flipper_length_mm":215,"body_mass_g":5650,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.1,"culmen_depth_mm":14.5,"flipper_length_mm":215,"body_mass_g":5000,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.6,"culmen_depth_mm":14.5,"flipper_length_mm":215,"body_mass_g":5400,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.8,"culmen_depth_mm":16.1,"flipper_length_mm":215,"body_mass_g":5500,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.3,"culmen_depth_mm":15.8,"flipper_length_mm":215,"body_mass_g":5050,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.2,"culmen_depth_mm":15.8,"flipper_length_mm":215,"body_mass_g":5300,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.2,"culmen_depth_mm":13.8,"flipper_length_mm":215,"body_mass_g":4750,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":42.9,"culmen_depth_mm":13.1,"flipper_length_mm":215,"body_mass_g":5000,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.1,"culmen_depth_mm":15.1,"flipper_length_mm":215,"body_mass_g":5100,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":46.8,"culmen_depth_mm":14.3,"flipper_length_mm":215,"body_mass_g":4850,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":47.2,"culmen_depth_mm":15.5,"flipper_length_mm":215,"body_mass_g":4975,"sex":"FEMALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":50.7,"culmen_depth_mm":15,"flipper_length_mm":223,"body_mass_g":5550,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":45.2,"culmen_depth_mm":16.4,"flipper_length_mm":223,"body_mass_g":5950,"sex":"MALE"} +{"species":"Gentoo penguin (Pygoscelis papua)","island":"Biscoe","culmen_length_mm":54.3,"culmen_depth_mm":15.7,"flipper_length_mm":231,"body_mass_g":5650,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":40.2,"culmen_depth_mm":17,"flipper_length_mm":176,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":34.4,"culmen_depth_mm":18.4,"flipper_length_mm":184,"body_mass_g":3325,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":37.2,"culmen_depth_mm":19.4,"flipper_length_mm":184,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":36.6,"culmen_depth_mm":17.8,"flipper_length_mm":185,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":35.1,"culmen_depth_mm":19.4,"flipper_length_mm":193,"body_mass_g":4200,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":36.7,"culmen_depth_mm":19.3,"flipper_length_mm":193,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":34.1,"culmen_depth_mm":18.1,"flipper_length_mm":193,"body_mass_g":3475} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":37.8,"culmen_depth_mm":17.1,"flipper_length_mm":186,"body_mass_g":3300} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":35.2,"culmen_depth_mm":15.9,"flipper_length_mm":186,"body_mass_g":3050,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.5,"culmen_depth_mm":17.4,"flipper_length_mm":186,"body_mass_g":3800,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":46,"culmen_depth_mm":21.5,"flipper_length_mm":194,"body_mass_g":4200,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":41.4,"culmen_depth_mm":18.5,"flipper_length_mm":202,"body_mass_g":3875,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":44.1,"culmen_depth_mm":18,"flipper_length_mm":210,"body_mass_g":4000,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":36.2,"culmen_depth_mm":16.1,"flipper_length_mm":187,"body_mass_g":3550,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":36.2,"culmen_depth_mm":17.2,"flipper_length_mm":187,"body_mass_g":3150,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":36.7,"culmen_depth_mm":18.8,"flipper_length_mm":187,"body_mass_g":3800,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":41.5,"culmen_depth_mm":18.3,"flipper_length_mm":195,"body_mass_g":4300,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":42.1,"culmen_depth_mm":19.1,"flipper_length_mm":195,"body_mass_g":4000,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":40.3,"culmen_depth_mm":18,"flipper_length_mm":195,"body_mass_g":3250,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.2,"culmen_depth_mm":19.6,"flipper_length_mm":195,"body_mass_g":4675,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.7,"culmen_depth_mm":19,"flipper_length_mm":195,"body_mass_g":3450,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":42.8,"culmen_depth_mm":18.5,"flipper_length_mm":195,"body_mass_g":4250,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":37.8,"culmen_depth_mm":17.3,"flipper_length_mm":180,"body_mass_g":3700} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.6,"culmen_depth_mm":17,"flipper_length_mm":188,"body_mass_g":2900,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":42.9,"culmen_depth_mm":17.6,"flipper_length_mm":196,"body_mass_g":4700,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.6,"culmen_depth_mm":17.2,"flipper_length_mm":196,"body_mass_g":3550,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.1,"culmen_depth_mm":18.7,"flipper_length_mm":181,"body_mass_g":3750,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.9,"culmen_depth_mm":17.8,"flipper_length_mm":181,"body_mass_g":3625,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":35.7,"culmen_depth_mm":17,"flipper_length_mm":189,"body_mass_g":3350,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":41.1,"culmen_depth_mm":18.6,"flipper_length_mm":189,"body_mass_g":3325,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":34.6,"culmen_depth_mm":17.2,"flipper_length_mm":189,"body_mass_g":3200,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":42.5,"culmen_depth_mm":20.7,"flipper_length_mm":197,"body_mass_g":4500,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":43.1,"culmen_depth_mm":19.2,"flipper_length_mm":197,"body_mass_g":3500,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":45.8,"culmen_depth_mm":18.9,"flipper_length_mm":197,"body_mass_g":4150,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":41.1,"culmen_depth_mm":17.6,"flipper_length_mm":182,"body_mass_g":3200,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":35.9,"culmen_depth_mm":16.6,"flipper_length_mm":190,"body_mass_g":3050,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":33.5,"culmen_depth_mm":19,"flipper_length_mm":190,"body_mass_g":3600,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.7,"culmen_depth_mm":18.4,"flipper_length_mm":190,"body_mass_g":3900,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":35.5,"culmen_depth_mm":17.5,"flipper_length_mm":190,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":42,"culmen_depth_mm":20.2,"flipper_length_mm":190,"body_mass_g":4250} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39.3,"culmen_depth_mm":20.6,"flipper_length_mm":190,"body_mass_g":3650,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.5,"culmen_depth_mm":17.9,"flipper_length_mm":190,"body_mass_g":3325,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":37.7,"culmen_depth_mm":19.8,"flipper_length_mm":198,"body_mass_g":3500,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":34.6,"culmen_depth_mm":21.1,"flipper_length_mm":198,"body_mass_g":4400,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":41.8,"culmen_depth_mm":19.4,"flipper_length_mm":198,"body_mass_g":4450,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.6,"culmen_depth_mm":21.2,"flipper_length_mm":191,"body_mass_g":3800,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":38.8,"culmen_depth_mm":17.6,"flipper_length_mm":191,"body_mass_g":3275,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":40.9,"culmen_depth_mm":16.8,"flipper_length_mm":191,"body_mass_g":3700,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":39,"culmen_depth_mm":17.1,"flipper_length_mm":191,"body_mass_g":3050,"sex":"FEMALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":40.6,"culmen_depth_mm":19,"flipper_length_mm":199,"body_mass_g":4000,"sex":"MALE"} +{"species":"Adelie Penguin (Pygoscelis adeliae)","island":"Torgersen","culmen_length_mm":37.3,"culmen_depth_mm":20.5,"flipper_length_mm":199,"body_mass_g":3775,"sex":"MALE"} diff --git a/tests/data/penguins_schema.json b/tests/data/penguins_schema.json new file mode 100644 index 00000000000..104323547fc --- /dev/null +++ b/tests/data/penguins_schema.json @@ -0,0 +1,37 @@ +[ + { + "mode": "NULLABLE", + "name": "species", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "island", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "culmen_length_mm", + "type": "FLOAT" + }, + { + "mode": "NULLABLE", + "name": "culmen_depth_mm", + "type": "FLOAT" + }, + { + "mode": "NULLABLE", + "name": "flipper_length_mm", + "type": "FLOAT" + }, + { + "mode": "NULLABLE", + "name": "body_mass_g", + "type": "FLOAT" + }, + { + "mode": "NULLABLE", + "name": "sex", + "type": "STRING" + } + ] diff --git a/tests/data/scalars.jsonl b/tests/data/scalars.jsonl new file mode 100644 index 00000000000..172a55ec11a --- /dev/null +++ b/tests/data/scalars.jsonl @@ -0,0 +1,9 @@ +{"bool_col": true, "bytes_col": "SGVsbG8sIFdvcmxkIQ==", "date_col": "2021-07-21", "datetime_col": "2021-07-21 11:39:45", "geography_col": "POINT(-122.0838511 37.3860517)", "int64_col": "123456789", "int64_too": "0", "numeric_col": "1.23456789", "float64_col": "1.25", "rowindex": 0, "rowindex_2": 0, "string_col": "Hello, World!", "time_col": "11:41:43.076160", "timestamp_col": "2021-07-21T17:43:43.945289Z"} +{"bool_col": false, "bytes_col": "44GT44KT44Gr44Gh44Gv", "date_col": "1991-02-03", "datetime_col": "1991-01-02 03:45:06", "geography_col": "POINT(-71.104 42.315)", "int64_col": "-987654321", "int64_too": "1", "numeric_col": "1.23456789", "float64_col": "2.51", "rowindex": 1, "rowindex_2": 1, "string_col": "こんにちは", "time_col": "11:14:34.701606", "timestamp_col": "2021-07-21T17:43:43.945289Z"} +{"bool_col": true, "bytes_col": "wqFIb2xhIE11bmRvIQ==", "date_col": "2023-03-01", "datetime_col": "2023-03-01 10:55:13", "geography_col": "POINT(-0.124474760143016 51.5007826749545)", "int64_col": "314159", "int64_too": "0", "numeric_col": "101.1010101", "float64_col": "2.5e10", "rowindex": 2, "rowindex_2": 2, "string_col": " ¡Hola Mundo! ", "time_col": "23:59:59.999999", "timestamp_col": "2023-03-01T10:55:13.250125Z"} +{"bool_col": null, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": null, "int64_too": "1", "numeric_col": null, "float64_col": null, "rowindex": 3, "rowindex_2": 3, "string_col": null, "time_col": null, "timestamp_col": null} +{"bool_col": false, "bytes_col": "44GT44KT44Gr44Gh44Gv", "date_col": "2021-07-21", "datetime_col": null, "geography_col": null, "int64_col": "-234892", "int64_too": "-2345", "numeric_col": null, "float64_col": null, "rowindex": 4, "rowindex_2": 4, "string_col": "Hello, World!", "time_col": null, "timestamp_col": null} +{"bool_col": false, "bytes_col": "R8O8dGVuIFRhZw==", "date_col": "1980-03-14", "datetime_col": "1980-03-14 15:16:17", "geography_col": null, "int64_col": "55555", "int64_too": "0", "numeric_col": "5.555555", "float64_col": "555.555", "rowindex": 5, "rowindex_2": 5, "string_col": "Güten Tag!", "time_col": "15:16:17.181921", "timestamp_col": "1980-03-14T15:16:17.181921Z"} +{"bool_col": true, "bytes_col": "SGVsbG8JQmlnRnJhbWVzIQc=", "date_col": "2023-05-23", "datetime_col": "2023-05-23 11:37:01", "geography_col": "MULTIPOINT (20 20, 10 40, 40 30, 30 10)", "int64_col": "101202303", "int64_too": "2", "numeric_col": "-10.090807", "float64_col": "-123.456", "rowindex": 6, "rowindex_2": 6, "string_col": "capitalize, This ", "time_col": "01:02:03.456789", "timestamp_col": "2023-05-23T11:42:55.000001Z"} +{"bool_col": true, "bytes_col": null, "date_col": "2038-01-20", "datetime_col": "2038-01-19 03:14:08", "geography_col": null, "int64_col": "-214748367", "int64_too": "2", "numeric_col": "11111111.1", "float64_col": "42.42", "rowindex": 7, "rowindex_2": 7, "string_col": " سلام", "time_col": "12:00:00.000001", "timestamp_col": "2038-01-19T03:14:17.999999Z"} +{"bool_col": false, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": "2", "int64_too": "1", "numeric_col": null, "float64_col": "6.87", "rowindex": 8, "rowindex_2": 8, "string_col": "T", "time_col": null, "timestamp_col": null} diff --git a/tests/data/scalars_schema.json b/tests/data/scalars_schema.json new file mode 100644 index 00000000000..1f5d8cdb650 --- /dev/null +++ b/tests/data/scalars_schema.json @@ -0,0 +1,75 @@ +[ + { + "mode": "NULLABLE", + "name": "bool_col", + "type": "BOOLEAN" + }, + { + "mode": "NULLABLE", + "name": "bytes_col", + "type": "BYTES" + }, + { + "mode": "NULLABLE", + "name": "date_col", + "type": "DATE" + }, + { + "mode": "NULLABLE", + "name": "datetime_col", + "type": "DATETIME" + }, + { + "mode": "NULLABLE", + "name": "geography_col", + "type": "GEOGRAPHY" + }, + { + "mode": "NULLABLE", + "name": "int64_col", + "type": "INTEGER", + "description": "Assumed to be have unique values." + }, + { + "mode": "REQUIRED", + "name": "int64_too", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "numeric_col", + "type": "NUMERIC" + }, + { + "mode": "NULLABLE", + "name": "float64_col", + "type": "FLOAT" + }, + { + "mode": "REQUIRED", + "name": "rowindex", + "type": "INTEGER", + "description": "Assumed to be sequential integers." + }, + { + "mode": "REQUIRED", + "name": "rowindex_2", + "type": "INTEGER", + "description": "Assumed to be a copy of the rowindex column." + }, + { + "mode": "NULLABLE", + "name": "string_col", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "time_col", + "type": "TIME" + }, + { + "mode": "NULLABLE", + "name": "timestamp_col", + "type": "TIMESTAMP" + } + ] diff --git a/tests/data/time_series.jsonl b/tests/data/time_series.jsonl new file mode 100644 index 00000000000..e0f9ca7ae25 --- /dev/null +++ b/tests/data/time_series.jsonl @@ -0,0 +1,366 @@ +{"parsed_date":"2017-07-01 00:00:00 UTC","total_visits":"2048"} +{"parsed_date":"2016-09-07 00:00:00 UTC","total_visits":"2562"} +{"parsed_date":"2016-10-25 00:00:00 UTC","total_visits":"3842"} +{"parsed_date":"2017-04-10 00:00:00 UTC","total_visits":"2563"} +{"parsed_date":"2017-01-09 00:00:00 UTC","total_visits":"2308"} +{"parsed_date":"2017-05-02 00:00:00 UTC","total_visits":"2564"} +{"parsed_date":"2016-11-11 00:00:00 UTC","total_visits":"3588"} +{"parsed_date":"2017-07-30 00:00:00 UTC","total_visits":"1799"} +{"parsed_date":"2017-06-10 00:00:00 UTC","total_visits":"1545"} +{"parsed_date":"2016-08-14 00:00:00 UTC","total_visits":"1801"} +{"parsed_date":"2017-05-14 00:00:00 UTC","total_visits":"1290"} +{"parsed_date":"2017-02-08 00:00:00 UTC","total_visits":"2570"} +{"parsed_date":"2017-06-01 00:00:00 UTC","total_visits":"2826"} +{"parsed_date":"2017-04-23 00:00:00 UTC","total_visits":"1548"} +{"parsed_date":"2016-11-04 00:00:00 UTC","total_visits":"3596"} +{"parsed_date":"2017-02-04 00:00:00 UTC","total_visits":"1549"} +{"parsed_date":"2016-12-09 00:00:00 UTC","total_visits":"2830"} +{"parsed_date":"2016-10-30 00:00:00 UTC","total_visits":"3086"} +{"parsed_date":"2017-03-28 00:00:00 UTC","total_visits":"2577"} +{"parsed_date":"2017-06-11 00:00:00 UTC","total_visits":"1555"} +{"parsed_date":"2016-12-17 00:00:00 UTC","total_visits":"2324"} +{"parsed_date":"2016-09-22 00:00:00 UTC","total_visits":"2581"} +{"parsed_date":"2017-01-29 00:00:00 UTC","total_visits":"1814"} +{"parsed_date":"2017-03-22 00:00:00 UTC","total_visits":"2582"} +{"parsed_date":"2017-02-21 00:00:00 UTC","total_visits":"2582"} +{"parsed_date":"2016-10-14 00:00:00 UTC","total_visits":"2838"} +{"parsed_date":"2017-04-27 00:00:00 UTC","total_visits":"2838"} +{"parsed_date":"2016-10-26 00:00:00 UTC","total_visits":"4375"} +{"parsed_date":"2016-08-22 00:00:00 UTC","total_visits":"2584"} +{"parsed_date":"2016-12-07 00:00:00 UTC","total_visits":"2840"} +{"parsed_date":"2017-01-20 00:00:00 UTC","total_visits":"2074"} +{"parsed_date":"2017-03-07 00:00:00 UTC","total_visits":"2586"} +{"parsed_date":"2017-05-16 00:00:00 UTC","total_visits":"3098"} +{"parsed_date":"2017-05-03 00:00:00 UTC","total_visits":"2588"} +{"parsed_date":"2017-05-01 00:00:00 UTC","total_visits":"2588"} +{"parsed_date":"2016-11-27 00:00:00 UTC","total_visits":"3356"} +{"parsed_date":"2017-04-29 00:00:00 UTC","total_visits":"1566"} +{"parsed_date":"2016-09-18 00:00:00 UTC","total_visits":"1822"} +{"parsed_date":"2017-03-23 00:00:00 UTC","total_visits":"2847"} +{"parsed_date":"2017-03-14 00:00:00 UTC","total_visits":"2338"} +{"parsed_date":"2016-12-21 00:00:00 UTC","total_visits":"2594"} +{"parsed_date":"2016-10-11 00:00:00 UTC","total_visits":"2850"} +{"parsed_date":"2017-01-24 00:00:00 UTC","total_visits":"3618"} +{"parsed_date":"2017-03-05 00:00:00 UTC","total_visits":"1827"} +{"parsed_date":"2017-01-19 00:00:00 UTC","total_visits":"2083"} +{"parsed_date":"2016-08-09 00:00:00 UTC","total_visits":"2851"} +{"parsed_date":"2017-04-08 00:00:00 UTC","total_visits":"1829"} +{"parsed_date":"2017-04-12 00:00:00 UTC","total_visits":"2341"} +{"parsed_date":"2016-09-29 00:00:00 UTC","total_visits":"2597"} +{"parsed_date":"2016-12-20 00:00:00 UTC","total_visits":"3110"} +{"parsed_date":"2017-01-15 00:00:00 UTC","total_visits":"1576"} +{"parsed_date":"2017-04-14 00:00:00 UTC","total_visits":"1834"} +{"parsed_date":"2017-02-28 00:00:00 UTC","total_visits":"2347"} +{"parsed_date":"2016-09-16 00:00:00 UTC","total_visits":"2603"} +{"parsed_date":"2016-10-18 00:00:00 UTC","total_visits":"3628"} +{"parsed_date":"2017-02-24 00:00:00 UTC","total_visits":"2093"} +{"parsed_date":"2017-05-17 00:00:00 UTC","total_visits":"3117"} +{"parsed_date":"2017-06-23 00:00:00 UTC","total_visits":"2095"} +{"parsed_date":"2016-11-12 00:00:00 UTC","total_visits":"3119"} +{"parsed_date":"2016-11-21 00:00:00 UTC","total_visits":"4143"} +{"parsed_date":"2017-02-27 00:00:00 UTC","total_visits":"2352"} +{"parsed_date":"2016-12-26 00:00:00 UTC","total_visits":"1586"} +{"parsed_date":"2017-04-25 00:00:00 UTC","total_visits":"2354"} +{"parsed_date":"2017-03-21 00:00:00 UTC","total_visits":"2611"} +{"parsed_date":"2016-12-22 00:00:00 UTC","total_visits":"2100"} +{"parsed_date":"2016-10-01 00:00:00 UTC","total_visits":"1589"} +{"parsed_date":"2016-09-24 00:00:00 UTC","total_visits":"1845"} +{"parsed_date":"2017-06-21 00:00:00 UTC","total_visits":"2357"} +{"parsed_date":"2016-09-02 00:00:00 UTC","total_visits":"2613"} +{"parsed_date":"2016-08-26 00:00:00 UTC","total_visits":"2359"} +{"parsed_date":"2016-10-12 00:00:00 UTC","total_visits":"2871"} +{"parsed_date":"2017-05-15 00:00:00 UTC","total_visits":"2360"} +{"parsed_date":"2017-06-12 00:00:00 UTC","total_visits":"2361"} +{"parsed_date":"2016-08-16 00:00:00 UTC","total_visits":"2873"} +{"parsed_date":"2017-04-30 00:00:00 UTC","total_visits":"1594"} +{"parsed_date":"2017-04-05 00:00:00 UTC","total_visits":"2619"} +{"parsed_date":"2016-08-12 00:00:00 UTC","total_visits":"2619"} +{"parsed_date":"2016-11-08 00:00:00 UTC","total_visits":"3899"} +{"parsed_date":"2016-08-13 00:00:00 UTC","total_visits":"1596"} +{"parsed_date":"2017-05-09 00:00:00 UTC","total_visits":"2108"} +{"parsed_date":"2017-02-23 00:00:00 UTC","total_visits":"2364"} +{"parsed_date":"2017-07-31 00:00:00 UTC","total_visits":"2620"} +{"parsed_date":"2017-06-25 00:00:00 UTC","total_visits":"1597"} +{"parsed_date":"2017-07-29 00:00:00 UTC","total_visits":"1597"} +{"parsed_date":"2016-09-17 00:00:00 UTC","total_visits":"1853"} +{"parsed_date":"2016-12-27 00:00:00 UTC","total_visits":"1855"} +{"parsed_date":"2017-05-20 00:00:00 UTC","total_visits":"1855"} +{"parsed_date":"2016-10-08 00:00:00 UTC","total_visits":"2114"} +{"parsed_date":"2016-10-27 00:00:00 UTC","total_visits":"4162"} +{"parsed_date":"2017-07-08 00:00:00 UTC","total_visits":"1859"} +{"parsed_date":"2016-08-24 00:00:00 UTC","total_visits":"2627"} +{"parsed_date":"2016-12-23 00:00:00 UTC","total_visits":"1604"} +{"parsed_date":"2017-02-02 00:00:00 UTC","total_visits":"2372"} +{"parsed_date":"2016-09-08 00:00:00 UTC","total_visits":"2628"} +{"parsed_date":"2017-04-02 00:00:00 UTC","total_visits":"1861"} +{"parsed_date":"2017-02-15 00:00:00 UTC","total_visits":"2629"} +{"parsed_date":"2017-07-05 00:00:00 UTC","total_visits":"2885"} +{"parsed_date":"2016-10-17 00:00:00 UTC","total_visits":"3397"} +{"parsed_date":"2017-02-20 00:00:00 UTC","total_visits":"2374"} +{"parsed_date":"2017-03-24 00:00:00 UTC","total_visits":"2374"} +{"parsed_date":"2017-04-20 00:00:00 UTC","total_visits":"2374"} +{"parsed_date":"2016-11-18 00:00:00 UTC","total_visits":"3654"} +{"parsed_date":"2017-07-25 00:00:00 UTC","total_visits":"2631"} +{"parsed_date":"2016-11-13 00:00:00 UTC","total_visits":"3144"} +{"parsed_date":"2017-03-18 00:00:00 UTC","total_visits":"1610"} +{"parsed_date":"2016-08-03 00:00:00 UTC","total_visits":"2890"} +{"parsed_date":"2016-08-19 00:00:00 UTC","total_visits":"2379"} +{"parsed_date":"2017-02-14 00:00:00 UTC","total_visits":"2379"} +{"parsed_date":"2017-07-11 00:00:00 UTC","total_visits":"2635"} +{"parsed_date":"2017-04-22 00:00:00 UTC","total_visits":"1612"} +{"parsed_date":"2016-10-07 00:00:00 UTC","total_visits":"2892"} +{"parsed_date":"2016-09-05 00:00:00 UTC","total_visits":"2125"} +{"parsed_date":"2016-09-23 00:00:00 UTC","total_visits":"2381"} +{"parsed_date":"2016-11-15 00:00:00 UTC","total_visits":"4685"} +{"parsed_date":"2017-01-28 00:00:00 UTC","total_visits":"1614"} +{"parsed_date":"2017-07-14 00:00:00 UTC","total_visits":"2382"} +{"parsed_date":"2017-01-07 00:00:00 UTC","total_visits":"1615"} +{"parsed_date":"2017-04-03 00:00:00 UTC","total_visits":"2383"} +{"parsed_date":"2017-03-20 00:00:00 UTC","total_visits":"2383"} +{"parsed_date":"2016-12-18 00:00:00 UTC","total_visits":"2128"} +{"parsed_date":"2017-03-17 00:00:00 UTC","total_visits":"2129"} +{"parsed_date":"2017-05-23 00:00:00 UTC","total_visits":"2129"} +{"parsed_date":"2016-11-30 00:00:00 UTC","total_visits":"4435"} +{"parsed_date":"2017-01-01 00:00:00 UTC","total_visits":"1364"} +{"parsed_date":"2017-01-02 00:00:00 UTC","total_visits":"1620"} +{"parsed_date":"2016-09-25 00:00:00 UTC","total_visits":"1877"} +{"parsed_date":"2016-08-07 00:00:00 UTC","total_visits":"1622"} +{"parsed_date":"2016-10-09 00:00:00 UTC","total_visits":"2134"} +{"parsed_date":"2017-03-01 00:00:00 UTC","total_visits":"2390"} +{"parsed_date":"2017-01-04 00:00:00 UTC","total_visits":"2390"} +{"parsed_date":"2017-06-06 00:00:00 UTC","total_visits":"2391"} +{"parsed_date":"2017-04-18 00:00:00 UTC","total_visits":"2391"} +{"parsed_date":"2017-04-06 00:00:00 UTC","total_visits":"2647"} +{"parsed_date":"2017-01-30 00:00:00 UTC","total_visits":"2392"} +{"parsed_date":"2016-10-16 00:00:00 UTC","total_visits":"2649"} +{"parsed_date":"2016-08-04 00:00:00 UTC","total_visits":"3161"} +{"parsed_date":"2016-10-21 00:00:00 UTC","total_visits":"3419"} +{"parsed_date":"2016-08-02 00:00:00 UTC","total_visits":"2140"} +{"parsed_date":"2017-03-06 00:00:00 UTC","total_visits":"2396"} +{"parsed_date":"2016-09-13 00:00:00 UTC","total_visits":"2396"} +{"parsed_date":"2016-09-14 00:00:00 UTC","total_visits":"2652"} +{"parsed_date":"2017-04-19 00:00:00 UTC","total_visits":"2397"} +{"parsed_date":"2017-06-19 00:00:00 UTC","total_visits":"2142"} +{"parsed_date":"2016-12-13 00:00:00 UTC","total_visits":"3166"} +{"parsed_date":"2017-06-20 00:00:00 UTC","total_visits":"2143"} +{"parsed_date":"2016-10-10 00:00:00 UTC","total_visits":"2911"} +{"parsed_date":"2017-07-06 00:00:00 UTC","total_visits":"2658"} +{"parsed_date":"2017-01-03 00:00:00 UTC","total_visits":"2403"} +{"parsed_date":"2017-01-08 00:00:00 UTC","total_visits":"1637"} +{"parsed_date":"2017-02-25 00:00:00 UTC","total_visits":"1638"} +{"parsed_date":"2017-05-24 00:00:00 UTC","total_visits":"2406"} +{"parsed_date":"2016-11-22 00:00:00 UTC","total_visits":"3942"} +{"parsed_date":"2017-05-06 00:00:00 UTC","total_visits":"1383"} +{"parsed_date":"2017-07-02 00:00:00 UTC","total_visits":"1895"} +{"parsed_date":"2016-12-01 00:00:00 UTC","total_visits":"4200"} +{"parsed_date":"2017-03-16 00:00:00 UTC","total_visits":"2409"} +{"parsed_date":"2016-12-12 00:00:00 UTC","total_visits":"3433"} +{"parsed_date":"2016-12-25 00:00:00 UTC","total_visits":"1386"} +{"parsed_date":"2017-02-26 00:00:00 UTC","total_visits":"1643"} +{"parsed_date":"2017-04-28 00:00:00 UTC","total_visits":"2411"} +{"parsed_date":"2016-08-11 00:00:00 UTC","total_visits":"2667"} +{"parsed_date":"2017-07-20 00:00:00 UTC","total_visits":"2668"} +{"parsed_date":"2017-05-21 00:00:00 UTC","total_visits":"1645"} +{"parsed_date":"2017-06-17 00:00:00 UTC","total_visits":"1391"} +{"parsed_date":"2016-12-29 00:00:00 UTC","total_visits":"1647"} +{"parsed_date":"2017-07-17 00:00:00 UTC","total_visits":"2671"} +{"parsed_date":"2017-01-16 00:00:00 UTC","total_visits":"1906"} +{"parsed_date":"2017-03-03 00:00:00 UTC","total_visits":"2162"} +{"parsed_date":"2016-11-14 00:00:00 UTC","total_visits":"4466"} +{"parsed_date":"2016-08-30 00:00:00 UTC","total_visits":"2675"} +{"parsed_date":"2016-08-27 00:00:00 UTC","total_visits":"1654"} +{"parsed_date":"2017-02-09 00:00:00 UTC","total_visits":"2678"} +{"parsed_date":"2017-06-03 00:00:00 UTC","total_visits":"1399"} +{"parsed_date":"2017-05-07 00:00:00 UTC","total_visits":"1400"} +{"parsed_date":"2016-11-02 00:00:00 UTC","total_visits":"3960"} +{"parsed_date":"2016-12-15 00:00:00 UTC","total_visits":"2937"} +{"parsed_date":"2017-04-01 00:00:00 UTC","total_visits":"2170"} +{"parsed_date":"2017-07-21 00:00:00 UTC","total_visits":"2427"} +{"parsed_date":"2016-08-06 00:00:00 UTC","total_visits":"1663"} +{"parsed_date":"2016-09-01 00:00:00 UTC","total_visits":"2687"} +{"parsed_date":"2017-06-28 00:00:00 UTC","total_visits":"2687"} +{"parsed_date":"2016-08-20 00:00:00 UTC","total_visits":"1664"} +{"parsed_date":"2017-04-26 00:00:00 UTC","total_visits":"4224"} +{"parsed_date":"2017-07-09 00:00:00 UTC","total_visits":"1921"} +{"parsed_date":"2017-07-28 00:00:00 UTC","total_visits":"2433"} +{"parsed_date":"2016-09-19 00:00:00 UTC","total_visits":"2689"} +{"parsed_date":"2017-07-24 00:00:00 UTC","total_visits":"2436"} +{"parsed_date":"2017-06-13 00:00:00 UTC","total_visits":"2181"} +{"parsed_date":"2016-09-15 00:00:00 UTC","total_visits":"2949"} +{"parsed_date":"2017-02-03 00:00:00 UTC","total_visits":"2182"} +{"parsed_date":"2016-09-10 00:00:00 UTC","total_visits":"1671"} +{"parsed_date":"2017-06-09 00:00:00 UTC","total_visits":"1927"} +{"parsed_date":"2017-01-11 00:00:00 UTC","total_visits":"2185"} +{"parsed_date":"2017-02-19 00:00:00 UTC","total_visits":"2187"} +{"parsed_date":"2017-01-17 00:00:00 UTC","total_visits":"2443"} +{"parsed_date":"2017-05-12 00:00:00 UTC","total_visits":"1932"} +{"parsed_date":"2016-12-16 00:00:00 UTC","total_visits":"2956"} +{"parsed_date":"2017-02-01 00:00:00 UTC","total_visits":"2445"} +{"parsed_date":"2016-11-26 00:00:00 UTC","total_visits":"3213"} +{"parsed_date":"2017-06-02 00:00:00 UTC","total_visits":"2190"} +{"parsed_date":"2016-08-05 00:00:00 UTC","total_visits":"2702"} +{"parsed_date":"2016-11-01 00:00:00 UTC","total_visits":"3728"} +{"parsed_date":"2017-01-05 00:00:00 UTC","total_visits":"2193"} +{"parsed_date":"2017-03-08 00:00:00 UTC","total_visits":"2449"} +{"parsed_date":"2016-08-28 00:00:00 UTC","total_visits":"1682"} +{"parsed_date":"2017-07-04 00:00:00 UTC","total_visits":"1938"} +{"parsed_date":"2017-03-10 00:00:00 UTC","total_visits":"2194"} +{"parsed_date":"2017-07-07 00:00:00 UTC","total_visits":"2450"} +{"parsed_date":"2016-10-29 00:00:00 UTC","total_visits":"2964"} +{"parsed_date":"2016-10-13 00:00:00 UTC","total_visits":"2964"} +{"parsed_date":"2016-12-04 00:00:00 UTC","total_visits":"3220"} +{"parsed_date":"2017-01-21 00:00:00 UTC","total_visits":"1685"} +{"parsed_date":"2017-06-29 00:00:00 UTC","total_visits":"2709"} +{"parsed_date":"2016-08-29 00:00:00 UTC","total_visits":"2454"} +{"parsed_date":"2016-12-19 00:00:00 UTC","total_visits":"3222"} +{"parsed_date":"2017-05-30 00:00:00 UTC","total_visits":"2199"} +{"parsed_date":"2017-02-10 00:00:00 UTC","total_visits":"2199"} +{"parsed_date":"2016-08-31 00:00:00 UTC","total_visits":"3223"} +{"parsed_date":"2017-06-18 00:00:00 UTC","total_visits":"1432"} +{"parsed_date":"2017-01-12 00:00:00 UTC","total_visits":"2203"} +{"parsed_date":"2017-05-18 00:00:00 UTC","total_visits":"2715"} +{"parsed_date":"2016-10-23 00:00:00 UTC","total_visits":"2971"} +{"parsed_date":"2016-09-04 00:00:00 UTC","total_visits":"1692"} +{"parsed_date":"2016-12-10 00:00:00 UTC","total_visits":"2207"} +{"parsed_date":"2016-12-11 00:00:00 UTC","total_visits":"2208"} +{"parsed_date":"2017-04-11 00:00:00 UTC","total_visits":"2464"} +{"parsed_date":"2016-09-21 00:00:00 UTC","total_visits":"2720"} +{"parsed_date":"2016-11-06 00:00:00 UTC","total_visits":"3232"} +{"parsed_date":"2017-01-26 00:00:00 UTC","total_visits":"2209"} +{"parsed_date":"2016-09-12 00:00:00 UTC","total_visits":"2465"} +{"parsed_date":"2017-04-21 00:00:00 UTC","total_visits":"2210"} +{"parsed_date":"2017-01-06 00:00:00 UTC","total_visits":"2210"} +{"parsed_date":"2017-04-04 00:00:00 UTC","total_visits":"2978"} +{"parsed_date":"2017-01-22 00:00:00 UTC","total_visits":"1700"} +{"parsed_date":"2017-07-26 00:00:00 UTC","total_visits":"2725"} +{"parsed_date":"2016-08-18 00:00:00 UTC","total_visits":"2725"} +{"parsed_date":"2016-09-27 00:00:00 UTC","total_visits":"2727"} +{"parsed_date":"2016-12-02 00:00:00 UTC","total_visits":"3751"} +{"parsed_date":"2017-05-05 00:00:00 UTC","total_visits":"1960"} +{"parsed_date":"2016-11-19 00:00:00 UTC","total_visits":"2984"} +{"parsed_date":"2016-11-09 00:00:00 UTC","total_visits":"3752"} +{"parsed_date":"2016-12-05 00:00:00 UTC","total_visits":"4265"} +{"parsed_date":"2017-05-11 00:00:00 UTC","total_visits":"2218"} +{"parsed_date":"2017-01-25 00:00:00 UTC","total_visits":"2986"} +{"parsed_date":"2017-03-11 00:00:00 UTC","total_visits":"1707"} +{"parsed_date":"2017-03-30 00:00:00 UTC","total_visits":"2731"} +{"parsed_date":"2016-10-20 00:00:00 UTC","total_visits":"3755"} +{"parsed_date":"2017-02-07 00:00:00 UTC","total_visits":"2476"} +{"parsed_date":"2017-02-22 00:00:00 UTC","total_visits":"2477"} +{"parsed_date":"2017-07-23 00:00:00 UTC","total_visits":"1966"} +{"parsed_date":"2016-11-03 00:00:00 UTC","total_visits":"4014"} +{"parsed_date":"2016-08-01 00:00:00 UTC","total_visits":"1711"} +{"parsed_date":"2017-01-13 00:00:00 UTC","total_visits":"1967"} +{"parsed_date":"2017-05-19 00:00:00 UTC","total_visits":"2223"} +{"parsed_date":"2016-11-20 00:00:00 UTC","total_visits":"3247"} +{"parsed_date":"2016-11-25 00:00:00 UTC","total_visits":"3759"} +{"parsed_date":"2017-03-25 00:00:00 UTC","total_visits":"1712"} +{"parsed_date":"2017-01-27 00:00:00 UTC","total_visits":"1969"} +{"parsed_date":"2017-06-26 00:00:00 UTC","total_visits":"2226"} +{"parsed_date":"2017-05-25 00:00:00 UTC","total_visits":"2228"} +{"parsed_date":"2017-01-31 00:00:00 UTC","total_visits":"2229"} +{"parsed_date":"2017-07-13 00:00:00 UTC","total_visits":"2741"} +{"parsed_date":"2017-03-15 00:00:00 UTC","total_visits":"2486"} +{"parsed_date":"2017-05-28 00:00:00 UTC","total_visits":"1463"} +{"parsed_date":"2017-03-09 00:00:00 UTC","total_visits":"2231"} +{"parsed_date":"2017-07-15 00:00:00 UTC","total_visits":"1721"} +{"parsed_date":"2016-11-24 00:00:00 UTC","total_visits":"3770"} +{"parsed_date":"2016-10-05 00:00:00 UTC","total_visits":"3770"} +{"parsed_date":"2016-12-31 00:00:00 UTC","total_visits":"1211"} +{"parsed_date":"2016-10-02 00:00:00 UTC","total_visits":"1724"} +{"parsed_date":"2017-07-22 00:00:00 UTC","total_visits":"1724"} +{"parsed_date":"2016-09-11 00:00:00 UTC","total_visits":"1725"} +{"parsed_date":"2017-06-15 00:00:00 UTC","total_visits":"2237"} +{"parsed_date":"2017-06-05 00:00:00 UTC","total_visits":"2493"} +{"parsed_date":"2017-02-06 00:00:00 UTC","total_visits":"2238"} +{"parsed_date":"2016-10-15 00:00:00 UTC","total_visits":"2495"} +{"parsed_date":"2016-08-21 00:00:00 UTC","total_visits":"1730"} +{"parsed_date":"2016-08-23 00:00:00 UTC","total_visits":"2754"} +{"parsed_date":"2017-06-30 00:00:00 UTC","total_visits":"2499"} +{"parsed_date":"2017-01-18 00:00:00 UTC","total_visits":"2245"} +{"parsed_date":"2016-08-10 00:00:00 UTC","total_visits":"2757"} +{"parsed_date":"2016-12-08 00:00:00 UTC","total_visits":"3013"} +{"parsed_date":"2016-11-28 00:00:00 UTC","total_visits":"4807"} +{"parsed_date":"2017-05-22 00:00:00 UTC","total_visits":"2248"} +{"parsed_date":"2016-09-20 00:00:00 UTC","total_visits":"2760"} +{"parsed_date":"2016-10-06 00:00:00 UTC","total_visits":"3016"} +{"parsed_date":"2016-09-06 00:00:00 UTC","total_visits":"2508"} +{"parsed_date":"2016-09-03 00:00:00 UTC","total_visits":"1741"} +{"parsed_date":"2016-12-06 00:00:00 UTC","total_visits":"3021"} +{"parsed_date":"2016-12-24 00:00:00 UTC","total_visits":"1231"} +{"parsed_date":"2016-10-28 00:00:00 UTC","total_visits":"3791"} +{"parsed_date":"2016-12-30 00:00:00 UTC","total_visits":"1232"} +{"parsed_date":"2017-05-29 00:00:00 UTC","total_visits":"1745"} +{"parsed_date":"2017-07-10 00:00:00 UTC","total_visits":"2769"} +{"parsed_date":"2017-06-22 00:00:00 UTC","total_visits":"2258"} +{"parsed_date":"2017-07-19 00:00:00 UTC","total_visits":"2514"} +{"parsed_date":"2016-10-03 00:00:00 UTC","total_visits":"2514"} +{"parsed_date":"2017-06-14 00:00:00 UTC","total_visits":"2517"} +{"parsed_date":"2016-10-22 00:00:00 UTC","total_visits":"3029"} +{"parsed_date":"2017-01-23 00:00:00 UTC","total_visits":"2262"} +{"parsed_date":"2017-04-24 00:00:00 UTC","total_visits":"2263"} +{"parsed_date":"2016-11-10 00:00:00 UTC","total_visits":"4055"} +{"parsed_date":"2016-09-26 00:00:00 UTC","total_visits":"2776"} +{"parsed_date":"2016-10-19 00:00:00 UTC","total_visits":"3544"} +{"parsed_date":"2017-03-04 00:00:00 UTC","total_visits":"1753"} +{"parsed_date":"2017-05-26 00:00:00 UTC","total_visits":"2009"} +{"parsed_date":"2017-02-13 00:00:00 UTC","total_visits":"2266"} +{"parsed_date":"2017-02-18 00:00:00 UTC","total_visits":"1755"} +{"parsed_date":"2017-03-02 00:00:00 UTC","total_visits":"2267"} +{"parsed_date":"2017-03-31 00:00:00 UTC","total_visits":"2268"} +{"parsed_date":"2017-01-10 00:00:00 UTC","total_visits":"2268"} +{"parsed_date":"2017-03-29 00:00:00 UTC","total_visits":"2525"} +{"parsed_date":"2017-03-27 00:00:00 UTC","total_visits":"2525"} +{"parsed_date":"2016-11-23 00:00:00 UTC","total_visits":"3805"} +{"parsed_date":"2017-05-27 00:00:00 UTC","total_visits":"1502"} +{"parsed_date":"2016-10-24 00:00:00 UTC","total_visits":"4063"} +{"parsed_date":"2016-12-14 00:00:00 UTC","total_visits":"3040"} +{"parsed_date":"2017-02-11 00:00:00 UTC","total_visits":"1761"} +{"parsed_date":"2017-07-27 00:00:00 UTC","total_visits":"2529"} +{"parsed_date":"2017-02-17 00:00:00 UTC","total_visits":"2785"} +{"parsed_date":"2017-04-15 00:00:00 UTC","total_visits":"1506"} +{"parsed_date":"2016-11-05 00:00:00 UTC","total_visits":"3042"} +{"parsed_date":"2016-10-04 00:00:00 UTC","total_visits":"4322"} +{"parsed_date":"2017-05-13 00:00:00 UTC","total_visits":"1251"} +{"parsed_date":"2017-04-16 00:00:00 UTC","total_visits":"1507"} +{"parsed_date":"2016-12-28 00:00:00 UTC","total_visits":"1763"} +{"parsed_date":"2016-08-15 00:00:00 UTC","total_visits":"3043"} +{"parsed_date":"2016-12-03 00:00:00 UTC","total_visits":"3044"} +{"parsed_date":"2017-06-27 00:00:00 UTC","total_visits":"2789"} +{"parsed_date":"2017-06-24 00:00:00 UTC","total_visits":"1510"} +{"parsed_date":"2017-07-16 00:00:00 UTC","total_visits":"1766"} +{"parsed_date":"2017-04-09 00:00:00 UTC","total_visits":"1766"} +{"parsed_date":"2017-06-07 00:00:00 UTC","total_visits":"2279"} +{"parsed_date":"2017-04-17 00:00:00 UTC","total_visits":"2279"} +{"parsed_date":"2016-09-28 00:00:00 UTC","total_visits":"2535"} +{"parsed_date":"2017-03-26 00:00:00 UTC","total_visits":"1768"} +{"parsed_date":"2017-05-10 00:00:00 UTC","total_visits":"2024"} +{"parsed_date":"2017-06-08 00:00:00 UTC","total_visits":"2280"} +{"parsed_date":"2017-05-08 00:00:00 UTC","total_visits":"2025"} +{"parsed_date":"2017-03-13 00:00:00 UTC","total_visits":"2537"} +{"parsed_date":"2016-11-17 00:00:00 UTC","total_visits":"4074"} +{"parsed_date":"2016-08-25 00:00:00 UTC","total_visits":"2539"} +{"parsed_date":"2017-02-16 00:00:00 UTC","total_visits":"2539"} +{"parsed_date":"2017-06-16 00:00:00 UTC","total_visits":"2028"} +{"parsed_date":"2016-11-16 00:00:00 UTC","total_visits":"4334"} +{"parsed_date":"2016-08-17 00:00:00 UTC","total_visits":"2799"} +{"parsed_date":"2017-03-19 00:00:00 UTC","total_visits":"1776"} +{"parsed_date":"2016-11-29 00:00:00 UTC","total_visits":"4337"} +{"parsed_date":"2017-02-05 00:00:00 UTC","total_visits":"1522"} +{"parsed_date":"2016-10-31 00:00:00 UTC","total_visits":"3827"} +{"parsed_date":"2017-05-31 00:00:00 UTC","total_visits":"2292"} +{"parsed_date":"2017-07-18 00:00:00 UTC","total_visits":"2804"} +{"parsed_date":"2017-03-12 00:00:00 UTC","total_visits":"1781"} +{"parsed_date":"2016-09-09 00:00:00 UTC","total_visits":"2549"} +{"parsed_date":"2017-01-14 00:00:00 UTC","total_visits":"1526"} +{"parsed_date":"2017-05-04 00:00:00 UTC","total_visits":"2806"} +{"parsed_date":"2016-11-07 00:00:00 UTC","total_visits":"3832"} +{"parsed_date":"2017-04-07 00:00:00 UTC","total_visits":"2297"} +{"parsed_date":"2017-07-12 00:00:00 UTC","total_visits":"2554"} +{"parsed_date":"2017-04-13 00:00:00 UTC","total_visits":"2300"} +{"parsed_date":"2017-08-01 00:00:00 UTC","total_visits":"2556"} +{"parsed_date":"2017-06-04 00:00:00 UTC","total_visits":"1534"} +{"parsed_date":"2017-02-12 00:00:00 UTC","total_visits":"1790"} +{"parsed_date":"2017-07-03 00:00:00 UTC","total_visits":"2046"} +{"parsed_date":"2016-09-30 00:00:00 UTC","total_visits":"2303"} +{"parsed_date":"2016-08-08 00:00:00 UTC","total_visits":"2815"} diff --git a/tests/data/time_series_schema.json b/tests/data/time_series_schema.json new file mode 100644 index 00000000000..857595b9e64 --- /dev/null +++ b/tests/data/time_series_schema.json @@ -0,0 +1,12 @@ +[ + { + "mode": "NULLABLE", + "name": "parsed_date", + "type": "TIMESTAMP" + }, + { + "mode": "NULLABLE", + "name": "total_visits", + "type": "INTEGER" + } + ] diff --git a/tests/system/__init__.py b/tests/system/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/system/__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/tests/system/conftest.py b/tests/system/conftest.py new file mode 100644 index 00000000000..03638698966 --- /dev/null +++ b/tests/system/conftest.py @@ -0,0 +1,667 @@ +# 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. + +from datetime import datetime +import hashlib +import logging +import pathlib +import typing +from typing import Dict, Optional + +import google.cloud.bigquery as bigquery +import google.cloud.bigquery_connection_v1 as bigquery_connection_v1 +import google.cloud.exceptions +import google.cloud.storage as storage # type: ignore +import ibis.backends.base +import pandas as pd +import pytest +import pytz +import test_utils.prefixer + +import bigframes +from tests.system.utils import convert_pandas_dtypes + +CURRENT_DIR = pathlib.Path(__file__).parent +DATA_DIR = CURRENT_DIR.parent / "data" +PERMANENT_DATASET = "bigframes_testing" +PERMANENT_DATASET_TOKYO = "bigframes_testing_tokyo" +TOKYO_LOCATION = "asia-northeast1" +prefixer = test_utils.prefixer.Prefixer("bigframes", "tests/system") + + +def _hash_digest_file(hasher, filepath): + with open(filepath, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + hasher.update(chunk) + + +@pytest.fixture(scope="session") +def tokyo_location() -> str: + return TOKYO_LOCATION + + +@pytest.fixture(scope="session") +def gcs_client() -> storage.Client: + # TODO(swast): Ensure same credentials and project are used as in the rest + # of our tests. + return storage.Client() + + +@pytest.fixture(scope="session") +def gcs_folder(gcs_client: storage.Client): + # TODO(swast): Allow bucket name from environment variable for testing by + # non-Googlers. + bucket = "bigframes-dev-testing" + prefix = prefixer.create_prefix() + path = f"gs://{bucket}/{prefix}/" + yield path + for blob in gcs_client.list_blobs(bucket, prefix=prefix): + blob = typing.cast(storage.Blob, blob) + blob.delete() + + +@pytest.fixture(scope="session") +def bigquery_client(session: bigframes.Session) -> bigquery.Client: + return session.bqclient + + +@pytest.fixture(scope="session") +def bigquery_client_tokyo(session_tokyo: bigframes.Session) -> bigquery.Client: + return session_tokyo.bqclient + + +@pytest.fixture(scope="session") +def ibis_client(session: bigframes.Session) -> ibis.backends.base.BaseBackend: + return session.ibis_client + + +@pytest.fixture(scope="session") +def bigqueryconnection_client( + session: bigframes.Session, +) -> bigquery_connection_v1.ConnectionServiceClient: + return session.bqconnectionclient + + +@pytest.fixture(scope="session") +def session() -> bigframes.Session: + return bigframes.Session() + + +@pytest.fixture(scope="session") +def session_tokyo(tokyo_location: str) -> bigframes.Session: + context = bigframes.BigQueryOptions( + location=tokyo_location, + use_regional_endpoints=True, + ) + return bigframes.Session(context=context) + + +@pytest.fixture(scope="session", autouse=True) +def cleanup_datasets(bigquery_client: bigquery.Client) -> None: + """Cleanup any datasets that were created but not cleaned up.""" + 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 dataset_id(bigquery_client: bigquery.Client): + """Create (and cleanup) a temporary dataset.""" + project_id = bigquery_client.project + dataset_id = f"{project_id}.{prefixer.create_prefix()}_dataset_id" + dataset = bigquery.Dataset(dataset_id) + bigquery_client.create_dataset(dataset) + yield dataset_id + bigquery_client.delete_dataset(dataset, delete_contents=True) + + +@pytest.fixture(scope="session") +def dataset_id_permanent(bigquery_client: bigquery.Client) -> str: + """Create a dataset if it doesn't exist.""" + project_id = bigquery_client.project + dataset_id = f"{project_id}.{PERMANENT_DATASET}" + dataset = bigquery.Dataset(dataset_id) + bigquery_client.create_dataset(dataset, exists_ok=True) + return dataset_id + + +@pytest.fixture(scope="session") +def dataset_id_permanent_tokyo( + bigquery_client_tokyo: bigquery.Client, tokyo_location: str +) -> str: + """Create a dataset in Tokyo if it doesn't exist.""" + project_id = bigquery_client_tokyo.project + dataset_id = f"{project_id}.{PERMANENT_DATASET_TOKYO}" + dataset = bigquery.Dataset(dataset_id) + dataset.location = tokyo_location + dataset = bigquery_client_tokyo.create_dataset(dataset, exists_ok=True) + assert dataset.location == tokyo_location + return dataset_id + + +@pytest.fixture(scope="session") +def scalars_schema(bigquery_client: bigquery.Client): + # TODO(swast): Add missing scalar data types such as BIGNUMERIC. + # See also: https://github.com/ibis-project/ibis-bigquery/pull/67 + schema = bigquery_client.schema_from_json(DATA_DIR / "scalars_schema.json") + return tuple(schema) + + +def load_test_data( + table_id: str, + bigquery_client: bigquery.Client, + schema_filename: str, + data_filename: str, + location: Optional[str], +) -> bigquery.LoadJob: + """Create a temporary table with test data""" + job_config = bigquery.LoadJobConfig() + job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON + job_config.schema = tuple( + bigquery_client.schema_from_json(DATA_DIR / schema_filename) + ) + job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE + with open(DATA_DIR / data_filename, "rb") as input_file: + # TODO(swast): Location is allowed to be None in BigQuery Client. + # Can remove after + # https://github.com/googleapis/python-bigquery/pull/1554 is released. + location = "US" if location is None else location + job = bigquery_client.load_table_from_file( + input_file, + table_id, + job_config=job_config, + location=location, + ) + # No cleanup necessary, as the surrounding dataset will delete contents. + return typing.cast(bigquery.LoadJob, job.result()) + + +def load_test_data_tables( + session: bigframes.Session, dataset_id_permanent: str +) -> Dict[str, str]: + """Returns cached references to the test data tables in BigQuery. If no matching table is found + for the hash of the data and schema, the table will be uploaded.""" + existing_table_ids = [ + table.table_id for table in session.bqclient.list_tables(dataset_id_permanent) + ] + table_mapping: Dict[str, str] = {} + for table_name, schema_filename, data_filename in [ + ("scalars", "scalars_schema.json", "scalars.jsonl"), + ("scalars_too", "scalars_schema.json", "scalars.jsonl"), + ("penguins", "penguins_schema.json", "penguins.jsonl"), + ("time_series", "time_series_schema.json", "time_series.jsonl"), + ]: + test_data_hash = hashlib.md5() + _hash_digest_file(test_data_hash, DATA_DIR / schema_filename) + _hash_digest_file(test_data_hash, DATA_DIR / data_filename) + test_data_hash.update(table_name.encode()) + target_table_id = f"{table_name}_{test_data_hash.hexdigest()}" + target_table_id_full = f"{dataset_id_permanent}.{target_table_id}" + if target_table_id not in existing_table_ids: + # matching table wasn't found in the permanent dataset - we need to upload it + logging.info( + f"Test data table {table_name} was not found in the permanent dataset, regenerating it..." + ) + load_test_data( + target_table_id_full, + session.bqclient, + schema_filename, + data_filename, + location=session._location, + ) + + table_mapping[table_name] = target_table_id_full + + return table_mapping + + +@pytest.fixture(scope="session") +def test_data_tables( + session: bigframes.Session, dataset_id_permanent: str +) -> Dict[str, str]: + return load_test_data_tables(session, dataset_id_permanent) + + +@pytest.fixture(scope="session") +def test_data_tables_tokyo( + session_tokyo: bigframes.Session, dataset_id_permanent_tokyo: str +) -> Dict[str, str]: + return load_test_data_tables(session_tokyo, dataset_id_permanent_tokyo) + + +@pytest.fixture(scope="session") +def scalars_table_id(test_data_tables) -> str: + return test_data_tables["scalars"] + + +@pytest.fixture(scope="session") +def scalars_table_id_2(test_data_tables) -> str: + return test_data_tables["scalars_too"] + + +@pytest.fixture(scope="session") +def scalars_table_tokyo(test_data_tables_tokyo) -> str: + return test_data_tables_tokyo["scalars"] + + +@pytest.fixture(scope="session") +def penguins_table_id(test_data_tables) -> str: + return test_data_tables["penguins"] + + +@pytest.fixture(scope="session") +def time_series_table_id(test_data_tables) -> str: + return test_data_tables["time_series"] + + +@pytest.fixture(scope="session") +def scalars_df_default_index( + scalars_df_index: bigframes.dataframe.DataFrame, + scalars_pandas_df_default_index: pd.DataFrame, +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + df = scalars_df_index.reset_index(drop=False) + # Ensure the order of the columns is the same. + df = typing.cast( + bigframes.dataframe.DataFrame, df[scalars_pandas_df_default_index.columns] + ) + return df + + +@pytest.fixture(scope="session") +def scalars_df_index( + scalars_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return session.read_gbq(scalars_table_id, index_col="rowindex") + + +@pytest.fixture(scope="session") +def scalars_df_2_default_index( + scalars_df_2_index: bigframes.dataframe.DataFrame, +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return scalars_df_2_index.reset_index(drop=False) + + +@pytest.fixture(scope="session") +def scalars_df_2_index( + scalars_table_id_2: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return session.read_gbq(scalars_table_id_2, index_col="rowindex") + + +@pytest.fixture(scope="session") +def scalars_pandas_df_default_index() -> pd.DataFrame: + """pd.DataFrame pointing at test data.""" + + df = pd.read_json( + DATA_DIR / "scalars.jsonl", + lines=True, + ) + convert_pandas_dtypes(df, bytes_col=True) + + df = df.set_index("rowindex", drop=False) + df.index.name = None + return df + + +@pytest.fixture(scope="session") +def scalars_pandas_df_index( + scalars_pandas_df_default_index: pd.DataFrame, +) -> pd.DataFrame: + """pd.DataFrame pointing at test data.""" + return scalars_pandas_df_default_index.set_index("rowindex").sort_index() + + +@pytest.fixture(scope="session") +def scalars_pandas_df_multi_index( + scalars_pandas_df_default_index: pd.DataFrame, +) -> pd.DataFrame: + """pd.DataFrame pointing at test data.""" + return scalars_pandas_df_default_index.set_index( + ["rowindex", "timestamp_col"] + ).sort_index() + + +@pytest.fixture(scope="session") +def scalars_dfs( + scalars_df_index, + scalars_pandas_df_index, +): + return scalars_df_index, scalars_pandas_df_index + + +@pytest.fixture(scope="session") +def penguins_df_default_index( + penguins_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return session.read_gbq(penguins_table_id) + + +@pytest.fixture(scope="session") +def time_series_df_default_index( + time_series_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return session.read_gbq(time_series_table_id) + + +@pytest.fixture(scope="session") +def new_time_series_pandas_df(): + """Additional data matching the time series dataset. The values are dummy ones used to basically check the prediction scores.""" + utc = pytz.utc + return pd.DataFrame( + { + "parsed_date": [ + datetime(2017, 8, 2, tzinfo=utc), + datetime(2017, 8, 3, tzinfo=utc), + datetime(2017, 8, 4, tzinfo=utc), + ], + "total_visits": [2500, 2500, 2500], + } + ) + + +@pytest.fixture(scope="session") +def new_time_series_df(session, new_time_series_pandas_df): + return session.read_pandas(new_time_series_pandas_df) + + +@pytest.fixture(scope="session") +def penguins_pandas_df_default_index() -> pd.DataFrame: + """Consistently ordered pandas dataframe for penguins test data""" + df = pd.read_json( + f"{DATA_DIR}/penguins.jsonl", + lines=True, + dtype={ + "species": pd.StringDtype(storage="pyarrow"), + "island": pd.StringDtype(storage="pyarrow"), + "culmen_length_mm": pd.Float64Dtype(), + "culmen_depth_mm": pd.Float64Dtype(), + "flipper_length_mm": pd.Float64Dtype(), + "sex": pd.StringDtype(storage="pyarrow"), + "body_mass_g": pd.Float64Dtype(), + }, + ) + df.index = df.index.astype("Int64") + return df + + +@pytest.fixture(scope="session") +def new_penguins_pandas_df(): + """Additional data matching the penguins dataset, with a new index""" + return pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "body_mass_g": [3750.0, 5200.0, 3325.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + + +@pytest.fixture(scope="session") +def new_penguins_df(session, new_penguins_pandas_df): + return session.read_pandas(new_penguins_pandas_df) + + +@pytest.fixture(scope="session") +def penguins_linear_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='linear_reg', + input_label_cols=['body_mass_g'], + data_split_method='NO_SPLIT' +) AS +SELECT + * +FROM + `{penguins_table_id}` +WHERE + body_mass_g IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_linear_reg_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_linear_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def penguins_logistic_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='logistic_reg', + input_label_cols=['sex'], + data_split_method='NO_SPLIT' +) AS SELECT + * +FROM `{penguins_table_id}` +WHERE + sex IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_logistic_reg_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_logistic_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def penguins_xgbregressor_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='BOOSTED_TREE_REGRESSOR', + num_parallel_tree=1, + booster_type='GBTREE', + early_stop=True, + data_split_method='NO_SPLIT', + subsample=1.0, + input_label_cols=['body_mass_g'] +) AS SELECT + * +FROM `{penguins_table_id}` +WHERE + body_mass_g IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_xgbregressor_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_xgbregressor_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def time_series_arima_plus_model_name( + session: bigframes.Session, dataset_id_permanent, time_series_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='ARIMA_PLUS', + time_series_timestamp_col = 'parsed_date', + time_series_data_col = 'total_visits' +) AS SELECT + * +FROM `{time_series_table_id}`""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.time_series_arima_plus_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "time_series_arima_plus_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def penguins_xgbclassifier_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type="BOOSTED_TREE_CLASSIFIER", + num_parallel_tree=1, + booster_type='GBTREE', + early_stop=True, + data_split_method='NO_SPLIT', + subsample=1.0, + input_label_cols=['sex'] +) AS SELECT + * +FROM `{penguins_table_id}` +WHERE + sex IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_classifier_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_classifier_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def penguins_randomforest_regressor_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='RANDOM_FOREST_REGRESSOR', + num_parallel_tree=100, + early_stop=True, + data_split_method='NO_SPLIT', + subsample=0.8, + input_label_cols=['body_mass_g'] +) AS SELECT + * +FROM `{penguins_table_id}` +WHERE + body_mass_g IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_randomforest_regressor_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_randomforest_regressor_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + +@pytest.fixture(scope="session") +def penguins_randomforest_classifier_model_name( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type="RANDOM_FOREST_CLASSIFIER", + num_parallel_tree=100, + early_stop=True, + data_split_method='NO_SPLIT', + subsample=0.8, + input_label_cols=['sex'] +) AS SELECT + * +FROM `{penguins_table_id}` +WHERE + sex IS NOT NULL""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_randomforest_classifier_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_randomforest_classifier_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name diff --git a/tests/system/large/__init__.py b/tests/system/large/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/system/large/__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/tests/system/large/ml/test_cluster.py b/tests/system/large/ml/test_cluster.py new file mode 100644 index 00000000000..bb64b19076f --- /dev/null +++ b/tests/system/large/ml/test_cluster.py @@ -0,0 +1,103 @@ +# 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 pandas + +import bigframes.ml.cluster +from tests.system.utils import assert_pandas_df_equal_ignore_ordering + + +def test_cluster_configure_fit_predict(session, penguins_df_default_index, dataset_id): + model = bigframes.ml.cluster.KMeans(n_clusters=3) + + df = penguins_df_default_index.dropna()[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + + # TODO(swast): How should we handle the default index? Currently, we get: + # "Column bigframes_index_0_z is not found in the input data to the + # EVALUATE function." + df = df.reset_index(drop=True) + + model.fit(df) + + pd_new_penguins = pandas.DataFrame.from_dict( + { + "test1": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Dream", + "culmen_length_mm": 37.5, + "culmen_depth_mm": 18.5, + "flipper_length_mm": 199, + "body_mass_g": 4475, + "sex": "MALE", + }, + "test2": { + "species": "Chinstrap penguin (Pygoscelis antarctica)", + "island": "Dream", + "culmen_length_mm": 55.8, + "culmen_depth_mm": 19.8, + "flipper_length_mm": 207, + "body_mass_g": 4000, + "sex": "MALE", + }, + "test3": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Biscoe", + "culmen_length_mm": 39.7, + "culmen_depth_mm": 18.9, + "flipper_length_mm": 184, + "body_mass_g": 3550, + "sex": "MALE", + }, + "test4": { + "species": "Gentoo penguin (Pygoscelis papua)", + "island": "Biscoe", + "culmen_length_mm": 43.8, + "culmen_depth_mm": 13.9, + "flipper_length_mm": 208, + "body_mass_g": 4300, + "sex": "FEMALE", + }, + }, + orient="index", + ) + pd_new_penguins.index.name = "observation" + + new_penguins = session.read_pandas(pd_new_penguins) + result = model.predict(new_penguins).compute() + expected = pandas.DataFrame( + {"CENTROID_ID": [2, 3, 1, 2]}, + dtype="Int64", + index=pandas.Index( + ["test1", "test2", "test3", "test4"], dtype="string[pyarrow]" + ), + ) + expected.index.name = "observation" + assert_pandas_df_equal_ignore_ordering(result, expected) + + # save, load, check n_clusters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_cluster_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_cluster_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.n_clusters == 3 diff --git a/tests/system/large/ml/test_compose.py b/tests/system/large/ml/test_compose.py new file mode 100644 index 00000000000..52f347151ed --- /dev/null +++ b/tests/system/large/ml/test_compose.py @@ -0,0 +1,71 @@ +# 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 pandas + +import bigframes.ml.cluster +import bigframes.ml.compose +import bigframes.ml.linear_model +import bigframes.ml.pipeline +import bigframes.ml.preprocessing + + +def test_columntransformer_standalone_fit_transform( + penguins_df_default_index, new_penguins_df +): + transformer = bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ) + + transformer.fit( + penguins_df_default_index[["species", "culmen_length_mm", "flipper_length_mm"]] + ) + result = transformer.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pandas.DataFrame( + { + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], + "scaled_flipper_length_mm": [-0.3495, -1.416, -0.9185], + }, + index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"), + ) + expected.scaled_culmen_length_mm = expected.scaled_culmen_length_mm.astype( + "Float64" + ) + expected.scaled_flipper_length_mm = expected.scaled_flipper_length_mm.astype( + "Float64" + ) + + pandas.testing.assert_frame_equal(result, expected, rtol=1e-3) diff --git a/tests/system/large/ml/test_core.py b/tests/system/large/ml/test_core.py new file mode 100644 index 00000000000..cd829f0ae2d --- /dev/null +++ b/tests/system/large/ml/test_core.py @@ -0,0 +1,183 @@ +# 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 pandas + +import bigframes.ml.core +import bigframes.ml.sql + + +def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_df): + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + + model = bigframes.ml.core.create_bqml_model( + train_X, train_y, options={"model_type": "linear_reg"} + ) + + # no data - report evaluation from the automatic data split + evaluate_result = model.evaluate().compute() + evaluate_expected = pandas.DataFrame( + { + "mean_absolute_error": [225.817334], + "mean_squared_error": [80540.705944], + "mean_squared_log_error": [0.004972], + "median_absolute_error": [173.080816], + "r2_score": [0.87529], + "explained_variance": [0.87529], + }, + dtype="Float64", + ) + evaluate_expected = evaluate_expected.reindex( + index=evaluate_expected.index.astype("Int64") + ) + pandas.testing.assert_frame_equal( + evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 + ) + + # evaluate on all training data + evaluate_result = model.evaluate(df).compute() + pandas.testing.assert_frame_equal( + evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + predictions = model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) + + new_name = f"{dataset_id}.my_model" + new_model = model.copy(new_name, True) + assert new_model.model_name == new_name + + fetch_result = session.bqclient.get_model(new_name) + assert fetch_result.model_type == "LINEAR_REGRESSION" + + +def test_bqml_manual_preprocessing_e2e( + session, dataset_id, penguins_df_default_index, new_penguins_df +): + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["body_mass_g"]] + transforms = [ + bigframes.ml.sql.ml_standard_scaler(column, column) + for column in train_X.columns.tolist() + ] + transforms.extend(train_y.columns.tolist()) + options = {"model_type": "linear_reg"} + model = bigframes.ml.core.create_bqml_model( + train_X, train_y, transforms=transforms, options=options + ) + + # no data - report evaluation from the automatic data split + evaluate_result = model.evaluate().compute() + evaluate_expected = pandas.DataFrame( + { + "mean_absolute_error": [309.477334], + "mean_squared_error": [152184.227218], + "mean_squared_log_error": [0.009524], + "median_absolute_error": [257.727777], + "r2_score": [0.764356], + "explained_variance": [0.764356], + }, + dtype="Float64", + ) + evaluate_expected = evaluate_expected.reindex( + index=evaluate_expected.index.astype("Int64") + ) + + pandas.testing.assert_frame_equal( + evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 + ) + + # evaluate on all training data + evaluate_result = model.evaluate(df).compute() + pandas.testing.assert_frame_equal( + evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + predictions = model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_body_mass_g": [3968.8, 3176.3, 3545.2]}, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) + + new_name = f"{dataset_id}.my_model" + new_model = model.copy(new_name, True) + assert new_model.model_name == new_name + + fetch_result = session.bqclient.get_model(new_name) + assert fetch_result.model_type == "LINEAR_REGRESSION" + + +def test_bqml_standalone_transform(penguins_df_default_index, new_penguins_df): + X = penguins_df_default_index[["culmen_length_mm", "species"]] + model = bigframes.ml.core.create_bqml_model( + X, + options={"model_type": "transform_only"}, + transforms=[ + "ML.STANDARD_SCALER(culmen_length_mm) OVER() AS scaled_culmen_length_mm", + "ML.ONE_HOT_ENCODER(species) OVER() AS onehotencoded_species", + ], + ) + + transformed = model.transform(new_penguins_df).compute() + expected = pandas.DataFrame( + { + "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + }, + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + expected["scaled_culmen_length_mm"] = expected["scaled_culmen_length_mm"].astype( + "Float64" + ) + pandas.testing.assert_frame_equal( + transformed[["scaled_culmen_length_mm", "onehotencoded_species"]], + expected, + check_exact=False, + rtol=0.1, + ) diff --git a/tests/system/large/ml/test_decomposition.py b/tests/system/large/ml/test_decomposition.py new file mode 100644 index 00000000000..57507c11342 --- /dev/null +++ b/tests/system/large/ml/test_decomposition.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. + +import pandas + +import bigframes.ml.decomposition + + +def test_decomposition_configure_fit_predict( + session, penguins_df_default_index, dataset_id +): + model = bigframes.ml.decomposition.PCA(n_components=3) + model.fit(penguins_df_default_index) + + pd_new_penguins = session.read_pandas( + pandas.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Gentoo penguin (Pygoscelis papua)", + "Adelie Penguin (Pygoscelis adeliae)", + ], + "island": ["Dream", "Biscoe", "Torgersen"], + "culmen_length_mm": [37.8, 46.5, 41.1], + "culmen_depth_mm": [18.1, 14.8, 18.6], + "flipper_length_mm": [193.0, 217.0, 189.0], + "body_mass_g": [3750.0, 5200.0, 3325.0], + "sex": ["MALE", "FEMALE", "MALE"], + } + ).set_index("tag_number") + ) + + result = model.predict(pd_new_penguins).compute() + expected = pandas.DataFrame( + { + "principal_component_1": [-1.459, 2.258, -1.685], + "principal_component_2": [-1.120, -1.351, -0.874], + "principal_component_3": [-0.646, 0.443, -0.704], + }, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + result.sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) + + # save, load, check n_components to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_pca_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_pca_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.n_components == 3 diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py new file mode 100644 index 00000000000..0d5da380d1e --- /dev/null +++ b/tests/system/large/ml/test_ensemble.py @@ -0,0 +1,458 @@ +# 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. + +from unittest import TestCase + +import pandas + +import bigframes.ml.ensemble + + +def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.XGBRegressor() + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [97.368139], + "mean_squared_error": [16284.877027], + "mean_squared_log_error": [0.0010189], + "median_absolute_error": [72.158691], + "r2_score": [0.974784], + "explained_variance": [0.974845], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pandas.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_xgbregressor_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_xgbregressor_model" + in reloaded_model._bqml_model.model_name + ) + + +def test_xgbregressor_dart_booster_multiple_params( + penguins_df_default_index, dataset_id +): + model = bigframes.ml.ensemble.XGBRegressor( + booster="dart", + tree_method="AUTO", + min_tree_child_weight=2, + colsample_bytree=0.95, + colsample_bylevel=0.95, + colsample_bynode=0.95, + num_parallel_tree=2, + max_depth=4, + subsample=0.95, + reg_alpha=0.0001, + reg_lambda=0.0001, + learning_rate=0.015, + max_iterations=4, + min_rel_progress=0.02, + ) + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "mean_absolute_error", + "mean_squared_error", + "mean_squared_log_error", + "median_absolute_error", + "r2_score", + "explained_variance", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_xgbregressor_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_xgbregressor_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.booster == "DART" + assert reloaded_model.dart_normalized_type == "TREE" + assert reloaded_model.tree_method == "AUTO" + assert reloaded_model.colsample_bytree == 0.95 + assert reloaded_model.colsample_bylevel == 0.95 + assert reloaded_model.colsample_bynode == 0.95 + assert reloaded_model.early_stop is True + assert reloaded_model.subsample == 0.95 + assert reloaded_model.reg_alpha == 0.0001 + assert reloaded_model.reg_lambda == 0.0001 + assert reloaded_model.learning_rate == 0.015 + assert reloaded_model.max_iterations == 4 + assert reloaded_model.min_rel_progress == 0.02 + assert reloaded_model.gamma == 0.0 + assert reloaded_model.max_depth == 4 + assert reloaded_model.min_tree_child_weight == 2 + assert reloaded_model.num_parallel_tree == 2 + + +def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.XGBClassifier() + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_xgbclassifier_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_xgbclassifier_model" + in reloaded_model._bqml_model.model_name + ) + + +def test_xgbclassifier_dart_booster_multiple_params( + penguins_df_default_index, dataset_id +): + model = bigframes.ml.ensemble.XGBClassifier( + booster="dart", + tree_method="AUTO", + min_tree_child_weight=2, + colsample_bytree=0.95, + colsample_bylevel=0.95, + colsample_bynode=0.95, + num_parallel_tree=2, + max_depth=4, + subsample=0.95, + reg_alpha=0.0001, + reg_lambda=0.0001, + learning_rate=0.015, + max_iterations=4, + min_rel_progress=0.02, + ) + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_xgbclassifier_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_xgbclassifier_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.booster == "DART" + assert reloaded_model.dart_normalized_type == "TREE" + assert reloaded_model.tree_method == "AUTO" + assert reloaded_model.colsample_bytree == 0.95 + assert reloaded_model.colsample_bylevel == 0.95 + assert reloaded_model.colsample_bynode == 0.95 + assert reloaded_model.early_stop is True + assert reloaded_model.subsample == 0.95 + assert reloaded_model.reg_alpha == 0.0001 + assert reloaded_model.reg_lambda == 0.0001 + assert reloaded_model.learning_rate == 0.015 + assert reloaded_model.max_iterations == 4 + assert reloaded_model.min_rel_progress == 0.02 + assert reloaded_model.gamma == 0.0 + assert reloaded_model.max_depth == 4 + assert reloaded_model.min_tree_child_weight == 2 + assert reloaded_model.num_parallel_tree == 2 + + +def test_randomforestregressor_default_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.RandomForestRegressor() + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "mean_absolute_error", + "mean_squared_error", + "mean_squared_log_error", + "median_absolute_error", + "r2_score", + "explained_variance", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_randomforestregressor_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_randomforestregressor_model" + in reloaded_model._bqml_model.model_name + ) + + +def test_randomforestregressor_multiple_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.RandomForestRegressor( + tree_method="AUTO", + min_tree_child_weight=2, + colsample_bytree=0.95, + colsample_bylevel=0.95, + colsample_bynode=0.95, + num_parallel_tree=90, + max_depth=14, + subsample=0.95, + reg_alpha=0.0001, + reg_lambda=0.0001, + min_rel_progress=0.02, + ) + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "mean_absolute_error", + "mean_squared_error", + "mean_squared_log_error", + "median_absolute_error", + "r2_score", + "explained_variance", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_randomforestregressor_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_randomforestregressor_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.tree_method == "AUTO" + assert reloaded_model.colsample_bytree == 0.95 + assert reloaded_model.colsample_bylevel == 0.95 + assert reloaded_model.colsample_bynode == 0.95 + assert reloaded_model.early_stop is True + assert reloaded_model.subsample == 0.95 + assert reloaded_model.reg_alpha == 0.0001 + assert reloaded_model.reg_lambda == 0.0001 + assert reloaded_model.min_rel_progress == 0.02 + assert reloaded_model.gamma == 0.0 + assert reloaded_model.max_depth == 14 + assert reloaded_model.min_tree_child_weight == 2 + assert reloaded_model.num_parallel_tree == 90 + assert reloaded_model.enable_global_explain is False + + +def test_randomforestclassifier_default_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.RandomForestClassifier() + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_randomforestclassifier_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_randomforestclassifier_model" + in reloaded_model._bqml_model.model_name + ) + + +def test_randomforestclassifier_multiple_params(penguins_df_default_index, dataset_id): + model = bigframes.ml.ensemble.RandomForestClassifier( + tree_method="AUTO", + min_tree_child_weight=2, + colsample_bytree=0.95, + colsample_bylevel=0.95, + colsample_bynode=0.95, + num_parallel_tree=90, + max_depth=14, + subsample=0.95, + reg_alpha=0.0001, + reg_lambda=0.0001, + min_rel_progress=0.02, + ) + + df = penguins_df_default_index.dropna().sample(n=70) + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_randomforestclassifier_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_randomforestclassifier_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.tree_method == "AUTO" + assert reloaded_model.colsample_bytree == 0.95 + assert reloaded_model.colsample_bylevel == 0.95 + assert reloaded_model.colsample_bynode == 0.95 + assert reloaded_model.early_stop is True + assert reloaded_model.subsample == 0.95 + assert reloaded_model.reg_alpha == 0.0001 + assert reloaded_model.reg_lambda == 0.0001 + assert reloaded_model.min_rel_progress == 0.02 + assert reloaded_model.gamma == 0.0 + assert reloaded_model.max_depth == 14 + assert reloaded_model.min_tree_child_weight == 2 + assert reloaded_model.num_parallel_tree == 90 + assert reloaded_model.enable_global_explain is False diff --git a/tests/system/large/ml/test_forecasting.py b/tests/system/large/ml/test_forecasting.py new file mode 100644 index 00000000000..98726eb289d --- /dev/null +++ b/tests/system/large/ml/test_forecasting.py @@ -0,0 +1,48 @@ +# 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 pandas as pd + +from bigframes.ml import forecasting + + +def test_arima_plus_model_fit_score( + time_series_df_default_index, dataset_id, new_time_series_df +): + model = forecasting.ARIMAPlus() + train_X = time_series_df_default_index[["parsed_date"]] + train_y = time_series_df_default_index[["total_visits"]] + model.fit(train_X, train_y) + + result = model.score( + new_time_series_df[["parsed_date"]], new_time_series_df[["total_visits"]] + ).compute() + expected = pd.DataFrame( + { + "mean_absolute_error": [154.742547], + "mean_squared_error": [26844.868855], + "root_mean_squared_error": [163.844038], + "mean_absolute_percentage_error": [6.189702], + "symmetric_mean_absolute_percentage_error": [6.097155], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pd.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load to ensure configuration was kept + reloaded_model = model.to_gbq(f"{dataset_id}.temp_configured_model", replace=True) + assert ( + f"{dataset_id}.temp_configured_model" in reloaded_model._bqml_model.model_name + ) diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py new file mode 100644 index 00000000000..33010a005cb --- /dev/null +++ b/tests/system/large/ml/test_linear_model.py @@ -0,0 +1,196 @@ +# 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 pandas as pd + +import bigframes.ml.linear_model + + +def test_linear_regression_configure_fit_score(penguins_df_default_index, dataset_id): + model = bigframes.ml.linear_model.LinearRegression(fit_intercept=False) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + expected = pd.DataFrame( + { + "mean_absolute_error": [225.735767], + "mean_squared_error": [80417.461828], + "mean_squared_log_error": [0.004967], + "median_absolute_error": [172.543702], + "r2_score": [0.87548], + "explained_variance": [0.87548], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pd.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq(f"{dataset_id}.temp_configured_model", replace=True) + assert ( + f"{dataset_id}.temp_configured_model" in reloaded_model._bqml_model.model_name + ) + + # TODO(yunmengxie): enable this once b/277242951 (fit_intercept missing from API) is fixed + # assert reloaded_model.fit_intercept == False + + +def test_linear_regression_manual_split_configure_fit_score( + penguins_df_default_index, dataset_id +): + model = bigframes.ml.linear_model.LinearRegression(fit_intercept=True) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + train_y = df[["body_mass_g"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + expected = pd.DataFrame( + { + "mean_absolute_error": [225.735767], + "mean_squared_error": [80417.461828], + "mean_squared_log_error": [0.004967], + "median_absolute_error": [172.543702], + "r2_score": [0.87548], + "explained_variance": [0.87548], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pd.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq(f"{dataset_id}.temp_configured_model", replace=True) + assert ( + f"{dataset_id}.temp_configured_model" in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.fit_intercept is True + + +def test_logistic_regression_auto_class_weights_configure_fit_score( + penguins_df_default_index, dataset_id +): + model = bigframes.ml.linear_model.LogisticRegression() + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + expected = pd.DataFrame( + { + "precision": [0.58085], + "recall": [0.582576], + "accuracy": [0.871257], + "f1_score": [0.58171], + "log_loss": [1.59285], + "roc_auc": [0.9602], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pd.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_logistic_reg_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_logistic_reg_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.fit_intercept is True + # TODO(gaotianxiang): enable this once (auto_class_weights missing from API) is fixed + # assert reloaded_model.auto_class_weights is True + + +def test_logistic_regression_manual_split_configure_fit_score( + penguins_df_default_index, dataset_id +): + model = bigframes.ml.linear_model.LogisticRegression(fit_intercept=True) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + train_y = df[["sex"]] + model.fit(train_X, train_y) + + # Check score to ensure the model was fitted + result = model.score(train_X, train_y).compute() + expected = pd.DataFrame( + { + "precision": [0.616753], + "recall": [0.618615], + "accuracy": [0.92515], + "f1_score": [0.617681], + "log_loss": [1.498832], + "roc_auc": [0.975807], + }, + dtype="Float64", + ) + expected = expected.reindex(index=expected.index.astype("Int64")) + pd.testing.assert_frame_equal(result, expected, check_exact=False, rtol=0.1) + + # save, load, check parameters to ensure configuration was kept + reloaded_model = model.to_gbq( + f"{dataset_id}.temp_configured_logistic_reg_model", replace=True + ) + assert ( + f"{dataset_id}.temp_configured_logistic_reg_model" + in reloaded_model._bqml_model.model_name + ) + assert reloaded_model.fit_intercept is True + assert reloaded_model.auto_class_weights is False diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py new file mode 100644 index 00000000000..9700ba2bf65 --- /dev/null +++ b/tests/system/large/ml/test_pipeline.py @@ -0,0 +1,524 @@ +# 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 pandas as pd + +from bigframes.ml import ( + cluster, + compose, + decomposition, + ensemble, + linear_model, + pipeline, + preprocessing, +) +from tests.system.utils import assert_pandas_df_equal_ignore_ordering + + +def test_pipeline_linear_regression_fit_score_predict( + session, penguins_df_default_index +): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("linreg", linear_model.LinearRegression()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["body_mass_g"]] + pl.fit(train_X, train_y) + + # Check score to ensure the model was fitted + score_result = pl.score(train_X, train_y).compute() + score_expected = pd.DataFrame( + { + "mean_absolute_error": [309.477334], + "mean_squared_error": [152184.227218], + "mean_squared_log_error": [0.009524], + "median_absolute_error": [257.727777], + "r2_score": [0.764356], + "explained_variance": [0.764356], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + {"predicted_body_mass_g": [3968.8, 3176.3, 3545.2]}, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) + + +def test_pipeline_logistic_regression_fit_score_predict( + session, penguins_df_default_index +): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("logreg", linear_model.LogisticRegression()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + pl.fit(train_X, train_y) + + # Check score to ensure the model was fitted + score_result = pl.score(train_X, train_y).compute() + score_expected = pd.DataFrame( + { + "precision": [0.537091], + "recall": [0.538636], + "accuracy": [0.805389], + "f1_score": [0.537716], + "log_loss": [1.445433], + "roc_auc": [0.917818], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + {"predicted_sex": ["MALE", "FEMALE", "FEMALE"]}, + dtype=pd.StringDtype(storage="pyarrow"), + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_sex"]], + expected, + ) + + +def test_pipeline_xgbregressor_fit_score_predict(session, penguins_df_default_index): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("xgbreg", ensemble.XGBRegressor()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["body_mass_g"]] + pl.fit(train_X, train_y) + + # Check score to ensure the model was fitted + score_result = pl.score(train_X, train_y).compute() + score_expected = pd.DataFrame( + { + "mean_absolute_error": [203.4001727989334], + "mean_squared_error": [74898.80551717622], + "mean_squared_log_error": [0.004394266810531861], + "median_absolute_error": [152.01806640625], + "r2_score": [0.8840255831308607], + "explained_variance": [0.8858505311591299], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + { + "predicted_body_mass_g": [ + 4287.34521484375, + 3198.351806640625, + 3385.34130859375, + ] + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) + + +def test_pipeline_random_forest_classifier_fit_score_predict( + session, penguins_df_default_index +): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("rfcls", ensemble.RandomForestClassifier()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + train_y = df[["sex"]] + pl.fit(train_X, train_y) + + # Check score to ensure the model was fitted + score_result = pl.score(train_X, train_y).compute() + score_expected = pd.DataFrame( + { + "precision": [0.587673], + "recall": [0.588781], + "accuracy": [0.88024], + "f1_score": [0.587644], + "log_loss": [0.859459], + "roc_auc": [0.971737], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + {"predicted_sex": ["MALE", "FEMALE", "FEMALE"]}, + dtype=pd.StringDtype(storage="pyarrow"), + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_sex"]], + expected, + ) + + +def test_pipeline_PCA_fit_predict(session, penguins_df_default_index): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("pca", decomposition.PCA()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + pl.fit(train_X) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + { + "principal_component_1": [-1.115259, -1.506141, -1.471174], + "principal_component_2": [-0.074824, 0.69664, 0.406104], + "principal_component_3": [0.500012, -0.544479, 0.075849], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[ + ["principal_component_1", "principal_component_2", "principal_component_3"] + ], + expected, + check_exact=False, + rtol=0.1, + ) + + +def test_pipeline_standard_scaler_kmeans_fit_predict( + session, penguins_pandas_df_default_index +): + """Test an unsupervised model with a non-BQML implementation of StandardScaler""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("kmeans", cluster.KMeans(n_clusters=2)), + ] + ) + + # kmeans is sensitive to the order with this configuration, so use ordered source data + df = session.read_pandas(penguins_pandas_df_default_index).dropna() + train_X = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + pl.fit(train_X) + + # predict new labels + pd_new_penguins = pd.DataFrame.from_dict( + { + "test1": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Dream", + "culmen_length_mm": 27.5, + "culmen_depth_mm": 8.5, + "flipper_length_mm": 99, + "body_mass_g": 4475, + "sex": "MALE", + }, + "test2": { + "species": "Chinstrap penguin (Pygoscelis antarctica)", + "island": "Dream", + "culmen_length_mm": 55.8, + "culmen_depth_mm": 29.8, + "flipper_length_mm": 307, + "body_mass_g": 4000, + "sex": "MALE", + }, + "test3": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Biscoe", + "culmen_length_mm": 19.7, + "culmen_depth_mm": 8.9, + "flipper_length_mm": 84, + "body_mass_g": 3550, + "sex": "MALE", + }, + "test4": { + "species": "Gentoo penguin (Pygoscelis papua)", + "island": "Biscoe", + "culmen_length_mm": 63.8, + "culmen_depth_mm": 33.9, + "flipper_length_mm": 298, + "body_mass_g": 4300, + "sex": "FEMALE", + }, + "test5": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Dream", + "culmen_length_mm": 27.5, + "culmen_depth_mm": 8.5, + "flipper_length_mm": 99, + "body_mass_g": 4475, + "sex": "MALE", + }, + "test6": { + "species": "Chinstrap penguin (Pygoscelis antarctica)", + "island": "Dream", + "culmen_length_mm": 55.8, + "culmen_depth_mm": 29.8, + "flipper_length_mm": 307, + "body_mass_g": 4000, + "sex": "MALE", + }, + }, + orient="index", + ) + pd_new_penguins.index.name = "observation" + + new_penguins = session.read_pandas(pd_new_penguins) + result = pl.predict(new_penguins).to_pandas().sort_index() + expected = pd.DataFrame( + {"CENTROID_ID": [1, 2, 1, 2, 1, 2]}, + dtype="Int64", + index=pd.Index( + ["test1", "test2", "test3", "test4", "test5", "test6"], + dtype="string[pyarrow]", + ), + ) + expected.index.name = "observation" + assert_pandas_df_equal_ignore_ordering(result, expected) + + +def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_index): + """Test a preprocessing step that manages heterogenous data with ColumnTransformer""" + pl = pipeline.Pipeline( + [ + ( + "preproc", + compose.ColumnTransformer( + [ + ( + "onehot", + preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ), + ), + ("linreg", linear_model.LinearRegression()), + ] + ) + + df = penguins_df_default_index.dropna() + train_X = df[["species", "culmen_length_mm", "flipper_length_mm"]] + train_y = df[["body_mass_g"]] + pl.fit(train_X, train_y) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Adelie Penguin (Pygoscelis adeliae)", + "Chinstrap penguin (Pygoscelis antarctica)", + ], + "island": ["Torgersen", "Torgersen", "Dream"], + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [18.8, 17.2, 18.1], + "flipper_length_mm": [196.0, 181.0, 188.0], + "sex": ["MALE", "FEMALE", "FEMALE"], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins).to_pandas() + expected = pd.DataFrame( + {"predicted_body_mass_g": [3909.2, 3436.0, 2860.0]}, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py new file mode 100644 index 00000000000..714252c9e00 --- /dev/null +++ b/tests/system/large/test_remote_function.py @@ -0,0 +1,776 @@ +# 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. + +from datetime import datetime +import importlib.util +import inspect +import math # must keep this at top level to test udf referring global import +import os.path +import shutil +import tempfile +import textwrap + +from google.api_core.exceptions import NotFound, ResourceExhausted +from google.cloud import functions_v2 +import pandas +import pytest +import test_utils.prefixer + +from bigframes.remote_function import ( + get_cloud_function_name, + get_remote_function_locations, +) +from tests.system.utils import assert_pandas_df_equal_ignore_ordering + +# Use this to control the number of cloud functions being deleted in a single +# test session. This should help soften the spike of the number of mutations per +# minute tracked against a quota limit (default 60, increased to 120 for +# bigframes-dev project) by the Cloud Functions API +# We are running pytest with "-n 20". Let's say each session lasts about a +# minute, so we are setting a limit of 120/20 = 6 deletions per session. +_MAX_NUM_FUNCTIONS_TO_DELETE_PER_SESSION = 6 + +# NOTE: Keep this import at the top level to test global var behavior with +# remote functions +_team_pi = "Team Pi" +_team_euler = "Team Euler" + + +def get_remote_function_endpoints(bigquery_client, dataset_id): + """Get endpoints used by the remote functions in a datset""" + endpoints = set() + routines = bigquery_client.list_routines(dataset=dataset_id) + for routine in routines: + rf_options = routine._properties.get("remoteFunctionOptions") + if not rf_options: + continue + rf_endpoint = rf_options.get("endpoint") + if rf_endpoint: + endpoints.add(rf_endpoint) + return endpoints + + +def get_cloud_functions(functions_client, project, location, name_prefix="bigframes-"): + """Get the cloud functions in the given project and location.""" + _, location = get_remote_function_locations(location) + parent = f"projects/{project}/locations/{location}" + request = functions_v2.ListFunctionsRequest(parent=parent) + page_result = functions_client.list_functions(request=request) + full_name_prefix = parent + f"/functions/{name_prefix}" + for response in page_result: + if not name_prefix or response.name.startswith(full_name_prefix): + yield response + + +def delete_cloud_function(functions_client, full_name): + """Delete a cloud function with the given fully qualified name.""" + request = functions_v2.DeleteFunctionRequest(name=full_name) + operation = functions_client.delete_function(request=request) + return operation + + +def cleanup_remote_function_assets( + bigquery_client, functions_client, remote_udf, ignore_failures=True +): + """Clean up the GCP assets behind a bigframes remote function.""" + try: + bigquery_client.delete_routine(remote_udf.bigframes_remote_function) + delete_cloud_function(functions_client, remote_udf.bigframes_cloud_function) + except Exception: + # By default don't raise exception in cleanup + if not ignore_failures: + raise + + +def make_uniq_udf(udf): + """Transform a udf to another with same behavior but a unique name.""" + prefixer = test_utils.prefixer.Prefixer(udf.__name__, "") + udf_uniq_name = prefixer.create_prefix() + udf_file_name = f"{udf_uniq_name}.py" + + # We are not using `tempfile.TemporaryDirectory()` because we want to keep + # the temp code around, otherwise `inspect.getsource()` complains. + tmpdir = tempfile.mkdtemp() + udf_file_path = os.path.join(tmpdir, udf_file_name) + with open(udf_file_path, "w") as f: + # TODO(shobs): Find a better way of modifying the udf, maybe regex? + source_key = f"def {udf.__name__}" + target_key = f"def {udf_uniq_name}" + source_code = textwrap.dedent(inspect.getsource(udf)) + target_code = source_code.replace(source_key, target_key, 1) + f.write(target_code) + spec = importlib.util.spec_from_file_location(udf_file_name, udf_file_path) + return getattr(spec.loader.load_module(), udf_uniq_name), tmpdir + + +@pytest.fixture(scope="module") +def bq_cf_connection() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us --project_id=bigframes-dev bigframes-rf-conn + """ + return "bigframes-rf-conn" + + +@pytest.fixture(scope="module") +def functions_client() -> functions_v2.FunctionServiceClient: + """Cloud Functions client""" + return functions_v2.FunctionServiceClient() + + +@pytest.fixture(scope="module", autouse=True) +def cleanup_cloud_functions(session, functions_client, dataset_id_permanent): + """Clean up stale cloud functions.""" + permanent_endpoints = get_remote_function_endpoints( + session.bqclient, dataset_id_permanent + ) + delete_count = 0 + for cloud_function in get_cloud_functions( + functions_client, session.bqclient.project, session.bqclient.location + ): + # Ignore bigframes cloud functions referred by the remote functions in + # the permanent dataset + if cloud_function.service_config.uri in permanent_endpoints: + continue + + # Ignore the functions less than one day old + age = datetime.now() - datetime.fromtimestamp( + cloud_function.update_time.timestamp() + ) + if age.days <= 0: + continue + + # Go ahead and delete + try: + delete_cloud_function(functions_client, cloud_function.name) + delete_count += 1 + if delete_count >= _MAX_NUM_FUNCTIONS_TO_DELETE_PER_SESSION: + break + except NotFound: + # This can happen when multiple pytest sessions are running in + # parallel. Two or more sessions may discover the same cloud + # function, but only one of them would be able to delete it + # successfully, while the other instance will run into this + # exception. Ignore this exception. + pass + except ResourceExhausted: + # This can happen if we are hitting GCP limits, e.g. + # google.api_core.exceptions.ResourceExhausted: 429 Quota exceeded + # for quota metric 'Per project mutation requests' and limit + # 'Per project mutation requests per minute per region' of service + # 'cloudfunctions.googleapis.com' for consumer + # 'project_number:1084210331973'. + # [reason: "RATE_LIMIT_EXCEEDED" domain: "googleapis.com" ... + # Let's stop further clean up and leave it to later. + break + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_multiply_with_ibis( + session, + scalars_table_id, + ibis_client, + dataset_id, + bq_cf_connection, + functions_client, +): + try: + + @session.remote_function( + [int, int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + ) + def multiply(x, y): + return x * y + + project_id, dataset_name, table_name = scalars_table_id.split(".") + if not ibis_client.dataset: + ibis_client.dataset = dataset_name + + col_name = "int64_col" + table = ibis_client.tables[table_name] + table = table.filter(table[col_name].notnull()).order_by("rowindex").head(10) + pandas_df_orig = table.execute() + + col = table[col_name] + col_2x = multiply(col, 2).name("int64_col_2x") + col_square = multiply(col, col).name("int64_col_square") + table = table.mutate([col_2x, col_square]) + pandas_df_new = table.execute() + + pandas.testing.assert_series_equal( + pandas_df_orig[col_name] * 2, + pandas_df_new["int64_col_2x"], + check_names=False, + ) + + pandas.testing.assert_series_equal( + pandas_df_orig[col_name] * pandas_df_orig[col_name], + pandas_df_new["int64_col_square"], + check_names=False, + ) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets(session.bqclient, functions_client, multiply) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_stringify_with_ibis( + session, + scalars_table_id, + ibis_client, + dataset_id, + bq_cf_connection, + functions_client, +): + try: + + @session.remote_function( + [int], + str, + dataset_id, + bq_cf_connection, + reuse=False, + ) + def stringify(x): + return f"I got {x}" + + project_id, dataset_name, table_name = scalars_table_id.split(".") + if not ibis_client.dataset: + ibis_client.dataset = dataset_name + + col_name = "int64_col" + table = ibis_client.tables[table_name] + table = table.filter(table[col_name].notnull()).order_by("rowindex").head(10) + pandas_df_orig = table.execute() + + col = table[col_name] + col_2x = stringify(col).name("int64_str_col") + table = table.mutate([col_2x]) + pandas_df_new = table.execute() + + pandas.testing.assert_series_equal( + pandas_df_orig[col_name].apply(lambda x: f"I got {x}"), + pandas_df_new["int64_str_col"], + check_names=False, + ) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets(session.bqclient, functions_client, stringify) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_decorator_with_bigframes_series( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + @session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets(session.bqclient, functions_client, square) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_explicit_with_bigframes_series( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def add_one(x): + return x + 1 + + remote_add_one = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + )(add_one) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(remote_add_one) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(add_one) + # TODO(shobs): Figure why pandas .apply() changes the dtype, e.g. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, remote_add_one + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_referring_outside_var( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + POSITIVE_SIGN = 1 + NEGATIVE_SIGN = -1 + NO_SIGN = 0 + + def sign(num): + if num > 0: + return POSITIVE_SIGN + elif num < 0: + return NEGATIVE_SIGN + return NO_SIGN + + remote_sign = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + )(sign) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(remote_sign) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(sign) + # TODO(shobs): Figure why pandas .apply() changes the dtype, e.g. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets(session.bqclient, functions_client, remote_sign) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_referring_outside_import( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + import math as mymath + + def circumference(radius): + return 2 * mymath.pi * radius + + remote_circumference = session.remote_function( + [float], + float, + dataset_id, + bq_cf_connection, + reuse=False, + )(circumference) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_float64_col = scalars_df["float64_col"] + bf_float64_col_filter = bf_float64_col.notnull() + bf_float64_col_filtered = bf_float64_col[bf_float64_col_filter] + bf_result_col = bf_float64_col_filtered.apply(remote_circumference) + bf_result = ( + bf_float64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_float64_col = scalars_pandas_df["float64_col"] + pd_float64_col_filter = pd_float64_col.notnull() + pd_float64_col_filtered = pd_float64_col[pd_float64_col_filter] + pd_result_col = pd_float64_col_filtered.apply(circumference) + # TODO(shobs): Figure why pandas .apply() changes the dtype, e.g. + # pd_float64_col_filtered.dtype is Float64Dtype() + # pd_float64_col_filtered.apply(lambda x: x).dtype is float64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Float64Dtype()) + pd_result = pd_float64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, remote_circumference + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_referring_global_var_and_import( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def find_team(num): + boundary = (math.pi + math.e) / 2 + if num >= boundary: + return _team_euler + return _team_pi + + remote_find_team = session.remote_function( + [float], + str, + dataset_id, + bq_cf_connection, + reuse=False, + )(find_team) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_float64_col = scalars_df["float64_col"] + bf_float64_col_filter = bf_float64_col.notnull() + bf_float64_col_filtered = bf_float64_col[bf_float64_col_filter] + bf_result_col = bf_float64_col_filtered.apply(remote_find_team) + bf_result = ( + bf_float64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_float64_col = scalars_pandas_df["float64_col"] + pd_float64_col_filter = pd_float64_col.notnull() + pd_float64_col_filtered = pd_float64_col[pd_float64_col_filter] + pd_result_col = pd_float64_col_filtered.apply(find_team) + # TODO(shobs): Figure if the dtype mismatch is by design: + # bf_result.dtype: string[pyarrow] + # pd_result.dtype: dtype('O'). + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.StringDtype(storage="pyarrow")) + pd_result = pd_float64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, remote_find_team + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_restore_with_bigframes_series( + session, + scalars_dfs, + dataset_id, + bq_cf_connection, + functions_client, +): + try: + + def add_one(x): + return x + 1 + + # Make a unique udf + add_one_uniq, add_one_uniq_dir = make_uniq_udf(add_one) + + # This is a bit of a hack but we need to remove the reference to a foreign + # module, otherwise the serialization would keep the foreign module + # reference and deserialization would fail with error like following: + # ModuleNotFoundError: No module named 'add_one_2nxcmd9j' + # TODO(shobs): Figure out if there is a better way of generating the unique + # function object, but for now let's just set it to same module as the + # original udf. + add_one_uniq.__module__ = add_one.__module__ + + # Expected cloud function name for the unique udf + add_one_uniq_cf_name = get_cloud_function_name(add_one_uniq) + + # There should be no cloud function yet for the unique udf + cloud_functions = list( + get_cloud_functions( + functions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix=add_one_uniq_cf_name, + ) + ) + assert len(cloud_functions) == 0 + + # The first time both the cloud function and the bq remote function don't + # exist and would be created + remote_add_one = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=True, + )(add_one_uniq) + + # There should have been excactly one cloud function created at this point + cloud_functions = list( + get_cloud_functions( + functions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix=add_one_uniq_cf_name, + ) + ) + assert len(cloud_functions) == 1 + + # We will test this twice + def inner_test(): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(remote_add_one) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(add_one_uniq) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + # Test that the remote function works as expected + inner_test() + + # Let's delete the cloud function while not touching the bq remote function + delete_operation = delete_cloud_function( + functions_client, cloud_functions[0].name + ) + delete_operation.result() + assert delete_operation.done() + + # There should be no cloud functions at this point for the uniq udf + cloud_functions = list( + get_cloud_functions( + functions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix=add_one_uniq_cf_name, + ) + ) + assert len(cloud_functions) == 0 + + # The second time bigframes detects that the required cloud function doesn't + # exist even though the remote function exists, and goes ahead and recreates + # the cloud function + remote_add_one = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=True, + )(add_one_uniq) + + # There should be excactly one cloud function again + cloud_functions = list( + get_cloud_functions( + functions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix=add_one_uniq_cf_name, + ) + ) + assert len(cloud_functions) == 1 + + # Test again after the cloud function is restored that the remote function + # works as expected + inner_test() + + # clean up the temp code + shutil.rmtree(add_one_uniq_dir) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, remote_add_one + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_mask_default_value( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def is_odd(num): + flag = False + try: + flag = num % 2 == 1 + except TypeError: + pass + return flag + + is_odd_remote = session.remote_function( + [int], + bool, + dataset_id, + bq_cf_connection, + reuse=False, + )(is_odd) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_result_col = bf_int64_col.mask(is_odd_remote) + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_result_col = pd_int64_col.mask(is_odd) + pd_result = pd_int64_col.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, is_odd_remote + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_mask_custom_value( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def is_odd(num): + flag = False + try: + flag = num % 2 == 1 + except TypeError: + pass + return flag + + is_odd_remote = session.remote_function( + [int], + bool, + dataset_id, + bq_cf_connection, + reuse=False, + )(is_odd) + + scalars_df, scalars_pandas_df = scalars_dfs + + # TODO(shobs): Revisit this test when NA handling of pandas' Series.mask is + # fixed https://github.com/pandas-dev/pandas/issues/52955, + # for now filter out the nulls and test the rest + bf_int64_col = scalars_df["int64_col"] + bf_result_col = bf_int64_col[bf_int64_col.notnull()].mask(is_odd_remote, -1) + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_result_col = pd_int64_col[pd_int64_col.notnull()].mask(is_odd, -1) + pd_result = pd_int64_col.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, is_odd_remote + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_udf_lambda( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + add_one_lambda = lambda x: x + 1 # noqa: E731 + + add_one_lambda_remote = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + )(add_one_lambda) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(add_one_lambda_remote) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(add_one_lambda) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, add_one_lambda_remote + ) diff --git a/tests/system/large/test_session.py b/tests/system/large/test_session.py new file mode 100644 index 00000000000..62fa5a83d33 --- /dev/null +++ b/tests/system/large/test_session.py @@ -0,0 +1,52 @@ +# 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 + +from bigframes import Session + + +@pytest.mark.parametrize( + ("query_or_table", "index_col"), + [ + pytest.param( + "bigquery-public-data.patents_view.ipcr_201708", + (), + id="1g_table_w_default_index", + ), + pytest.param( + "bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2011", + (), + id="30g_table_w_default_index", + ), + # TODO(chelsealin): Disable the long run tests until we have propertily + # ordering support to avoid materializating any data. + # # Adding default index to large tables would take much longer time, + # # e.g. ~5 mins for a 100G table, ~20 mins for a 1T table. + # pytest.param( + # "bigquery-public-data.stackoverflow.post_history", + # ["id"], + # id="100g_table_w_unique_column_index", + # ), + # pytest.param( + # "bigquery-public-data.wise_all_sky_data_release.all_wise", + # ["cntr"], + # id="1t_table_w_unique_column_index", + # ), + ], +) +def test_read_gbq_for_large_tables(session: Session, query_or_table, index_col): + """Verify read_gbq() is able to read large tables.""" + df = session.read_gbq(query_or_table, index_col=index_col) + assert len(df.columns) != 0 diff --git a/tests/system/small/__init__.py b/tests/system/small/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/system/small/__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/tests/system/small/ml/__init__.py b/tests/system/small/ml/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/system/small/ml/__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/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py new file mode 100644 index 00000000000..5c3abf25dcc --- /dev/null +++ b/tests/system/small/ml/conftest.py @@ -0,0 +1,294 @@ +# 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 hashlib +import logging +from typing import cast +import uuid + +import google.cloud.exceptions +import pandas as pd +import pytest + +from bigframes.ml import core, forecasting, imported, llm +import bigframes.ml.cluster +import bigframes.ml.core +import bigframes.ml.ensemble +import bigframes.ml.linear_model + + +@pytest.fixture(scope="session") +def ml_connection() -> str: + return "bigframes-dev.us.bigframes-ml" + + +@pytest.fixture(scope="session") +def penguins_bqml_linear_model( + session, penguins_linear_model_name +) -> bigframes.ml.core.BqmlModel: + model = session.bqclient.get_model(penguins_linear_model_name) + return bigframes.ml.core.BqmlModel(session, model) + + +@pytest.fixture(scope="function") +def ephemera_penguins_bqml_linear_model( + penguins_bqml_linear_model, +) -> bigframes.ml.linear_model.LinearRegression: + model = penguins_bqml_linear_model + return model.copy( + f"{model._model.project}.{model._model.dataset_id}.{uuid.uuid4().hex}" + ) + + +@pytest.fixture(scope="session") +def penguins_linear_model( + session, penguins_linear_model_name: str +) -> bigframes.ml.linear_model.LinearRegression: + return cast( + bigframes.ml.linear_model.LinearRegression, + session.read_gbq_model(penguins_linear_model_name), + ) + + +@pytest.fixture(scope="function") +def ephemera_penguins_linear_model( + ephemera_penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, +) -> bigframes.ml.linear_model.LinearRegression: + bf_model = bigframes.ml.linear_model.LinearRegression() + bf_model._bqml_model = ephemera_penguins_bqml_linear_model + return bf_model + + +@pytest.fixture(scope="session") +def penguins_logistic_model( + session, penguins_logistic_model_name +) -> bigframes.ml.linear_model.LogisticRegression: + return cast( + bigframes.ml.linear_model.LogisticRegression, + session.read_gbq_model(penguins_logistic_model_name), + ) + + +@pytest.fixture(scope="session") +def penguins_xgbregressor_model( + session, penguins_xgbregressor_model_name +) -> bigframes.ml.ensemble.XGBRegressor: + return cast( + bigframes.ml.ensemble.XGBRegressor, + session.read_gbq_model(penguins_xgbregressor_model_name), + ) + + +@pytest.fixture(scope="session") +def penguins_xgbclassifier_model( + session, penguins_xgbclassifier_model_name +) -> bigframes.ml.ensemble.XGBClassifier: + return cast( + bigframes.ml.ensemble.XGBClassifier, + session.read_gbq_model(penguins_xgbclassifier_model_name), + ) + + +@pytest.fixture(scope="session") +def penguins_randomforest_regressor_model( + session, penguins_randomforest_regressor_model_name +) -> bigframes.ml.ensemble.RandomForestRegressor: + return cast( + bigframes.ml.ensemble.RandomForestRegressor, + session.read_gbq_model(penguins_randomforest_regressor_model_name), + ) + + +@pytest.fixture(scope="session") +def penguins_randomforest_classifier_model( + session, penguins_randomforest_classifier_model_name +) -> bigframes.ml.ensemble.RandomForestClassifier: + return cast( + bigframes.ml.ensemble.RandomForestClassifier, + session.read_gbq_model(penguins_randomforest_classifier_model_name), + ) + + +@pytest.fixture(scope="session") +def penguins_kmeans_model( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> bigframes.ml.cluster.KMeans: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='kmeans', + num_clusters=3 +) AS SELECT + culmen_length_mm, + culmen_depth_mm, + flipper_length_mm, + sex +FROM `{penguins_table_id}`""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_cluster_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + return session.read_gbq_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_kmeans_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + return session.read_gbq_model(model_name) + + +@pytest.fixture(scope="session") +def penguins_pca_model( + session: bigframes.Session, dataset_id_permanent, penguins_table_id +) -> bigframes.ml.decomposition.PCA: + + # TODO(yunmengxie): Create a shared method to get different types of pretrained models. + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='pca', + num_principal_components=3 +) AS SELECT + * +FROM `{penguins_table_id}`""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = ( + f"{dataset_id_permanent}.penguins_pca_{hashlib.md5(sql.encode()).hexdigest()}" + ) + sql = sql.replace("$model_name", model_name) + + try: + return session.read_gbq_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_pca_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + return session.read_gbq_model(model_name) + + +@pytest.fixture(scope="session") +def llm_text_pandas_df(): + """Additional data matching the penguins dataset, with a new index""" + return pd.DataFrame( + { + "prompt": [ + "What is BigQuery?", + "What is BQML?", + "What is BigQuery DataFrame?", + ], + } + ) + + +@pytest.fixture(scope="session") +def onnx_iris_pandas_df(): + """Data matching the iris dataset.""" + return pd.DataFrame( + { + "sepal_length": [4.9, 5.1, 34.7], + "sepal_width": [3.0, 5.1, 24.7], + "petal_length": [1.4, 1.5, 13.3], + "petal_width": [0.4, 0.2, 18.3], + "species": [ + "setosa", + "setosa", + "virginica", + ], + } + ) + + +@pytest.fixture(scope="session") +def onnx_iris_df(session, onnx_iris_pandas_df): + return session.read_pandas(onnx_iris_pandas_df) + + +@pytest.fixture(scope="session") +def llm_text_df(session, llm_text_pandas_df): + return session.read_pandas(llm_text_pandas_df) + + +@pytest.fixture(scope="session") +def bqml_palm2_text_generator_model(session, ml_connection) -> core.BqmlModel: + options = { + "remote_service_type": "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1", + } + return core.create_bqml_remote_model( + session=session, connection_name=ml_connection, options=options + ) + + +@pytest.fixture(scope="session") +def palm2_text_generator_model(session, ml_connection) -> llm.PaLM2TextGenerator: + return llm.PaLM2TextGenerator(session=session, connection_name=ml_connection) + + +@pytest.fixture(scope="function") +def ephemera_palm2_text_generator_model( + session, ml_connection +) -> llm.PaLM2TextGenerator: + return llm.PaLM2TextGenerator(session=session, connection_name=ml_connection) + + +@pytest.fixture(scope="session") +def palm2_embedding_generator_model( + session, ml_connection +) -> llm.PaLM2EmbeddingGenerator: + return llm.PaLM2EmbeddingGenerator(session=session, connection_name=ml_connection) + + +@pytest.fixture(scope="session") +def time_series_bqml_arima_plus_model( + session, time_series_arima_plus_model_name +) -> core.BqmlModel: + model = session.bqclient.get_model(time_series_arima_plus_model_name) + return bigframes.ml.core.BqmlModel(session, model) + + +@pytest.fixture(scope="session") +def time_series_arima_plus_model( + session, time_series_arima_plus_model_name +) -> forecasting.ARIMAPlus: + return cast( + forecasting.ARIMAPlus, + session.read_gbq_model(time_series_arima_plus_model_name), + ) + + +@pytest.fixture(scope="session") +def imported_tensorflow_model(session) -> imported.TensorFlowModel: + return imported.TensorFlowModel( + session=session, + model_path="gs://cloud-training-demos/txtclass/export/exporter/1549825580/*", + ) + + +@pytest.fixture(scope="function") +def ephemera_imported_tensorflow_model(session) -> imported.TensorFlowModel: + return imported.TensorFlowModel( + session=session, + model_path="gs://cloud-training-demos/txtclass/export/exporter/1549825580/*", + ) + + +@pytest.fixture(scope="session") +def imported_onnx_model(session) -> imported.OnnxModel: + return imported.OnnxModel( + session=session, + model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx", + ) diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py new file mode 100644 index 00000000000..4aefc5fa697 --- /dev/null +++ b/tests/system/small/ml/test_cluster.py @@ -0,0 +1,79 @@ +# 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 pandas + +import bigframes.ml.cluster +from tests.system.utils import assert_pandas_df_equal_ignore_ordering + + +def test_model_predict(session, penguins_kmeans_model: bigframes.ml.cluster.KMeans): + pd_new_penguins = pandas.DataFrame.from_dict( + { + "test1": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Dream", + "culmen_length_mm": 37.5, + "culmen_depth_mm": 18.5, + "flipper_length_mm": 199, + "body_mass_g": 4475, + "sex": "MALE", + }, + "test2": { + "species": "Chinstrap penguin (Pygoscelis antarctica)", + "island": "Dream", + "culmen_length_mm": 55.8, + "culmen_depth_mm": 19.8, + "flipper_length_mm": 207, + "body_mass_g": 4000, + "sex": "MALE", + }, + "test3": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Biscoe", + "culmen_length_mm": 39.7, + "culmen_depth_mm": 18.9, + "flipper_length_mm": 184, + "body_mass_g": 3550, + "sex": "MALE", + }, + "test4": { + "species": "Gentoo penguin (Pygoscelis papua)", + "island": "Biscoe", + "culmen_length_mm": 43.8, + "culmen_depth_mm": 13.9, + "flipper_length_mm": 208, + "body_mass_g": 4300, + "sex": "FEMALE", + }, + }, + orient="index", + ) + pd_new_penguins.index.name = "observation" + + new_penguins = session.read_pandas(pd_new_penguins) + result = penguins_kmeans_model.predict(new_penguins).compute() + expected = pandas.DataFrame( + {"CENTROID_ID": [2, 3, 1, 2]}, + dtype="Int64", + index=pandas.Index( + ["test1", "test2", "test3", "test4"], dtype="string[pyarrow]" + ), + ) + expected.index.name = "observation" + assert_pandas_df_equal_ignore_ordering(result, expected) + + +def test_loaded_config(penguins_kmeans_model): + assert penguins_kmeans_model.n_clusters == 3 diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py new file mode 100644 index 00000000000..d63e401dd7a --- /dev/null +++ b/tests/system/small/ml/test_core.py @@ -0,0 +1,193 @@ +# 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. + +from datetime import datetime +import typing +from unittest import TestCase + +import pandas as pd +import pyarrow as pa +import pytz + +import bigframes +import bigframes.ml.core + + +def test_model_eval( + penguins_bqml_linear_model, +): + result = penguins_bqml_linear_model.evaluate().compute() + expected = pd.DataFrame( + { + "mean_absolute_error": [227.01223], + "mean_squared_error": [81838.159892], + "mean_squared_log_error": [0.00507], + "median_absolute_error": [173.080816], + "r2_score": [0.872377], + "explained_variance": [0.872377], + }, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_model_eval_with_data(penguins_bqml_linear_model, penguins_df_default_index): + result = penguins_bqml_linear_model.evaluate( + penguins_df_default_index.dropna() + ).compute() + expected = pd.DataFrame( + { + "mean_absolute_error": [225.817334], + "mean_squared_error": [80540.705944], + "mean_squared_log_error": [0.004972], + "median_absolute_error": [173.080816], + "r2_score": [0.87529], + "explained_variance": [0.87529], + }, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_model_predict( + penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, new_penguins_df +): + predictions = penguins_bqml_linear_model.predict(new_penguins_df).compute() + expected = pd.DataFrame( + {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]].sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) + + +def test_model_predict_with_unnamed_index( + penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, new_penguins_df +): + + # This will result in an index that lacks a name, which the ML library will + # need to persist through the call to ML.PREDICT + new_penguins_df = new_penguins_df.reset_index() + + # remove the middle tag number to ensure we're really keeping the unnamed index + new_penguins_df = typing.cast( + bigframes.dataframe.DataFrame, + new_penguins_df[new_penguins_df.tag_number != 1672], + ) + + predictions = penguins_bqml_linear_model.predict(new_penguins_df).compute() + + expected = pd.DataFrame( + {"predicted_body_mass_g": [4030.1, 3177.9]}, + dtype="Float64", + index=pd.Index([0, 2], dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]].sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) + + +def test_model_generate_text( + bqml_palm2_text_generator_model: bigframes.ml.core.BqmlModel, llm_text_df +): + options = { + "temperature": 0.5, + "max_output_tokens": 100, + "top_k": 20, + "top_p": 0.5, + "flatten_json_output": True, + } + df = bqml_palm2_text_generator_model.generate_text( + llm_text_df, options=options + ).compute() + + TestCase().assertSequenceEqual(df.shape, (3, 4)) + TestCase().assertSequenceEqual( + [ + "ml_generate_text_llm_result", + "ml_generate_text_rai_result", + "ml_generate_text_status", + "prompt", + ], + df.columns.to_list(), + ) + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +def test_model_forecast(time_series_bqml_arima_plus_model: bigframes.ml.core.BqmlModel): + utc = pytz.utc + forecast = time_series_bqml_arima_plus_model.forecast().compute()[ + ["forecast_timestamp", "forecast_value"] + ] + expected = pd.DataFrame( + { + "forecast_timestamp": [ + datetime(2017, 8, 2, tzinfo=utc), + datetime(2017, 8, 3, tzinfo=utc), + datetime(2017, 8, 4, tzinfo=utc), + ], + "forecast_value": [2724.472284, 2593.368389, 2353.613034], + } + ) + expected["forecast_value"] = expected["forecast_value"].astype(pd.Float64Dtype()) + expected["forecast_timestamp"] = expected["forecast_timestamp"].astype( + pd.ArrowDtype(pa.timestamp("us", tz="UTC")) + ) + pd.testing.assert_frame_equal( + forecast, + expected, + rtol=0.1, + check_index_type=False, + ) + + +def test_model_register(ephemera_penguins_bqml_linear_model): + model = ephemera_penguins_bqml_linear_model + model.register() + + model_name = "bigframes_" + model.model.model_id + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert model_name in model.model.training_runs[-1]["vertexAiModelId"] + + +def test_model_register_with_params(ephemera_penguins_bqml_linear_model): + model_name = "bigframes_system_test_model" + model = ephemera_penguins_bqml_linear_model + model.register(model_name) + + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert model_name in model.model.training_runs[-1]["vertexAiModelId"] diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py new file mode 100644 index 00000000000..7e0421129eb --- /dev/null +++ b/tests/system/small/ml/test_decomposition.py @@ -0,0 +1,55 @@ +# 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 pandas + +import bigframes.ml.decomposition + + +def test_model_predict(session, penguins_pca_model: bigframes.ml.decomposition.PCA): + new_penguins = session.read_pandas( + pandas.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Gentoo penguin (Pygoscelis papua)", + "Adelie Penguin (Pygoscelis adeliae)", + ], + "island": ["Dream", "Biscoe", "Torgersen"], + "culmen_length_mm": [37.8, 46.5, 41.1], + "culmen_depth_mm": [18.1, 14.8, 18.6], + "flipper_length_mm": [193.0, 217.0, 189.0], + "body_mass_g": [3750.0, 5200.0, 3325.0], + "sex": ["MALE", "FEMALE", "MALE"], + } + ).set_index("tag_number") + ) + + predictions = penguins_pca_model.predict(new_penguins).compute() + expected = pandas.DataFrame( + { + "principal_component_1": [-1.459, 2.258, -1.685], + "principal_component_2": [-1.120, -1.351, -0.874], + "principal_component_3": [-0.646, 0.443, -0.704], + }, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + predictions.sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) diff --git a/tests/system/small/ml/test_ensemble.py b/tests/system/small/ml/test_ensemble.py new file mode 100644 index 00000000000..61c60c249f2 --- /dev/null +++ b/tests/system/small/ml/test_ensemble.py @@ -0,0 +1,427 @@ +# 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. + +from unittest import TestCase + +import google.api_core.exceptions +import pandas +import pytest + +import bigframes.ml.ensemble + + +def test_xgbregressor_model_score( + penguins_xgbregressor_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = penguins_xgbregressor_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [108.77582], + "mean_squared_error": [20943.272738], + "mean_squared_log_error": [0.00135], + "median_absolute_error": [86.313477], + "r2_score": [0.967571], + "explained_variance": [0.967609], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_xgbregressor_model_predict( + penguins_xgbregressor_model: bigframes.ml.ensemble.XGBRegressor, new_penguins_df +): + result = penguins_xgbregressor_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_body_mass_g": ["4293.1538089", "3410.0271", "3357.944"]}, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + result.sort_index(), + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + ) + + +def test_to_gbq_saved_xgbregressor_model_scores( + penguins_xgbregressor_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_xgbregressor_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [109.016973], + "mean_squared_error": [20867.299758], + "mean_squared_log_error": [0.00135], + "median_absolute_error": [86.490234], + "r2_score": [0.967458], + "explained_variance": [0.967504], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_to_xgbregressor_model_gbq_replace(penguins_xgbregressor_model, dataset_id): + penguins_xgbregressor_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_xgbregressor_model.to_gbq(f"{dataset_id}.test_penguins_model") + + +def test_xgbclassifier_model_score( + penguins_xgbclassifier_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = penguins_xgbclassifier_model.score(test_X, test_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + +def test_xgbclassifier_model_predict( + penguins_xgbclassifier_model: bigframes.ml.ensemble.XGBClassifier, new_penguins_df +): + result = penguins_xgbclassifier_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, + dtype="string[pyarrow]", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + result.sort_index(), + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + ) + + +def test_to_gbq_saved_xgbclassifier_model_scores( + penguins_xgbclassifier_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_xgbclassifier_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "precision": [1.0], + "recall": [1.0], + "accuracy": [1.0], + "f1_score": [1.0], + "log_loss": [0.331442], + "roc_auc": [1.0], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + assert saved_model.max_depth == 6 + assert saved_model.max_iterations == 20 + + +def test_to_xgbclassifier_model_gbq_replace(penguins_xgbclassifier_model, dataset_id): + penguins_xgbclassifier_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_xgbclassifier_model.to_gbq(f"{dataset_id}.test_penguins_model") + + +def test_randomforestregressor_model_score( + penguins_randomforest_regressor_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = penguins_randomforest_regressor_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [317.031042], + "mean_squared_error": [159713.053504], + "mean_squared_log_error": [0.008449], + "median_absolute_error": [258.385742], + "r2_score": [0.752698], + "explained_variance": [0.756173], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigFramese + check_index_type=False, + ) + + +def test_randomforestregressor_model_predict( + penguins_randomforest_regressor_model: bigframes.ml.ensemble.RandomForestRegressor, + new_penguins_df, +): + result = penguins_randomforest_regressor_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_body_mass_g": ["3897.341797", "3458.385742", "3458.385742"]}, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + result.sort_index(), + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + ) + + +def test_to_gbq_saved_randomforestregressor_model_scores( + penguins_randomforest_regressor_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_randomforest_regressor_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [319.239235], + "mean_squared_error": [161913.126651], + "mean_squared_log_error": [0.008611], + "median_absolute_error": [266.614258], + "r2_score": [0.747504], + "explained_variance": [0.750358], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigFramese + check_index_type=False, + ) + + +def test_to_randomforestregressor_model_gbq_replace( + penguins_randomforest_regressor_model, dataset_id +): + penguins_randomforest_regressor_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_randomforest_regressor_model.to_gbq( + f"{dataset_id}.test_penguins_model" + ) + + +def test_randomforestclassifier_model_score( + penguins_randomforest_classifier_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = penguins_randomforest_classifier_model.score(test_X, test_y).compute() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + +def test_randomforestclassifier_model_predict( + penguins_randomforest_classifier_model: bigframes.ml.ensemble.RandomForestClassifier, + new_penguins_df, +): + result = penguins_randomforest_classifier_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, + dtype="string[pyarrow]", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + result.sort_index(), + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + ) + + +def test_to_gbq_saved_randomforestclassifier_model_scores( + penguins_randomforest_classifier_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_randomforest_classifier_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "precision": [0.636746], + "recall": [0.638636], + "accuracy": [0.95509], + "f1_score": [0.637688], + "log_loss": [0.886307], + "roc_auc": [0.966543], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_to_randomforestclassifier_model_gbq_replace( + penguins_randomforest_classifier_model, dataset_id +): + penguins_randomforest_classifier_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_randomforest_classifier_model.to_gbq( + f"{dataset_id}.test_penguins_model" + ) diff --git a/tests/system/small/ml/test_forecasting.py b/tests/system/small/ml/test_forecasting.py new file mode 100644 index 00000000000..153299f5989 --- /dev/null +++ b/tests/system/small/ml/test_forecasting.py @@ -0,0 +1,44 @@ +# 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. + +from datetime import datetime + +import pandas as pd +import pyarrow as pa +import pytz + + +def test_model_predict(time_series_arima_plus_model): + utc = pytz.utc + predictions = time_series_arima_plus_model.predict().compute() + expected = pd.DataFrame( + { + "forecast_timestamp": [ + datetime(2017, 8, 2, tzinfo=utc), + datetime(2017, 8, 3, tzinfo=utc), + datetime(2017, 8, 4, tzinfo=utc), + ], + "forecast_value": [2724.472284, 2593.368389, 2353.613034], + } + ) + expected["forecast_value"] = expected["forecast_value"].astype(pd.Float64Dtype()) + expected["forecast_timestamp"] = expected["forecast_timestamp"].astype( + pd.ArrowDtype(pa.timestamp("us", tz="UTC")) + ) + pd.testing.assert_frame_equal( + predictions, + expected, + rtol=0.1, + check_index_type=False, + ) diff --git a/tests/system/small/ml/test_imported.py b/tests/system/small/ml/test_imported.py new file mode 100644 index 00000000000..b2a9abc46e7 --- /dev/null +++ b/tests/system/small/ml/test_imported.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 numpy as np +import pandas as pd + + +def test_tensorflow_create_model(imported_tensorflow_model): + # Model creation doesn't return error + assert imported_tensorflow_model is not None + + +def test_tensorflow_model_predict(imported_tensorflow_model, llm_text_df): + df = llm_text_df.rename(columns={"prompt": "input"}) + result = imported_tensorflow_model.predict(df).compute() + # The values are non-human-readable. As they are a dense layer of Neural Network. + # And since it is pretrained and imported, the model is a opaque-box. + # We may want to switch to better test model and cases. + value = np.array( + [9.375373792863684e-07, 0.00015779426030348986, 0.9998412132263184] + ) + expected = pd.DataFrame( + { + "dense_1": [value, value, value], + }, + ) + expected.set_index(expected.index.astype("Int64"), inplace=True) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + atol=0.1, + ) + + +def test_onnx_create_model(imported_onnx_model): + # Model creation doesn't return error + assert imported_onnx_model is not None + + +def test_onnx_model_predict(imported_onnx_model, onnx_iris_df): + result = imported_onnx_model.predict(onnx_iris_df).compute() + value1 = np.array([0.9999993443489075, 0.0, 0.0]) + value2 = np.array([0.0, 0.0, 0.9999993443489075]) + expected = pd.DataFrame( + { + "label": pd.array([0, 0, 2]).astype("Int64"), + "probabilities": [value1, value1, value2], + }, + index=pd.Index([0, 1, 2], dtype="Int64"), + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + atol=0.1, + ) diff --git a/tests/system/small/ml/test_linear_model.py b/tests/system/small/ml/test_linear_model.py new file mode 100644 index 00000000000..d5d6957eec3 --- /dev/null +++ b/tests/system/small/ml/test_linear_model.py @@ -0,0 +1,209 @@ +# 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 pandas +import pytest + + +def test_linear_reg_model_score(penguins_linear_model, penguins_df_default_index): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + test_y = df[["body_mass_g"]] + result = penguins_linear_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [225.817334], + "mean_squared_error": [80540.705944], + "mean_squared_log_error": [0.004972], + "median_absolute_error": [173.080816], + "r2_score": [0.87529], + "explained_variance": [0.87529], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_linear_reg_model_predict(penguins_linear_model, new_penguins_df): + predictions = penguins_linear_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, + dtype="Float64", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + predictions.sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) + + +def test_to_gbq_saved_linear_reg_model_scores( + penguins_linear_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_linear_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + test_y = df[["body_mass_g"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "mean_absolute_error": [227.01223], + "mean_squared_error": [81838.159892], + "mean_squared_log_error": [0.00507], + "median_absolute_error": [173.080816], + "r2_score": [0.872377], + "explained_variance": [0.872377], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_to_gbq_replace(penguins_linear_model, dataset_id): + penguins_linear_model.to_gbq(f"{dataset_id}.test_penguins_model", replace=True) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_linear_model.to_gbq(f"{dataset_id}.test_penguins_model") + + +def test_logistic_model_score(penguins_logistic_model, penguins_df_default_index): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = penguins_logistic_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "precision": [0.616753], + "recall": [0.618615], + "accuracy": [0.92515], + "f1_score": [0.617681], + "log_loss": [1.498832], + "roc_auc": [0.975807], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_logsitic_model_predict(penguins_logistic_model, new_penguins_df): + predictions = penguins_logistic_model.predict(new_penguins_df).compute() + expected = pandas.DataFrame( + {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, + dtype="string[pyarrow]", + index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pandas.testing.assert_frame_equal( + predictions.sort_index(), + expected, + check_exact=False, + rtol=0.1, + ) + + +def test_to_gbq_saved_logsitic_model_score( + penguins_logistic_model, dataset_id, penguins_df_default_index +): + saved_model = penguins_logistic_model.to_gbq( + f"{dataset_id}.test_penguins_model", replace=True + ) + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df[["sex"]] + result = saved_model.score(test_X, test_y).compute() + expected = pandas.DataFrame( + { + "precision": [0.616753], + "recall": [0.618615], + "accuracy": [0.92515], + "f1_score": [0.617681], + "log_loss": [1.498832], + "roc_auc": [0.975807], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_to_logistic_model_gbq_replace(penguins_logistic_model, dataset_id): + penguins_logistic_model.to_gbq(f"{dataset_id}.test_penguins_model", replace=True) + with pytest.raises(google.api_core.exceptions.Conflict): + penguins_logistic_model.to_gbq(f"{dataset_id}.test_penguins_model") diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py new file mode 100644 index 00000000000..b579e754e5a --- /dev/null +++ b/tests/system/small/ml/test_llm.py @@ -0,0 +1,72 @@ +# 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. + +from unittest import TestCase + +import numpy as np + + +def test_create_text_generator_model(palm2_text_generator_model): + # Model creation doesn't return error + assert palm2_text_generator_model is not None + + +def test_text_generator_predict_default_params_success( + palm2_text_generator_model, llm_text_df +): + df = palm2_text_generator_model.predict(llm_text_df).compute() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +def test_text_generator_predict_arbitrary_col_label_success( + palm2_text_generator_model, llm_text_df +): + llm_text_df = llm_text_df.rename(columns={"prompt": "arbitrary"}) + df = palm2_text_generator_model.predict(llm_text_df).compute() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +def test_text_generator_predict_with_params_success( + palm2_text_generator_model, llm_text_df +): + df = palm2_text_generator_model.predict( + llm_text_df, temperature=0.5, max_output_tokens=100, top_k=20, top_p=0.5 + ).compute() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +def test_create_embedding_generator_model(palm2_embedding_generator_model): + # Model creation doesn't return error + assert palm2_embedding_generator_model is not None + + +def test_embedding_generator_predict_success( + palm2_embedding_generator_model, llm_text_df +): + df = palm2_embedding_generator_model.predict(llm_text_df).compute() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_embed_text_embedding" in df.columns + series = df["ml_embed_text_embedding"] + value = series[0] + assert isinstance(value, np.ndarray) + assert value.size == 768 diff --git a/tests/system/small/ml/test_metrics.py b/tests/system/small/ml/test_metrics.py new file mode 100644 index 00000000000..5b4c486b570 --- /dev/null +++ b/tests/system/small/ml/test_metrics.py @@ -0,0 +1,615 @@ +# 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 math + +import numpy as np +import pandas as pd +import pytest +import sklearn.metrics as sklearn_metrics # type: ignore + +import bigframes.ml.metrics + + +def test_r2_score_perfect_fit(session): + pd_df = pd.DataFrame({"y_true": [1, 7, 3, 2, 5], "y_pred": [1, 7, 3, 2, 5]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred"]]) == 1.0 + + +def test_r2_score_bad_fit(session): + pd_df = pd.DataFrame({"y_true": [1, 2, 3, 4, 5], "y_pred": [5, 4, 3, 2, 1]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred"]]) == -3.0 + + +def test_r2_score_force_finite(session): + pd_df = pd.DataFrame( + { + "y_true": [1, 1, 1, 1, 1], + "y_pred_1": [5, 4, 3, 2, 1], + "y_pred_2": [1, 1, 1, 1, 1], + } + ) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.r2_score( + df[["y_true"]], df[["y_pred_1"]], force_finite=False + ) == float("-inf") + assert bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred_1"]]) == 0.0 + assert math.isnan( + bigframes.ml.metrics.r2_score( + df[["y_true"]], df[["y_pred_2"]], force_finite=False + ) + ) + assert bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred_2"]]) == 1.0 + + +def test_r2_score_ok_fit_matches_sklearn(session): + pd_df = pd.DataFrame({"y_true": [1, 2, 3, 4, 5], "y_pred": [2, 3, 4, 3, 6]}) + + df = session.read_pandas(pd_df) + bf_result = bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred"]]) + sklearn_result = sklearn_metrics.r2_score(pd_df[["y_true"]], pd_df[["y_pred"]]) + assert math.isclose(bf_result, sklearn_result) + + +def test_accuracy_score_perfect_fit(session): + pd_df = pd.DataFrame({"y_true": [1, 7, 3, 2, 5], "y_pred": [1, 7, 3, 2, 5]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.accuracy_score(df[["y_true"]], df[["y_pred"]]) == 1.0 + + +def test_accuracy_score_bad_fit(session): + pd_df = pd.DataFrame({"y_true": [0, 2, 1, 3, 4], "y_pred": [0, 1, 2, 3, 4]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.accuracy_score(df[["y_true"]], df[["y_pred"]]) == 0.6 + + +def test_accuracy_score_not_normailze(session): + pd_df = pd.DataFrame({"y_true": [0, 2, 1, 3, 4], "y_pred": [0, 1, 2, 3, 4]}) + + df = session.read_pandas(pd_df) + assert ( + bigframes.ml.metrics.accuracy_score( + df[["y_true"]], df[["y_pred"]], normalize=False + ) + == 3 + ) + + +@pytest.mark.skipif(sklearn_metrics is None, reason="requires sklearn") +def test_accuracy_score_fit_matches_sklearn(session): + pd_df = pd.DataFrame({"y_true": [1, 2, 3, 4, 5], "y_pred": [2, 3, 4, 3, 6]}) + + df = session.read_pandas(pd_df) + bf_result = bigframes.ml.metrics.accuracy_score(df[["y_true"]], df[["y_pred"]]) + sklearn_result = sklearn_metrics.accuracy_score( + pd_df[["y_true"]], pd_df[["y_pred"]] + ) + assert math.isclose(bf_result, sklearn_result) + + +def test_roc_curve_binary_classification_prediction_returns_expected(session): + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + } + ) + + df = session.read_pandas(pd_df) + fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( + df[["y_true"]], df[["y_score"]], drop_intermediate=False + ) + + pd_fpr = fpr.compute() + pd_tpr = tpr.compute() + pd_thresholds = thresholds.compute() + + pd.testing.assert_series_equal( + # skip testing the first value, as it is redundant and inconsistent across sklearn versions + pd_thresholds[1:], + pd.Series( + [0.9, 0.8, 0.65, 0.6, 0.5, 0.45, 0.4, 0.35, 0.3, 0.1], + dtype="Float64", + name="thresholds", + ), + check_index=False, + ) + pd.testing.assert_series_equal( + pd_fpr, + pd.Series( + [0.0, 0.0, 0.0, 0.25, 0.25, 0.5, 0.5, 0.75, 0.75, 0.75, 1.0], + dtype="Float64", + name="fpr", + ), + check_index_type=False, + ) + pd.testing.assert_series_equal( + pd_tpr, + pd.Series( + [ + 0.0, + 0.16666667, + 0.33333333, + 0.33333333, + 0.5, + 0.5, + 0.66666667, + 0.66666667, + 0.83333333, + 1.0, + 1.0, + ], + dtype="Float64", + name="tpr", + ), + check_index_type=False, + ) + + +def test_roc_curve_binary_classification_prediction_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + } + ) + + df = session.read_pandas(pd_df) + fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( + df[["y_true"]], df[["y_score"]], drop_intermediate=False + ) + expected_fpr, expected_tpr, expected_thresholds = sklearn_metrics.roc_curve( + pd_df[["y_true"]], pd_df[["y_score"]], drop_intermediate=False + ) + + # sklearn returns float64 np arrays + np_fpr = fpr.compute().astype("float64").array + np_tpr = tpr.compute().astype("float64").array + np_thresholds = thresholds.compute().astype("float64").array + + np.testing.assert_array_equal( + # skip testing the first value, as it is redundant and inconsistent across sklearn versions + np_thresholds[1:], + expected_thresholds[1:], + ) + np.testing.assert_array_equal( + np_fpr, + expected_fpr, + ) + np.testing.assert_array_equal( + np_tpr, + expected_tpr, + ) + + +def test_roc_curve_binary_classification_decision_returns_expected(session): + # Instead of operating on probabilities, assume a 70% decision threshold + # has been applied, and operate on the final output + y_score = [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45] + decisions_70pct = [1 if s > 0.7 else 0 for s in y_score] + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": decisions_70pct, + } + ) + + df = session.read_pandas(pd_df) + fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( + df[["y_true"]], df[["y_score"]], drop_intermediate=False + ) + + pd_fpr = fpr.compute() + pd_tpr = tpr.compute() + pd_thresholds = thresholds.compute() + + pd.testing.assert_series_equal( + # skip testing the first value, as it is redundant and inconsistent across sklearn versions + pd_thresholds[1:], + pd.Series( + [1.0, 0.0], + dtype="Float64", + name="thresholds", + ), + check_index=False, + ) + pd.testing.assert_series_equal( + pd_fpr, + pd.Series( + [0.0, 0.0, 1.0], + dtype="Float64", + name="fpr", + ), + check_index_type=False, + ) + pd.testing.assert_series_equal( + pd_tpr, + pd.Series( + [ + 0.0, + 0.33333333, + 1.0, + ], + dtype="Float64", + name="tpr", + ), + check_index_type=False, + ) + + +def test_roc_curve_binary_classification_decision_matches_sklearn(session): + # Instead of operating on probabilities, assume a 70% decision threshold + # has been applied, and operate on the final output + y_score = [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45] + decisions_70pct = [1 if s > 0.7 else 0 for s in y_score] + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": decisions_70pct, + } + ) + + df = session.read_pandas(pd_df) + fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( + df[["y_true"]], df[["y_score"]], drop_intermediate=False + ) + expected_fpr, expected_tpr, expected_thresholds = sklearn_metrics.roc_curve( + pd_df[["y_true"]], pd_df[["y_score"]], drop_intermediate=False + ) + + # sklearn returns float64 np arrays + np_fpr = fpr.compute().astype("float64").array + np_tpr = tpr.compute().astype("float64").array + np_thresholds = thresholds.compute().astype("float64").array + + np.testing.assert_array_equal( + # skip testing the first value, as it is redundant and inconsistent across sklearn versions + np_thresholds[1:], + expected_thresholds[1:], + ) + np.testing.assert_array_equal( + np_fpr, + expected_fpr, + ) + np.testing.assert_array_equal( + np_tpr, + expected_tpr, + ) + + +def test_roc_auc_score_returns_expected(session): + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + } + ) + + df = session.read_pandas(pd_df) + score = bigframes.ml.metrics.roc_auc_score(df[["y_true"]], df[["y_score"]]) + + assert score == 0.625 + + +def test_roc_auc_score_returns_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + } + ) + + df = session.read_pandas(pd_df) + score = bigframes.ml.metrics.roc_auc_score(df[["y_true"]], df[["y_score"]]) + expected_score = sklearn_metrics.roc_auc_score( + pd_df[["y_true"]], pd_df[["y_score"]] + ) + + assert score == expected_score + + +def test_auc_invalid_x_size(session): + pd_df = pd.DataFrame({"x": [0], "y": [0]}) + df = session.read_pandas(pd_df) + with pytest.raises(ValueError): + bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) + + +def test_auc_nondecreasing_x(session): + pd_df = pd.DataFrame({"x": [0, 0, 0.5, 0.5, 1], "y": [0, 0.5, 0.5, 1, 1]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) == 0.75 + + +def test_auc_nonincreasing_x(session): + pd_df = pd.DataFrame({"x": [0, 0, -0.5, -0.5, -1], "y": [0, 0.5, 0.5, 1, 1]}) + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) == 0.75 + + +def test_auc_nonincreasing_x_negative(session): + pd_df = pd.DataFrame({"x": [0, 0, -0.5, -0.5, -1], "y": [0, -0.5, -0.5, -1, -1]}) + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) == -0.75 + + +def test_confusion_matrix(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + confusion_matrix = bigframes.ml.metrics.confusion_matrix( + df[["y_true"]], df[["y_pred"]] + ) + expected_pd_df = pd.DataFrame( + { + 0: [2, 0, 1], + 1: [0, 0, 0], + 2: [0, 1, 2], + } + ).astype("int64") + pd.testing.assert_frame_equal( + confusion_matrix, expected_pd_df, check_index_type=False + ) + + +def test_confusion_matrix_column_index(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 3, 3, 3, 4, 1], + "y_pred": [4, 1, 2, 2, 4, 1], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + confusion_matrix = bigframes.ml.metrics.confusion_matrix( + df[["y_true"]], df[["y_pred"]] + ) + expected_pd_df = ( + pd.DataFrame( + {1: [1, 0, 1, 0], 2: [0, 0, 2, 0], 3: [0, 0, 0, 0], 4: [0, 1, 0, 1]} + ) + .astype("int64") + .set_index([pd.Index([1, 2, 3, 4])]) + ) + pd.testing.assert_frame_equal( + confusion_matrix, expected_pd_df, check_index_type=False + ) + + +def test_confusion_matrix_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 3, 3, 3, 4, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + confusion_matrix = bigframes.ml.metrics.confusion_matrix( + df[["y_true"]], df[["y_pred"]] + ) + expected_confusion_matrix = sklearn_metrics.confusion_matrix( + pd_df[["y_true"]], pd_df[["y_pred"]] + ) + expected_pd_df = pd.DataFrame(expected_confusion_matrix) + pd.testing.assert_frame_equal( + confusion_matrix, expected_pd_df, check_index_type=False + ) + + +def test_confusion_matrix_str_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": ["cat", "ant", "cat", "cat", "ant", "bird"], + "y_pred": ["ant", "ant", "cat", "cat", "ant", "cat"], + } + ).astype("str") + df = session.read_pandas(pd_df) + confusion_matrix = bigframes.ml.metrics.confusion_matrix( + df[["y_true"]], df[["y_pred"]] + ) + expected_confusion_matrix = sklearn_metrics.confusion_matrix( + pd_df[["y_true"]], pd_df[["y_pred"]] + ) + expected_pd_df = pd.DataFrame(expected_confusion_matrix).set_index( + [pd.Index(["ant", "bird", "cat"])] + ) + expected_pd_df.columns = pd.Index(["ant", "bird", "cat"]) + pd.testing.assert_frame_equal( + confusion_matrix, expected_pd_df, check_index_type=False + ) + + +def test_recall_score(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + recall = bigframes.ml.metrics.recall_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = [1.000000, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_recall = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal(recall, expected_recall, check_index_type=False) + + +def test_recall_score_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + recall = bigframes.ml.metrics.recall_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.recall_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = [0, 1, 2] + expected_recall = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal(recall, expected_recall, check_index_type=False) + + +def test_recall_score_str_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": ["cat", "ant", "cat", "cat", "ant", "bird"], + "y_pred": ["ant", "ant", "cat", "cat", "ant", "cat"], + } + ).astype("str") + df = session.read_pandas(pd_df) + recall = bigframes.ml.metrics.recall_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.recall_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = ["ant", "bird", "cat"] + expected_recall = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal(recall, expected_recall, check_index_type=False) + + +def test_precision_score(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + precision_score = bigframes.ml.metrics.precision_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = [0.666667, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_precision = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal( + precision_score, expected_precision, check_index_type=False + ) + + +def test_precision_score_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + precision_score = bigframes.ml.metrics.precision_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.precision_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = [0, 1, 2] + expected_precision = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal( + precision_score, expected_precision, check_index_type=False + ) + + +def test_precision_score_str_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": ["cat", "ant", "cat", "cat", "ant", "bird"], + "y_pred": ["ant", "ant", "cat", "cat", "ant", "cat"], + } + ).astype("str") + df = session.read_pandas(pd_df) + precision_score = bigframes.ml.metrics.precision_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.precision_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = ["ant", "bird", "cat"] + expected_precision = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal( + precision_score, expected_precision, check_index_type=False + ) + + +def test_f1_score(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + f1_score = bigframes.ml.metrics.f1_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = [0.8, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_f1 = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal(f1_score, expected_f1, check_index_type=False) + + +def test_f1_score_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + f1_score = bigframes.ml.metrics.f1_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.f1_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = [0, 1, 2] + expected_f1 = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal(f1_score, expected_f1, check_index_type=False) + + +def test_f1_score_str_matches_sklearn(session): + pd_df = pd.DataFrame( + { + "y_true": ["cat", "ant", "cat", "cat", "ant", "bird"], + "y_pred": ["ant", "ant", "cat", "cat", "ant", "cat"], + } + ).astype("str") + df = session.read_pandas(pd_df) + f1_score = bigframes.ml.metrics.f1_score( + df[["y_true"]], df[["y_pred"]], average=None + ) + expected_values = sklearn_metrics.f1_score( + pd_df[["y_true"]], pd_df[["y_pred"]], average=None + ) + expected_index = ["ant", "bird", "cat"] + expected_f1 = pd.Series(expected_values, index=expected_index) + pd.testing.assert_series_equal(f1_score, expected_f1, check_index_type=False) diff --git a/tests/system/small/ml/test_model_selection.py b/tests/system/small/ml/test_model_selection.py new file mode 100644 index 00000000000..add455ab9ac --- /dev/null +++ b/tests/system/small/ml/test_model_selection.py @@ -0,0 +1,214 @@ +# 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 pandas as pd +import pytest + +import bigframes.ml.model_selection + + +def test_train_test_split_default_correct_shape(penguins_df_default_index): + X = penguins_df_default_index[ + [ + "species", + "island", + "culmen_length_mm", + ] + ] + y = penguins_df_default_index[["body_mass_g"]] + X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( + X, y + ) + + # even though the default seed is random, it should always result in this shape + assert X_train.shape == (258, 3) + assert X_test.shape == (86, 3) + assert y_train.shape == (258, 1) + assert y_test.shape == (86, 1) + + +def test_train_test_double_split_correct_shape(penguins_df_default_index): + X = penguins_df_default_index[ + [ + "species", + "island", + "culmen_length_mm", + ] + ] + y = penguins_df_default_index[["body_mass_g"]] + X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( + X, y, test_size=0.2, train_size=0.4 + ) + + # should have 20% in test, 40% in train, 40% dropped + assert X_train.shape == (138, 3) + assert X_test.shape == (69, 3) + assert y_train.shape == (138, 1) + assert y_test.shape == (69, 1) + + +def test_train_test_three_dataframes_correct_shape(penguins_df_default_index): + A = penguins_df_default_index[ + [ + "species", + "culmen_length_mm", + ] + ] + B = penguins_df_default_index[ + [ + "island", + ] + ] + C = penguins_df_default_index[["culmen_depth_mm", "body_mass_g"]] + ( + A_train, + A_test, + B_train, + B_test, + C_train, + C_test, + ) = bigframes.ml.model_selection.train_test_split(A, B, C) + + assert A_train.shape == (258, 2) + assert A_test.shape == (86, 2) + assert B_train.shape == (258, 1) + assert B_test.shape == (86, 1) + assert C_train.shape == (258, 2) + assert C_test.shape == (86, 2) + + +def test_train_test_split_seeded_correct_rows( + session, penguins_pandas_df_default_index +): + # Note that we're using `penguins_pandas_df_default_index` as this test depends + # on a stable row order being present end to end + # filter down to the chunkiest penguins, to keep our test code a reasonable size + all_data = penguins_pandas_df_default_index[ + penguins_pandas_df_default_index.body_mass_g > 5500 + ] + + # Note that bigframes loses the index if it doesn't have a name + all_data.index.name = "rowindex" + + df = session.read_pandas(all_data) + + X = df[ + [ + "species", + "island", + "culmen_length_mm", + ] + ] + y = df[["body_mass_g"]] + X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( + X, y, random_state=42 + ) + + X_train = X_train.to_pandas().sort_index() + X_test = X_test.to_pandas().sort_index() + y_train = y_train.to_pandas().sort_index() + y_test = y_test.to_pandas().sort_index() + + train_index = pd.Index( + [ + 144, + 146, + 148, + 168, + 183, + 186, + 217, + 225, + 226, + 237, + 244, + 245, + 260, + 262, + 263, + 266, + 268, + 269, + 289, + 290, + 291, + ], + dtype="Int64", + name="rowindex", + ) + test_index = pd.Index( + [161, 221, 240, 257, 264, 267, 278], dtype="Int64", name="rowindex" + ) + + all_data.index.name = "_" + pd.testing.assert_frame_equal( + X_train, + all_data[ + [ + "species", + "island", + "culmen_length_mm", + ] + ].loc[train_index], + ) + pd.testing.assert_frame_equal( + X_test, + all_data[ + [ + "species", + "island", + "culmen_length_mm", + ] + ].loc[test_index], + ) + pd.testing.assert_frame_equal( + y_train, + all_data[ + [ + "body_mass_g", + ] + ].loc[train_index], + ) + pd.testing.assert_frame_equal( + y_test, + all_data[ + [ + "body_mass_g", + ] + ].loc[test_index], + ) + + +@pytest.mark.parametrize( + ("train_size", "test_size"), + [ + (0.0, 0.5), + (-0.5, 0.7), + (0.5, 1.2), + (0.6, 0.6), + ], +) +def test_train_test_split_value_error(penguins_df_default_index, train_size, test_size): + X = penguins_df_default_index[ + [ + "species", + "island", + "culmen_length_mm", + ] + ] + y = penguins_df_default_index[["body_mass_g"]] + with pytest.raises(ValueError): + bigframes.ml.model_selection.train_test_split( + X, y, train_size=train_size, test_size=test_size + ) diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py new file mode 100644 index 00000000000..f31b93b4cc2 --- /dev/null +++ b/tests/system/small/ml/test_preprocessing.py @@ -0,0 +1,89 @@ +# 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 math + +import pandas as pd + +import bigframes.ml.preprocessing + + +def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): + # TODO(bmil): add a second test that compares output to sklearn.preprocessing.StandardScaler + scaler = bigframes.ml.preprocessing.StandardScaler() + scaler.fit( + penguins_df_default_index[ + "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm" + ] + ) + + result = scaler.transform( + penguins_df_default_index[ + "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm" + ] + ).to_pandas() + + # If standard-scaled correctly, mean should be 0.0 and standard deviation 1.0 + for column in result.columns: + assert math.isclose(result[column].mean(), 0.0, abs_tol=1e-3) + assert math.isclose(result[column].std(), 1.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "scaled_culmen_depth_mm": [0.8349, 0.02473, 0.4805], + "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], + "scaled_flipper_length_mm": [-0.3495, -1.4163, -0.9185], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_one_hot_encoder_encodes(penguins_df_default_index, new_penguins_df): + encoder = bigframes.ml.preprocessing.OneHotEncoder() + encoder.fit(penguins_df_default_index["species", "sex"]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "onehotencoded_sex": [ + [{"index": 2, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + ], + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + }, + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) diff --git a/tests/system/small/ml/test_register.py b/tests/system/small/ml/test_register.py new file mode 100644 index 00000000000..ad378d35084 --- /dev/null +++ b/tests/system/small/ml/test_register.py @@ -0,0 +1,87 @@ +# 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. + +from typing import cast + +from bigframes.ml import core, imported, linear_model, llm + + +def test_linear_reg_register( + ephemera_penguins_linear_model: linear_model.LinearRegression, +): + model = ephemera_penguins_linear_model + model.register() + + model_name = "bigframes_" + cast( + str, cast(core.BqmlModel, model._bqml_model).model.model_id + ) + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert ( + model_name + in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ + "vertexAiModelId" + ] + ) + + +def test_linear_reg_register_with_params( + ephemera_penguins_linear_model: linear_model.LinearRegression, +): + model_name = "bigframes_system_test_linear_reg_model" + model = ephemera_penguins_linear_model + model.register(model_name) + + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert ( + model_name + in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ + "vertexAiModelId" + ] + ) + + +def test_palm2_text_generator_register( + ephemera_palm2_text_generator_model: llm.PaLM2TextGenerator, +): + model = ephemera_palm2_text_generator_model + model.register() + + model_name = "bigframes_" + cast( + str, cast(core.BqmlModel, model._bqml_model).model.model_id + ) + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert ( + model_name + in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ + "vertexAiModelId" + ] + ) + + +def test_imported_tensorflow_register( + ephemera_imported_tensorflow_model: imported.TensorFlowModel, +): + model = ephemera_imported_tensorflow_model + model.register() + + model_name = "bigframes_" + cast( + str, cast(core.BqmlModel, model._bqml_model).model.model_id + ) + # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. + assert ( + model_name + in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ + "vertexAiModelId" + ] + ) diff --git a/tests/system/small/operations/__init__.py b/tests/system/small/operations/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/system/small/operations/__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/tests/system/small/operations/test_datetimes.py b/tests/system/small/operations/test_datetimes.py new file mode 100644 index 00000000000..25d3eb69adb --- /dev/null +++ b/tests/system/small/operations/test_datetimes.py @@ -0,0 +1,198 @@ +# 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 pandas as pd +import pytest + +import bigframes.series +from tests.system.utils import assert_series_equal_ignoring_order + +DATETIME_COL_NAMES = [("datetime_col",), ("timestamp_col",)] + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_day(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.day.compute() + pd_result = scalars_pandas_df[col_name].dt.day + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_date(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.date.compute() + pd_result = scalars_pandas_df[col_name].dt.date + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_dayofweek(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.dayofweek.compute() + pd_result = scalars_pandas_df[col_name].dt.dayofweek + + assert_series_equal_ignoring_order(pd_result, bf_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_hour(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.hour.compute() + pd_result = scalars_pandas_df[col_name].dt.hour + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_minute(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.minute.compute() + pd_result = scalars_pandas_df[col_name].dt.minute + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_month(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.month.compute() + pd_result = scalars_pandas_df[col_name].dt.month + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_quarter(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.quarter.compute() + pd_result = scalars_pandas_df[col_name].dt.quarter + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_second(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.second.compute() + pd_result = scalars_pandas_df[col_name].dt.second + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_time(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.time.compute() + pd_result = scalars_pandas_df[col_name].dt.time + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + DATETIME_COL_NAMES, +) +def test_year(scalars_dfs, col_name): + if pd.__version__.startswith("1."): + pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.dt.year.compute() + pd_result = scalars_pandas_df[col_name].dt.year + + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) diff --git a/tests/system/small/operations/test_strings.py b/tests/system/small/operations/test_strings.py new file mode 100644 index 00000000000..76736e3cddd --- /dev/null +++ b/tests/system/small/operations/test_strings.py @@ -0,0 +1,230 @@ +# 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 pandas as pd +import pytest + +import bigframes.series + +from ...utils import assert_series_equal_ignoring_order + + +def test_find(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.find("W").compute() + pd_result = scalars_pandas_df[col_name].str.find("W") + + # One of type mismatches to be documented. Here, the `bf_result.dtype` is `Int64` but + # the `pd_result.dtype` is `float64`: https://github.com/pandas-dev/pandas/issues/51948 + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +def test_len(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.len().compute() + pd_result = scalars_pandas_df[col_name].str.len() + + # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is `Int64` but + # the `pd_result.dtype` is `float64`: https://github.com/pandas-dev/pandas/issues/51948 + assert_series_equal_ignoring_order( + pd_result.astype(pd.Int64Dtype()), + bf_result, + ) + + +def test_lower(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.lower().compute() + pd_result = scalars_pandas_df[col_name].str.lower() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_reverse(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.reverse().compute() + pd_result = scalars_pandas_df[col_name].copy() + for i in pd_result.index: + cell = pd_result.loc[i] + if pd.isna(cell): + pd_result.loc[i] = None + else: + pd_result.loc[i] = cell[::-1] + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ["start", "stop"], [(0, 1), (3, 5), (100, 101), (None, 1), (0, 12), (0, None)] +) +def test_slice(scalars_dfs, start, stop): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.slice(start, stop).compute() + pd_series = scalars_pandas_df[col_name] + pd_result = pd_series.str.slice(start, stop) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_strip(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.strip().compute() + pd_result = scalars_pandas_df[col_name].str.strip() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_upper(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.upper().compute() + pd_result = scalars_pandas_df[col_name].str.upper() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_isnumeric(session): + pandas_df = pd.DataFrame( + { + "numeric_string_col": [ + "٠١٢٣٤٥٦٧٨٩", + "", + "0", + "字", + "五", + "0123456789", + pd.NA, + "abc 123 mixed letters and numbers", + "no numbers here", + "123a", + "23!", + " 45", + "a45", + ] + } + ) + + df = session.read_pandas(pandas_df) + + pd_result = pandas_df.numeric_string_col.str.isnumeric() + bf_result = df.numeric_string_col.str.isnumeric().compute() + + assert_series_equal_ignoring_order( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_rstrip(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.rstrip().compute() + pd_result = scalars_pandas_df[col_name].str.rstrip() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_lstrip(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.lstrip().compute() + pd_result = scalars_pandas_df[col_name].str.lstrip() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize(["repeats"], [(5,), (0,), (1,)]) +def test_repeat(scalars_dfs, repeats): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.repeat(repeats).compute() + pd_result = scalars_pandas_df[col_name].str.repeat(repeats) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_capitalize(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.capitalize().compute() + pd_result = scalars_pandas_df[col_name].str.capitalize() + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_cat_with_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_filter: bigframes.series.Series = scalars_df["bool_col"] + bf_left: bigframes.series.Series = scalars_df[col_name][bf_filter] + bf_right: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_left.str.cat(others=bf_right).compute() + pd_filter = scalars_pandas_df["bool_col"] + pd_left = scalars_pandas_df[col_name][pd_filter] + pd_right = scalars_pandas_df[col_name] + pd_result = pd_left.str.cat(others=pd_right) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py new file mode 100644 index 00000000000..1833c674fdd --- /dev/null +++ b/tests/system/small/test_dataframe.py @@ -0,0 +1,1694 @@ +# 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 operator + +import geopandas as gpd # type: ignore +import numpy as np +import pandas as pd +import pandas.testing +import pyarrow as pa # type: ignore +import pytest + +import bigframes +import bigframes._config.display_options as display_options +import bigframes.dataframe as dataframe +from tests.system.utils import ( + assert_pandas_df_equal_ignore_ordering, + assert_series_equal_ignoring_order, +) + + +def test_df_construct_copy(scalars_dfs): + columns = ["int64_col", "string_col", "float64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = dataframe.DataFrame(scalars_df, columns=columns).compute() + pd_result = pd.DataFrame(scalars_pandas_df, columns=columns) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_construct_pandas(scalars_dfs): + columns = ["int64_too", "int64_col", "float64_col", "bool_col", "string_col"] + _, scalars_pandas_df = scalars_dfs + bf_result = dataframe.DataFrame(scalars_pandas_df, columns=columns).compute() + pd_result = pd.DataFrame(scalars_pandas_df, columns=columns) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_construct_pandas_set_dtype(scalars_dfs): + columns = [ + "int64_too", + "int64_col", + "float64_col", + "bool_col", + ] + _, scalars_pandas_df = scalars_dfs + bf_result = dataframe.DataFrame( + scalars_pandas_df, columns=columns, dtype="Float64" + ).compute() + pd_result = pd.DataFrame(scalars_pandas_df, columns=columns, dtype="Float64") + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_construct_from_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = dataframe.DataFrame( + {"a": scalars_df["int64_col"], "b": scalars_df["string_col"]}, + dtype="string[pyarrow]", + ).compute() + pd_result = pd.DataFrame( + {"a": scalars_pandas_df["int64_col"], "b": scalars_pandas_df["string_col"]}, + dtype="string[pyarrow]", + ) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_get_column(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + series = scalars_df[col_name] + bf_result = series.compute() + pd_result = scalars_pandas_df[col_name] + assert_series_equal_ignoring_order(bf_result, pd_result) + + +def test_hasattr(scalars_dfs): + scalars_df, _ = scalars_dfs + assert hasattr(scalars_df, "int64_col") + assert hasattr(scalars_df, "head") + assert not hasattr(scalars_df, "not_exist") + + +def test_head_with_custom_column_labels(scalars_df_index, scalars_pandas_df_index): + rename_mapping = { + "int64_col": "Integer Column", + "string_col": "言語列", + } + bf_df = scalars_df_index.rename(columns=rename_mapping).head(3) + bf_result = bf_df.compute() + pd_result = scalars_pandas_df_index.rename(columns=rename_mapping).head(3) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_tail_with_custom_column_labels(scalars_df_index, scalars_pandas_df_index): + rename_mapping = { + "int64_col": "Integer Column", + "string_col": "言語列", + } + bf_df = scalars_df_index.rename(columns=rename_mapping).tail(3) + bf_result = bf_df.compute() + pd_result = scalars_pandas_df_index.rename(columns=rename_mapping).tail(3) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_get_column_by_attr(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + series = scalars_df.int64_col + bf_result = series.compute() + pd_result = scalars_pandas_df.int64_col + assert_series_equal_ignoring_order(bf_result, pd_result) + + +def test_get_columns(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_names = ["bool_col", "float64_col", "int64_col"] + df_subset = scalars_df.get(col_names) + df_pandas = df_subset.compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df[col_names].columns + ) + + +def test_get_columns_default(scalars_dfs): + scalars_df, _ = scalars_dfs + col_names = ["not", "column", "names"] + result = scalars_df.get(col_names, "default_val") + assert result == "default_val" + + +def test_drop_column(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + df_pandas = scalars_df.drop(columns=col_name).compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df.drop(columns=col_name).columns + ) + + +def test_drop_columns(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_names = ["int64_col", "geography_col", "time_col"] + df_pandas = scalars_df.drop(columns=col_names).compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df.drop(columns=col_names).columns + ) + + +def test_drop_with_custom_column_labels(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + rename_mapping = { + "int64_col": "Integer Column", + "string_col": "言語列", + } + dropped_columns = [ + "言語列", + "timestamp_col", + ] + bf_df = scalars_df.rename(columns=rename_mapping).drop(columns=dropped_columns) + bf_result = bf_df.compute() + pd_result = scalars_pandas_df.rename(columns=rename_mapping).drop( + columns=dropped_columns + ) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_rename(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"bool_col": "boolean_col"} + df_pandas = scalars_df.rename(columns=col_name_dict).compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns + ) + + +def test_repr_w_all_rows(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + # Remove columns with flaky formatting, like NUMERIC columns (which use the + # object dtype). Also makes a copy so that mutating the index name doesn't + # break other tests. + scalars_df = scalars_df.drop(columns=["numeric_col"]) + scalars_pandas_df = scalars_pandas_df.drop(columns=["numeric_col"]) + + if scalars_pandas_df.index.name is None: + # Note: Not quite the same as no index / default index, but hopefully + # simulates it well enough while being consistent enough for string + # comparison to work. + scalars_df = scalars_df.set_index("rowindex", drop=False).sort_index() + scalars_df.index.name = None + + # When there are 10 or fewer rows, the outputs should be identical. + actual = repr(scalars_df.head(10)) + + with display_options.pandas_repr(bigframes.options.display): + expected = repr(scalars_pandas_df.head(10)) + + assert actual == expected + + +def test_repr_html_w_all_rows(scalars_dfs): + scalars_df, _ = scalars_dfs + # get a pandas df of the expected format + df, _ = scalars_df._block.compute() + pandas_df = df.set_axis(scalars_df._block.column_labels, axis=1) + pandas_df.index.name = scalars_df.index.name + + # When there are 10 or fewer rows, the outputs should be identical except for the extra note. + actual = scalars_df.head(10)._repr_html_() + with display_options.pandas_repr(bigframes.options.display): + pandas_repr = pandas_df.head(10)._repr_html_() + + expected = ( + pandas_repr + + f"[{len(pandas_df.index)} rows x {len(pandas_df.columns)} columns in total]" + ) + assert actual == expected + + +def test_df_column_name_with_space(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"bool_col": "bool col"} + df_pandas = scalars_df.rename(columns=col_name_dict).compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns + ) + + +def test_df_column_name_duplicate(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"int64_too": "int64_col"} + df_pandas = scalars_df.rename(columns=col_name_dict).compute() + pd.testing.assert_index_equal( + df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns + ) + + +def test_get_df_column_name_duplicate(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"int64_too": "int64_col"} + + bf_result = scalars_df.rename(columns=col_name_dict)["int64_col"].compute() + pd_result = scalars_pandas_df.rename(columns=col_name_dict)["int64_col"] + pd.testing.assert_index_equal(bf_result.columns, pd_result.columns) + + +def test_filter_df(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_bool_series = scalars_df["bool_col"] + bf_result = scalars_df[bf_bool_series].compute() + + pd_bool_series = scalars_pandas_df["bool_col"] + pd_result = scalars_pandas_df[pd_bool_series] + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_new_column(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + kwargs = {"new_col": 2} + df = scalars_df.assign(**kwargs) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign(**kwargs) + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_new_column_w_loc(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df.loc[:, "new_col"] = 2 + pd_df.loc[:, "new_col"] = 2 + bf_result = bf_df.compute() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_new_column_w_setitem(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df["new_col"] = 2 + pd_df["new_col"] = 2 + bf_result = bf_df.compute() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_existing_column(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + kwargs = {"int64_col": 2} + df = scalars_df.assign(**kwargs) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign(**kwargs) + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["int64_col"] = pd_result["int64_col"].astype("Int64") + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + column_name = "int64_col" + df = scalars_df.assign(new_col=scalars_df[column_name]) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign(new_col=scalars_pandas_df[column_name]) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_series_overwrite(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + column_name = "int64_col" + df = scalars_df.assign(**{column_name: scalars_df[column_name] + 3}) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign( + **{column_name: scalars_pandas_df[column_name] + 3} + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_sequential(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + kwargs = {"int64_col": 2, "new_col": 3, "new_col2": 4} + df = scalars_df.assign(**kwargs) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign(**kwargs) + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["int64_col"] = pd_result["int64_col"].astype("Int64") + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + pd_result["new_col2"] = pd_result["new_col2"].astype("Int64") + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +# Require an index so that the self-join is consistent each time. +def test_assign_same_table_different_index_performs_self_join( + scalars_df_index, scalars_pandas_df_index +): + column_name = "int64_col" + bf_df = scalars_df_index.assign( + alternative_index=scalars_df_index["rowindex_2"] + 2 + ) + pd_df = scalars_pandas_df_index.assign( + alternative_index=scalars_pandas_df_index["rowindex_2"] + 2 + ) + bf_df_2 = bf_df.set_index("alternative_index") + pd_df_2 = pd_df.set_index("alternative_index") + bf_result = bf_df.assign(new_col=bf_df_2[column_name] * 10).compute() + pd_result = pd_df.assign(new_col=pd_df_2[column_name] * 10) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +# Different table expression must have Index +def test_assign_different_df( + scalars_df_index, scalars_df_2_index, scalars_pandas_df_index +): + column_name = "int64_col" + df = scalars_df_index.assign(new_col=scalars_df_2_index[column_name]) + bf_result = df.compute() + # Doesn't matter to pandas if it comes from the same DF or a different DF. + pd_result = scalars_pandas_df_index.assign( + new_col=scalars_pandas_df_index[column_name] + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_assign_different_df_w_loc( + scalars_df_index, scalars_df_2_index, scalars_pandas_df_index +): + bf_df = scalars_df_index.copy() + bf_df2 = scalars_df_2_index.copy() + pd_df = scalars_pandas_df_index.copy() + assert "int64_col" in bf_df.columns + assert "int64_col" in pd_df.columns + bf_df.loc[:, "int64_col"] = bf_df2.loc[:, "int64_col"] + 1 + pd_df.loc[:, "int64_col"] = pd_df.loc[:, "int64_col"] + 1 + bf_result = bf_df.compute() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["int64_col"] = pd_result["int64_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_different_df_w_setitem( + scalars_df_index, scalars_df_2_index, scalars_pandas_df_index +): + bf_df = scalars_df_index.copy() + bf_df2 = scalars_df_2_index.copy() + pd_df = scalars_pandas_df_index.copy() + assert "int64_col" in bf_df.columns + assert "int64_col" in pd_df.columns + bf_df["int64_col"] = bf_df2["int64_col"] + 1 + pd_df["int64_col"] = pd_df["int64_col"] + 1 + bf_result = bf_df.compute() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["int64_col"] = pd_result["int64_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_callable_lambda(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + kwargs = {"new_col": lambda x: x["int64_col"] + x["int64_too"]} + df = scalars_df.assign(**kwargs) + bf_result = df.compute() + pd_result = scalars_pandas_df.assign(**kwargs) + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_dropna(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + df = scalars_df.dropna() + bf_result = df.compute() + pd_result = scalars_pandas_df.dropna() + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + on = "rowindex_2" + left_columns = ["int64_col", "float64_col", "rowindex_2"] + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_columns] + # Offset the rows somewhat so that outer join can have an effect. + right = scalars_df[right_columns].assign(rowindex_2=scalars_df["rowindex_2"] + 2) + + df = left.merge(right, merge_how, on, sort=True) + bf_result = df.compute() + + pd_result = scalars_pandas_df[left_columns].merge( + scalars_pandas_df[right_columns].assign( + rowindex_2=scalars_pandas_df["rowindex_2"] + 2 + ), + merge_how, + on, + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge_custom_col_name(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + left_columns = ["int64_col", "float64_col"] + right_columns = ["int64_col", "bool_col", "string_col"] + on = "int64_col" + rename_columns = {"float64_col": "f64_col"} + + left = scalars_df[left_columns] + left = left.rename(columns=rename_columns) + right = scalars_df[right_columns] + df = left.merge(right, merge_how, on, sort=True) + bf_result = df.compute() + + pandas_left_df = scalars_pandas_df[left_columns] + pandas_left_df = pandas_left_df.rename(columns=rename_columns) + pandas_right_df = scalars_pandas_df[right_columns] + pd_result = pandas_left_df.merge(pandas_right_df, merge_how, on, sort=True) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge_left_on_right_on(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + left_columns = ["int64_col", "float64_col", "int64_too"] + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_columns] + right = scalars_df[right_columns] + + df = left.merge( + right, merge_how, left_on="int64_too", right_on="rowindex_2", sort=True + ) + bf_result = df.compute() + + pd_result = scalars_pandas_df[left_columns].merge( + scalars_pandas_df[right_columns], + merge_how, + left_on="int64_too", + right_on="rowindex_2", + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_get_dtypes(scalars_df_default_index): + dtypes = scalars_df_default_index.dtypes + pd.testing.assert_series_equal( + dtypes, + pd.Series( + { + "bool_col": pd.BooleanDtype(), + "bytes_col": np.dtype("O"), + "date_col": pd.ArrowDtype(pa.date32()), + "datetime_col": pd.ArrowDtype(pa.timestamp("us")), + "geography_col": gpd.array.GeometryDtype(), + "int64_col": pd.Int64Dtype(), + "int64_too": pd.Int64Dtype(), + "numeric_col": np.dtype("O"), + "float64_col": pd.Float64Dtype(), + "rowindex": pd.Int64Dtype(), + "rowindex_2": pd.Int64Dtype(), + "string_col": pd.StringDtype(storage="pyarrow"), + "time_col": pd.ArrowDtype(pa.time64("us")), + "timestamp_col": pd.ArrowDtype(pa.timestamp("us", tz="UTC")), + } + ), + ) + + +def test_get_dtypes_array_struct(session): + """We may upgrade struct and array to proper arrow dtype support in future. For now, + we return python objects""" + df = session.read_gbq( + """SELECT + [1, 3, 2] AS array_column, + STRUCT( + "a" AS string_field, + 1.2 AS float_field) AS struct_column""" + ) + + dtypes = df.dtypes + pd.testing.assert_series_equal( + dtypes, + pd.Series({"array_column": np.dtype("O"), "struct_column": np.dtype("O")}), + ) + + +def test_shape(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df.shape + pd_result = scalars_pandas_df.shape + + assert bf_result == pd_result + + +def test_len(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = len(scalars_df) + pd_result = len(scalars_pandas_df) + + assert bf_result == pd_result + + +def test_size(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df.size + pd_result = scalars_pandas_df.size + + assert bf_result == pd_result + + +def test_ndim(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df.ndim + pd_result = scalars_pandas_df.ndim + + assert bf_result == pd_result + + +def test_empty_false(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df.empty + pd_result = scalars_pandas_df.empty + + assert bf_result == pd_result + + +def test_empty_true(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df[[]].empty + pd_result = scalars_pandas_df[[]].empty + + assert bf_result == pd_result + + +@pytest.mark.parametrize( + ("drop",), + ((True,), (False,)), +) +def test_reset_index(scalars_df_index, scalars_pandas_df_index, drop): + df = scalars_df_index.reset_index(drop=drop) + assert df.index.name is None + + bf_result = df.compute() + pd_result = scalars_pandas_df_index.reset_index(drop=drop) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + # reset_index should maintain the original ordering. + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_reset_index_then_filter( + scalars_df_index, + scalars_pandas_df_index, +): + bf_filter = scalars_df_index["bool_col"].fillna(True) + bf_df = scalars_df_index.reset_index()[bf_filter] + bf_result = bf_df.compute() + pd_filter = scalars_pandas_df_index["bool_col"].fillna(True) + pd_result = scalars_pandas_df_index.reset_index()[pd_filter] + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + # reset_index should maintain the original ordering and index keys + # post-filter will have gaps. + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_reset_index_with_unnamed_index( + scalars_df_index, + scalars_pandas_df_index, +): + scalars_df_index = scalars_df_index.copy() + scalars_pandas_df_index = scalars_pandas_df_index.copy() + + scalars_df_index.index.name = None + scalars_pandas_df_index.index.name = None + df = scalars_df_index.reset_index(drop=False) + assert df.index.name is None + + # reset_index(drop=False) creates a new column "index". + assert df.columns[0] == "index" + + bf_result = df.compute() + pd_result = scalars_pandas_df_index.reset_index(drop=False) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + # reset_index should maintain the original ordering. + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_reset_index_with_unnamed_index_and_index_column( + scalars_df_index, + scalars_pandas_df_index, +): + scalars_df_index = scalars_df_index.copy() + scalars_pandas_df_index = scalars_pandas_df_index.copy() + + scalars_df_index.index.name = None + scalars_pandas_df_index.index.name = None + df = scalars_df_index.assign(index=scalars_df_index["int64_col"]).reset_index( + drop=False + ) + assert df.index.name is None + + # reset_index(drop=False) creates a new column "level_0" if the "index" column already exists. + assert df.columns[0] == "level_0" + + bf_result = df.compute() + pd_result = scalars_pandas_df_index.assign( + index=scalars_pandas_df_index["int64_col"] + ).reset_index(drop=False) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + # reset_index should maintain the original ordering. + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("drop",), + ( + (True,), + (False,), + ), +) +@pytest.mark.parametrize( + ("append",), + ( + (True,), + (False,), + ), +) +@pytest.mark.parametrize( + ("index_column",), + (("int64_too",), ("string_col",), ("timestamp_col",)), +) +def test_set_index(scalars_dfs, index_column, drop, append): + scalars_df, scalars_pandas_df = scalars_dfs + df = scalars_df.set_index(index_column, append=append, drop=drop) + bf_result = df.compute() + pd_result = scalars_pandas_df.set_index(index_column, append=append, drop=drop) + + # Sort to disambiguate when there are duplicate index labels. + # Note: Doesn't use assert_pandas_df_equal_ignore_ordering because we get + # "ValueError: 'timestamp_col' is both an index level and a column label, + # which is ambiguous" when trying to sort by a column with the same name as + # the index. + bf_result = bf_result.sort_values("rowindex_2") + pd_result = pd_result.sort_values("rowindex_2") + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_abs(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + columns = ["int64_col", "int64_too", "float64_col"] + + bf_result = scalars_df[columns].abs().compute() + pd_result = scalars_pandas_df[columns].abs() + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_df_isnull(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + columns = ["int64_col", "int64_too", "string_col", "bool_col"] + bf_result = scalars_df[columns].isnull().compute() + pd_result = scalars_pandas_df[columns].isnull() + + # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is + # `BooleanDtype` but the `pd_result.dtype` is `bool`. + pd_result["int64_col"] = pd_result["int64_col"].astype(pd.BooleanDtype()) + pd_result["int64_too"] = pd_result["int64_too"].astype(pd.BooleanDtype()) + pd_result["string_col"] = pd_result["string_col"].astype(pd.BooleanDtype()) + pd_result["bool_col"] = pd_result["bool_col"].astype(pd.BooleanDtype()) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_df_notnull(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + columns = ["int64_col", "int64_too", "string_col", "bool_col"] + bf_result = scalars_df[columns].notnull().compute() + pd_result = scalars_pandas_df[columns].notnull() + + # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is + # `BooleanDtype` but the `pd_result.dtype` is `bool`. + pd_result["int64_col"] = pd_result["int64_col"].astype(pd.BooleanDtype()) + pd_result["int64_too"] = pd_result["int64_too"].astype(pd.BooleanDtype()) + pd_result["string_col"] = pd_result["string_col"].astype(pd.BooleanDtype()) + pd_result["bool_col"] = pd_result["bool_col"].astype(pd.BooleanDtype()) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("op"), + [ + operator.add, + operator.sub, + operator.mul, + operator.truediv, + operator.floordiv, + operator.gt, + operator.ge, + operator.lt, + operator.le, + ], + ids=[ + "add", + "subtract", + "multiply", + "true_divide", + "floor_divide", + "gt", + "ge", + "lt", + "le", + ], +) +# TODO(garrettwu): deal with NA values +@pytest.mark.parametrize(("other_scalar"), [1, 2.5, 0, 0.0]) +@pytest.mark.parametrize(("reverse_operands"), [True, False]) +def test_scalar_binop(scalars_dfs, op, other_scalar, reverse_operands): + scalars_df, scalars_pandas_df = scalars_dfs + columns = ["int64_col", "float64_col"] + + maybe_reversed_op = (lambda x, y: op(y, x)) if reverse_operands else op + + bf_result = maybe_reversed_op(scalars_df[columns], other_scalar).compute() + pd_result = maybe_reversed_op(scalars_pandas_df[columns], other_scalar) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize(("other_scalar"), [1, -2]) +def test_mod(scalars_dfs, other_scalar): + # Zero case excluded as pandas produces 0 result for Int64 inputs rather than NA/NaN. + # This is likely a pandas bug as mod 0 is undefined in other dtypes, and most programming languages. + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = (scalars_df[["int64_col", "int64_too"]] % other_scalar).compute() + pd_result = scalars_pandas_df[["int64_col", "int64_too"]] % other_scalar + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_scalar_binop_str_exception(scalars_dfs): + scalars_df, _ = scalars_dfs + columns = ["string_col"] + with pytest.raises(TypeError): + (scalars_df[columns] + 1).compute() + + +@pytest.mark.parametrize( + ("op"), + [ + (lambda x, y: x.add(y, axis="index")), + (lambda x, y: x.radd(y, axis="index")), + (lambda x, y: x.sub(y, axis="index")), + (lambda x, y: x.rsub(y, axis="index")), + (lambda x, y: x.mul(y, axis="index")), + (lambda x, y: x.rmul(y, axis="index")), + (lambda x, y: x.truediv(y, axis="index")), + (lambda x, y: x.rtruediv(y, axis="index")), + (lambda x, y: x.floordiv(y, axis="index")), + (lambda x, y: x.floordiv(y, axis="index")), + (lambda x, y: x.gt(y, axis="index")), + (lambda x, y: x.ge(y, axis="index")), + (lambda x, y: x.lt(y, axis="index")), + (lambda x, y: x.le(y, axis="index")), + ], + ids=[ + "add", + "radd", + "sub", + "rsub", + "mul", + "rmul", + "truediv", + "rtruediv", + "floordiv", + "rfloordiv", + "gt", + "ge", + "lt", + "le", + ], +) +def test_series_binop_axis_index( + scalars_dfs, + op, +): + scalars_df, scalars_pandas_df = scalars_dfs + df_columns = ["int64_col", "float64_col"] + series_column = "int64_too" + + bf_result = op(scalars_df[df_columns], scalars_df[series_column]).compute() + pd_result = op(scalars_pandas_df[df_columns], scalars_pandas_df[series_column]) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("op"), + [ + (lambda x, y: x.add(y, axis="index")), + (lambda x, y: x.radd(y, axis="index")), + (lambda x, y: x.sub(y, axis="index")), + (lambda x, y: x.rsub(y, axis="index")), + (lambda x, y: x.mul(y, axis="index")), + (lambda x, y: x.rmul(y, axis="index")), + (lambda x, y: x.truediv(y, axis="index")), + (lambda x, y: x.rtruediv(y, axis="index")), + (lambda x, y: x.floordiv(y, axis="index")), + (lambda x, y: x.floordiv(y, axis="index")), + (lambda x, y: x.gt(y, axis="index")), + (lambda x, y: x.ge(y, axis="index")), + (lambda x, y: x.lt(y, axis="index")), + (lambda x, y: x.le(y, axis="index")), + ], + ids=[ + "add", + "radd", + "sub", + "rsub", + "mul", + "rmul", + "truediv", + "rtruediv", + "floordiv", + "rfloordiv", + "gt", + "ge", + "lt", + "le", + ], +) +def test_dataframe_binop_axis_index_throws_not_implemented( + scalars_dfs, + op, +): + scalars_df, scalars_pandas_df = scalars_dfs + df_columns = ["int64_col", "float64_col"] + other_df_columns = ["int64_too"] + + with pytest.raises(NotImplementedError): + op(scalars_df[df_columns], scalars_df[other_df_columns]).compute() + + +# Differnt table will only work for explicit index, since default index orders are arbitrary. +def test_series_binop_add_different_table( + scalars_df_index, scalars_pandas_df_index, scalars_df_2_index +): + df_columns = ["int64_col", "float64_col"] + series_column = "int64_too" + + bf_result = ( + scalars_df_index[df_columns] + .add(scalars_df_2_index[series_column], axis="index") + .compute() + ) + pd_result = scalars_pandas_df_index[df_columns].add( + scalars_pandas_df_index[series_column], axis="index" + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +# TODO(garrettwu): Test series binop with different index + +all_joins = pytest.mark.parametrize( + ("how",), + ( + ("outer",), + ("left",), + ("right",), + ("inner",), + ), +) + + +@all_joins +def test_join_same_table(scalars_dfs, how): + bf_df, pd_df = scalars_dfs + if how == "right" and pd_df.index.name != "rowindex": + pytest.skip("right join not supported without an index") + + bf_df_a = bf_df[["string_col", "int64_col"]] + bf_df_b = bf_df[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, how=how).compute() + pd_df_a = pd_df[["string_col", "int64_col"]] + pd_df_b = pd_df[["float64_col"]] + pd_result = pd_df_a.join(pd_df_b, how=how) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@all_joins +def test_join_different_table( + scalars_df_index, scalars_df_2_index, scalars_pandas_df_index, how +): + bf_df_a = scalars_df_index[["string_col", "int64_col"]] + bf_df_b = scalars_df_2_index.dropna()[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, how=how).compute() + pd_df_a = scalars_pandas_df_index[["string_col", "int64_col"]] + pd_df_b = scalars_pandas_df_index.dropna()[["float64_col"]] + pd_result = pd_df_a.join(pd_df_b, how=how) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_join_duplicate_columns_raises_not_implemented(scalars_dfs): + scalars_df, _ = scalars_dfs + df_a = scalars_df[["string_col", "float64_col"]] + df_b = scalars_df[["float64_col"]] + with pytest.raises(NotImplementedError): + df_a.join(df_b, how="outer").compute() + + +@pytest.mark.parametrize( + ("by", "ascending", "na_position"), + [ + ("int64_col", True, "first"), + (["bool_col", "int64_col"], True, "last"), + ("int64_col", False, "first"), + (["bool_col", "int64_col"], [False, True], "last"), + (["bool_col", "int64_col"], [True, False], "first"), + ], +) +def test_dataframe_sort_values( + scalars_df_index, scalars_pandas_df_index, by, ascending, na_position +): + # Test needs values to be unique + bf_result = scalars_df_index.sort_values( + by, ascending=ascending, na_position=na_position + ).compute() + pd_result = scalars_pandas_df_index.sort_values( + by, ascending=ascending, na_position=na_position + ) + + pandas.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("operator", "columns"), + [ + pytest.param(lambda x: x.cumsum(), ["float64_col", "int64_too"]), + pytest.param(lambda x: x.cumprod(), ["float64_col", "int64_too"]), + pytest.param( + lambda x: x.cumprod(), + ["string_col"], + marks=pytest.mark.xfail( + raises=ValueError, + ), + ), + ], + ids=[ + "cumsum", + "cumprod", + "non-numeric", + ], +) +def test_dataframe_numeric_analytic_op( + scalars_df_index, scalars_pandas_df_index, operator, columns +): + # TODO: Add nullable ints (pandas 1.x has poor behavior on these) + bf_series = operator(scalars_df_index[columns]) + pd_series = operator(scalars_pandas_df_index[columns]) + bf_result = bf_series.compute() + pd.testing.assert_frame_equal(pd_series, bf_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x: x.cummin()), + (lambda x: x.cummax()), + (lambda x: x.shift(2)), + (lambda x: x.shift(-2)), + ], + ids=[ + "cummin", + "cummax", + "shiftpostive", + "shiftnegative", + ], +) +def test_dataframe_general_analytic_op( + scalars_df_index, scalars_pandas_df_index, operator +): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col"] + bf_series = operator(scalars_df_index[col_names]) + pd_series = operator(scalars_pandas_df_index[col_names]) + bf_result = bf_series.compute() + pd.testing.assert_frame_equal( + pd_series, + bf_result, + ) + + +def test_ipython_key_completions_with_drop(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_names = "string_col" + bf_dataframe = scalars_df.drop(columns=col_names) + pd_dataframe = scalars_pandas_df.drop(columns=col_names) + expected = pd_dataframe.columns.tolist() + + results = bf_dataframe._ipython_key_completions_() + + assert col_names not in results + assert results == expected + # _ipython_key_completions_ is called with square brackets + # so only column names are relevant with tab completion + assert "to_gbq" not in results + assert "merge" not in results + assert "drop" not in results + + +def test_ipython_key_completions_with_rename(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"string_col": "a_renamed_column"} + bf_dataframe = scalars_df.rename(columns=col_name_dict) + pd_dataframe = scalars_pandas_df.rename(columns=col_name_dict) + expected = pd_dataframe.columns.tolist() + + results = bf_dataframe._ipython_key_completions_() + + assert "string_col" not in results + assert "a_renamed_column" in results + assert results == expected + # _ipython_key_completions_ is called with square brackets + # so only column names are relevant with tab completion + assert "to_gbq" not in results + assert "merge" not in results + assert "drop" not in results + + +def test__dir__with_drop(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_names = "string_col" + bf_dataframe = scalars_df.drop(columns=col_names) + pd_dataframe = scalars_pandas_df.drop(columns=col_names) + expected = pd_dataframe.columns.tolist() + + results = dir(bf_dataframe) + + assert col_names not in results + assert frozenset(expected) <= frozenset(results) + # __dir__ is called with a '.' and displays all methods, columns names, etc. + assert "to_gbq" in results + assert "merge" in results + assert "drop" in results + + +def test__dir__with_rename(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name_dict = {"string_col": "a_renamed_column"} + bf_dataframe = scalars_df.rename(columns=col_name_dict) + pd_dataframe = scalars_pandas_df.rename(columns=col_name_dict) + expected = pd_dataframe.columns.tolist() + + results = dir(bf_dataframe) + + assert "string_col" not in results + assert "a_renamed_column" in results + assert frozenset(expected) <= frozenset(results) + # __dir__ is called with a '.' and displays all methods, columns names, etc. + assert "to_gbq" in results + assert "merge" in results + assert "drop" in results + + +@pytest.mark.parametrize( + ("start", "stop", "step"), + [ + (0, 0, None), + (None, None, None), + (1, None, None), + (None, 4, None), + (None, None, 2), + (None, 50000000000, 1), + (5, 4, None), + (3, None, 2), + (1, 7, 2), + (1, 7, 50000000000), + ], +) +def test_iloc_slice(scalars_df_index, scalars_pandas_df_index, start, stop, step): + bf_result = scalars_df_index.iloc[start:stop:step].compute() + pd_result = scalars_pandas_df_index.iloc[start:stop:step] + + # Pandas may assign non-object dtype to empty series and series index + # dtypes of empty columns are a known area of divergence from pandas + for column in pd_result.columns: + if ( + pd_result[column].empty and column != "geography_col" + ): # for empty geography_col, bigframes assigns non-object dtype + pd_result[column] = pd_result[column].astype("object") + pd_result.index = pd_result.index.astype("object") + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_iloc_slice_zero_step(scalars_df_index): + with pytest.raises(ValueError): + scalars_df_index.iloc[0:0:0] + + +def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.iloc[1:].iloc[1:].compute() + pd_result = scalars_pandas_df_index.iloc[1:].iloc[1:] + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + "index", + [0, 5], +) +def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): + bf_result = scalars_df_index.iloc[index] + pd_result = scalars_pandas_df_index.iloc[index] + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_iloc_single_integer_out_of_bound_error( + scalars_df_index, scalars_pandas_df_index +): + with pytest.raises(IndexError, match="single positional indexer is out-of-bounds"): + scalars_df_index.iloc[99] + + +def test_loc_bool_series(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.loc[scalars_df_index.bool_col].compute() + pd_result = scalars_pandas_df_index.loc[scalars_pandas_df_index.bool_col] + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_loc_select_column(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.loc[:, "int64_col"].compute() + pd_result = scalars_pandas_df_index.loc[:, "int64_col"] + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("op"), + [ + (lambda x: x.sum(numeric_only=True)), + (lambda x: x.mean(numeric_only=True)), + (lambda x: x.min(numeric_only=True)), + (lambda x: x.max(numeric_only=True)), + (lambda x: x.std(numeric_only=True)), + (lambda x: x.var(numeric_only=True)), + (lambda x: x.count(numeric_only=False)), + (lambda x: x.nunique()), + ], + ids=["sum", "mean", "min", "max", "std", "var", "count", "nunique"], +) +def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op): + col_names = ["int64_too", "float64_col", "string_col", "int64_col", "bool_col"] + bf_series = op(scalars_df_index[col_names]) + pd_series = op(scalars_pandas_df_index[col_names]) + bf_result = bf_series.compute() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_series = pd_series.astype("Float64") + # Pandas has object index type + pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) + + +@pytest.mark.parametrize( + ("op"), + [ + (lambda x: x.all(bool_only=True)), + (lambda x: x.any(bool_only=True)), + ], + ids=["all", "any"], +) +def test_dataframe_bool_aggregates(scalars_df_index, scalars_pandas_df_index, op): + # Pandas will drop nullable 'boolean' dtype so we convert first to bool, then cast back later + scalars_pandas_df_index = scalars_pandas_df_index.assign( + bool_col=scalars_pandas_df_index.bool_col.fillna(False).astype("bool") + ) + bf_series = op(scalars_df_index) + pd_series = op(scalars_pandas_df_index).astype("boolean") + bf_result = bf_series.compute() + + # Pandas has object index type + pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) + + +def test_dataframe_prod(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col"] + bf_series = scalars_df_index[col_names].prod() + pd_series = scalars_pandas_df_index[col_names].prod() + bf_result = bf_series.compute() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_series = pd_series.astype("Float64") + # Pandas has object index type + pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) + + +@pytest.mark.parametrize( + ("frac", "n", "random_state"), + [ + (None, 4, None), + (0.5, None, None), + (None, 4, 10), + (0.5, None, 10), + (None, None, None), + ], + ids=[ + "n_wo_random_state", + "frac_wo_random_state", + "n_w_random_state", + "frac_w_random_state", + "n_default", + ], +) +def test_sample(scalars_dfs, frac, n, random_state): + scalars_df, _ = scalars_dfs + df = scalars_df.sample(frac=frac, n=n, random_state=random_state) + bf_result = df.compute() + + n = 1 if n is None else n + expected_sample_size = round(frac * scalars_df.shape[0]) if frac is not None else n + assert bf_result.shape[0] == expected_sample_size + assert bf_result.shape[1] == scalars_df.shape[1] + + +def test_sample_raises_value_error(scalars_dfs): + scalars_df, _ = scalars_dfs + with pytest.raises( + ValueError, match="Only one of 'n' or 'frac' parameter can be specified." + ): + scalars_df.sample(frac=0.5, n=4) + + +@pytest.mark.parametrize( + ("axis",), + [ + (0,), + (1,), + ], +) +def test_df_add_prefix(scalars_df_index, scalars_pandas_df_index, axis): + if pd.__version__.startswith("1."): + pytest.skip("add_prefix axis parameter not supported in pandas 1.x.") + bf_result = scalars_df_index.add_prefix("prefix_", axis).compute() + + pd_result = scalars_pandas_df_index.add_prefix("prefix_", axis) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + check_index_type=False, + ) + + +@pytest.mark.parametrize( + ("axis",), + [ + (0,), + (1,), + ], +) +def test_df_add_suffix(scalars_df_index, scalars_pandas_df_index, axis): + if pd.__version__.startswith("1."): + pytest.skip("add_prefix axis parameter not supported in pandas 1.x.") + bf_result = scalars_df_index.add_suffix("_suffix", axis).compute() + + pd_result = scalars_pandas_df_index.add_suffix("_suffix", axis) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + check_index_type=False, + ) + + +def test_df_values(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.values + + pd_result = scalars_pandas_df_index.values + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + pd.testing.assert_frame_equal( + pd.DataFrame(bf_result), pd.DataFrame(pd_result), check_dtype=False + ) + + +def test_df_to_numpy(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.to_numpy() + + pd_result = scalars_pandas_df_index.to_numpy() + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + pd.testing.assert_frame_equal( + pd.DataFrame(bf_result), pd.DataFrame(pd_result), check_dtype=False + ) + + +def test_df___array__(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.__array__() + + pd_result = scalars_pandas_df_index.__array__() + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + pd.testing.assert_frame_equal( + pd.DataFrame(bf_result), pd.DataFrame(pd_result), check_dtype=False + ) + + +def test_getattr_not_implemented(scalars_df_index): + with pytest.raises(NotImplementedError): + scalars_df_index.asof() + + +def test_getattr_attribute_error(scalars_df_index): + with pytest.raises(AttributeError): + scalars_df_index.not_a_method() + + +def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index): + index_list = scalars_pandas_df_index.string_col.iloc[[0, 1, 1, 5]].values + + scalars_df_index = scalars_df_index.set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.set_index("string_col") + + bf_result = scalars_df_index.loc[index_list] + pd_result = scalars_pandas_df_index.loc[index_list] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): + index_list = [3, 2, 1, 3, 2, 1] + + bf_result = scalars_df_index.loc[index_list] + pd_result = scalars_pandas_df_index.loc[index_list] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_iloc_list(scalars_df_index, scalars_pandas_df_index): + index_list = [0, 0, 0, 5, 4, 7] + + bf_result = scalars_df_index.iloc[index_list] + pd_result = scalars_pandas_df_index.iloc[index_list] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_iloc_empty_list(scalars_df_index, scalars_pandas_df_index): + index_list = [] + + bf_result = scalars_df_index.iloc[index_list] + pd_result = scalars_pandas_df_index.iloc[index_list] + + bf_result = bf_result.compute() + assert bf_result.shape == pd_result.shape # types are known to be different + + +def test_rename_axis(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.rename_axis("newindexname") + pd_result = scalars_pandas_df_index.rename_axis("newindexname") + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): + pd_string_series = scalars_pandas_df_index.string_col.iloc[[0, 5, 1, 1, 5]] + bf_string_series = scalars_df_index.string_col.iloc[[0, 5, 1, 1, 5]] + + scalars_df_index = scalars_df_index.set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.set_index("string_col") + + bf_result = scalars_df_index.loc[bf_string_series] + pd_result = scalars_pandas_df_index.loc[pd_string_series] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_bf_index_integer_index(scalars_df_index, scalars_pandas_df_index): + pd_index = scalars_pandas_df_index.iloc[[0, 5, 1, 1, 5]].index + bf_index = scalars_df_index.iloc[[0, 5, 1, 1, 5]].index + + bf_result = scalars_df_index.loc[bf_index] + pd_result = scalars_pandas_df_index.loc[pd_index] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_bf_index_integer_index_renamed_col( + scalars_df_index, scalars_pandas_df_index +): + scalars_df_index = scalars_df_index.rename(columns={"int64_col": "rename"}) + scalars_pandas_df_index = scalars_pandas_df_index.rename( + columns={"int64_col": "rename"} + ) + + pd_index = scalars_pandas_df_index.iloc[[0, 5, 1, 1, 5]].index + bf_index = scalars_df_index.iloc[[0, 5, 1, 1, 5]].index + + bf_result = scalars_df_index.loc[bf_index] + pd_result = scalars_pandas_df_index.loc[pd_index] + + pd.testing.assert_frame_equal( + bf_result.compute(), + pd_result, + ) + + +@pytest.mark.parametrize( + ("subset"), + [ + None, + ["bool_col", "int64_too"], + ], +) +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + (False,), + ], +) +def test_df_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, subset): + columns = ["bool_col", "int64_too", "int64_col"] + bf_series = scalars_df_index[columns].drop_duplicates(subset, keep=keep).compute() + pd_series = scalars_pandas_df_index[columns].drop_duplicates(subset, keep=keep) + pd.testing.assert_frame_equal( + pd_series, + bf_series, + ) + + +@pytest.mark.parametrize( + ("subset"), + [ + None, + ["bool_col"], + ], +) +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + (False,), + ], +) +def test_df_duplicated(scalars_df_index, scalars_pandas_df_index, keep, subset): + columns = ["bool_col", "int64_too", "int64_col"] + bf_series = scalars_df_index[columns].duplicated(subset, keep=keep).compute() + pd_series = scalars_pandas_df_index[columns].duplicated(subset, keep=keep) + pd.testing.assert_series_equal(pd_series, bf_series, check_dtype=False) + + +@pytest.mark.parametrize( + ("subset", "normalize", "ascending", "dropna"), + [ + (None, False, False, False), + (None, True, True, True), + ("bool_col", True, False, True), + ], +) +def test_df_value_counts(scalars_dfs, subset, normalize, ascending, dropna): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = ( + scalars_df[["string_col", "bool_col"]] + .value_counts(subset, normalize=normalize, ascending=ascending, dropna=dropna) + .compute() + ) + pd_result = scalars_pandas_df[["string_col", "bool_col"]].value_counts( + subset, normalize=normalize, ascending=ascending, dropna=dropna + ) + + # Older pandas version may not have these values, bigframes tries to emulate 2.0+ + pd_result.name = "count" + pd_result.index.names = bf_result.index.names + + pd.testing.assert_series_equal( + bf_result, pd_result, check_dtype=False, check_index_type=False + ) + + +@pytest.mark.parametrize( + ("na_option", "method", "ascending", "numeric_only"), + [ + ("keep", "average", True, True), + ("top", "min", False, False), + ("bottom", "max", False, False), + ("top", "first", False, False), + ("bottom", "dense", False, False), + ], +) +@pytest.mark.skipif( + True, reason="Blocked by possible pandas rank() regression (b/283278923)" +) +def test_df_rank_with_nulls( + scalars_df_index, + scalars_pandas_df_index, + na_option, + method, + ascending, + numeric_only, +): + unsupported_columns = ["geography_col"] + bf_result = ( + scalars_df_index.drop(columns=unsupported_columns) + .rank( + na_option=na_option, + method=method, + ascending=ascending, + numeric_only=numeric_only, + ) + .compute() + ) + pd_result = ( + scalars_pandas_df_index.drop(columns=unsupported_columns) + .rank( + na_option=na_option, + method=method, + ascending=ascending, + numeric_only=numeric_only, + ) + .astype(pd.Float64Dtype()) + ) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_bool_interpretation_error(scalars_df_index): + with pytest.raises(ValueError): + True if scalars_df_index else False diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py new file mode 100644 index 00000000000..a4318a8658f --- /dev/null +++ b/tests/system/small/test_dataframe_io.py @@ -0,0 +1,325 @@ +# 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. + +from typing import Tuple + +import google.api_core.exceptions +import pandas as pd +import pytest + +from tests.system.utils import ( + assert_pandas_df_equal_ignore_ordering, + convert_pandas_dtypes, +) + +try: + import pandas_gbq # type: ignore +except ImportError: + pandas_gbq = None + +import bigframes +import bigframes.dataframe + + +def test_to_pandas_w_correct_dtypes(scalars_df_default_index): + """Verify to_pandas() APIs returns the expected dtypes.""" + actual = scalars_df_default_index.to_pandas().dtypes + expected = scalars_df_default_index.dtypes + + pd.testing.assert_series_equal(actual, expected) + + +def test_to_pandas_array_struct_correct_result(session): + """In future, we should support arrays and structs with arrow types. + For now we fall back to the current connector behavior of converting + to Python objects""" + df = session.read_gbq( + """SELECT + [1, 3, 2] AS array_column, + STRUCT( + "a" AS string_field, + 1.2 AS float_field) AS struct_column""" + ) + + result = df.to_pandas() + expected = pd.DataFrame( + { + "array_column": [[1, 3, 2]], + "struct_column": [{"string_field": "a", "float_field": 1.2}], + } + ) + expected.index = expected.index.astype("Int64") + pd.testing.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +def test_to_csv_index( + scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame], + gcs_folder: str, + index: bool, +): + """Test the `to_csv` API with the `index` parameter.""" + scalars_df, scalars_pandas_df = scalars_dfs + index_col = None + if scalars_df.index.name is not None: + path = gcs_folder + f"test_index_df_to_csv_index_{index}" + if index: + index_col = scalars_df.index.name + else: + path = gcs_folder + f"test_default_index_df_to_csv_index_{index}" + + # TODO(swast): Support "date_format" parameter and make sure our + # DATETIME/TIMESTAMP column export is the same format as pandas by default. + scalars_df.to_csv(path, index=index) + + # Pandas dataframes dtypes from read_csv are not fully compatible with + # BigQuery-backed dataframes, so manually convert the dtypes specifically + # here. + dtype = scalars_df.reset_index().dtypes.to_dict() + dtype.pop("timestamp_col") + dtype.pop("geography_col") + gcs_df = pd.read_csv( + path, dtype=dtype, parse_dates=["timestamp_col"], index_col=index_col + ) + convert_pandas_dtypes(gcs_df, bytes_col=True) + + assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +@pytest.mark.skipif(pandas_gbq is None, reason="required by pd.read_gbq") +def test_to_gbq_index(scalars_dfs, dataset_id, index): + """Test the `to_gbq` API with the `index` parameter.""" + scalars_df, scalars_pandas_df = scalars_dfs + destination_table = f"{dataset_id}.test_index_df_to_gbq_{index}" + df_in = scalars_df.copy() + if index: + index_col = "index" + df_in.index.name = index_col + else: + index_col = None + + df_in.to_gbq(destination_table, if_exists="replace", index=index) + df_out = pd.read_gbq(destination_table, index_col=index_col) + + if index: + df_out = df_out.sort_index() + else: + df_out = df_out.sort_values("rowindex_2").reset_index(drop=True) + + convert_pandas_dtypes(df_out, bytes_col=False) + expected = scalars_pandas_df.copy() + expected.index.name = index_col + pd.testing.assert_frame_equal(df_out, expected, check_index_type=False) + + +@pytest.mark.parametrize( + ("if_exists", "expected_index"), + [ + pytest.param("replace", 1), + pytest.param("append", 2), + pytest.param( + "fail", + 0, + marks=pytest.mark.xfail( + raises=google.api_core.exceptions.Conflict, + ), + ), + pytest.param( + "unknown", + 0, + marks=pytest.mark.xfail( + raises=ValueError, + ), + ), + ], +) +@pytest.mark.skipif(pandas_gbq is None, reason="required by pd.read_gbq") +def test_to_gbq_if_exists( + scalars_df_default_index, + scalars_pandas_df_default_index, + dataset_id, + if_exists, + expected_index, +): + """Test the `to_gbq` API with the `if_exists` parameter.""" + destination_table = f"{dataset_id}.test_to_gbq_if_exists_{if_exists}" + + scalars_df_default_index.to_gbq(destination_table) + scalars_df_default_index.to_gbq(destination_table, if_exists=if_exists) + + gcs_df = pd.read_gbq(destination_table) + assert len(gcs_df.index) == expected_index * len( + scalars_pandas_df_default_index.index + ) + pd.testing.assert_index_equal( + gcs_df.columns, scalars_pandas_df_default_index.columns + ) + + +def test_to_gbq_w_invalid_destination_table(scalars_df_index): + with pytest.raises(ValueError): + scalars_df_index.to_gbq("table_id") + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +def test_to_json_index_invalid_orient( + scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame], + gcs_folder: str, + index: bool, +): + scalars_df, scalars_pandas_df = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + f"test_index_df_to_json_index_{index}" + else: + path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + with pytest.raises(ValueError): + scalars_df.to_json(path, index=index, lines=True) + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +def test_to_json_index_invalid_lines( + scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame], + gcs_folder: str, + index: bool, +): + scalars_df, scalars_pandas_df = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + f"test_index_df_to_json_index_{index}" + else: + path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + with pytest.raises(NotImplementedError): + scalars_df.to_json(path, index=index) + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +def test_to_json_index_records_orient( + scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame], + gcs_folder: str, + index: bool, +): + """Test the `to_json` API with the `index` parameter.""" + scalars_df, scalars_pandas_df = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + f"test_index_df_to_json_index_{index}" + else: + path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + + """ Test the `to_json` API with `orient` is `records` and `lines` is True""" + scalars_df.to_json(path, index=index, orient="records", lines=True) + + gcs_df = pd.read_json(path, lines=True, convert_dates=["datetime_col"]) + convert_pandas_dtypes(gcs_df, bytes_col=True) + if index and scalars_df.index.name is not None: + gcs_df = gcs_df.set_index(scalars_df.index.name) + + assert len(gcs_df.index) == len(scalars_pandas_df.index) + pd.testing.assert_index_equal(gcs_df.columns, scalars_pandas_df.columns) + assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + + +@pytest.mark.parametrize( + ("index"), + [True, False], +) +def test_to_parquet_index(scalars_dfs, gcs_folder, index): + """Test the `to_parquet` API with the `index` parameter.""" + scalars_df, scalars_pandas_df = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + f"test_index_df_to_parquet_{index}" + else: + path = gcs_folder + f"test_default_index_df_to_parquet_{index}" + + # TODO(b/268693993): Type GEOGRAPHY is not currently supported for parquet. + scalars_df = scalars_df.drop(columns="geography_col") + scalars_pandas_df = scalars_pandas_df.drop(columns="geography_col") + + # TODO(swast): Do a bit more processing on the input DataFrame to ensure + # the exported results are from the generated query, not just the source + # table. + scalars_df.to_parquet(path, index=index) + + gcs_df = pd.read_parquet(path) + convert_pandas_dtypes(gcs_df, bytes_col=False) + if index and scalars_df.index.name is not None: + gcs_df = gcs_df.set_index(scalars_df.index.name) + + assert len(gcs_df.index) == len(scalars_pandas_df.index) + pd.testing.assert_index_equal(gcs_df.columns, scalars_pandas_df.columns) + assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + + +def test_to_sql_query_named_index_included( + session, scalars_df_index, scalars_pandas_df_index +): + sql, index_columns = scalars_df_index.to_sql_query(always_include_index=True) + assert len(index_columns) == 1 + index_column, is_named = index_columns[0] + assert index_column == "rowindex" + assert is_named + + roundtrip = session.read_gbq(sql, index_col=[index_column]) + assert_pandas_df_equal_ignore_ordering( + roundtrip.to_pandas(), scalars_pandas_df_index + ) + + +def test_to_sql_query_unnamed_index_excluded( + session, scalars_df_default_index, scalars_pandas_df_default_index +): + # The .sql property should return SQL without the unnamed indexes + sql, index_columns = scalars_df_default_index.to_sql_query( + always_include_index=False + ) + assert len(index_columns) == 0 + + roundtrip = session.read_gbq(sql) + assert_pandas_df_equal_ignore_ordering( + roundtrip.to_pandas(), scalars_pandas_df_default_index + ) + + +def test_to_sql_query_unnamed_index_always_include( + session, + scalars_df_default_index: bigframes.dataframe.DataFrame, + scalars_pandas_df_default_index, +): + sql, index_columns = scalars_df_default_index.to_sql_query( + always_include_index=True + ) + assert len(index_columns) == 1 + index_column, is_named = index_columns[0] + assert index_column == "bigframes_index_0" + assert not is_named + + roundtrip = session.read_gbq(sql, index_col=[index_column]) + roundtrip.index.name = None + assert_pandas_df_equal_ignore_ordering( + roundtrip.to_pandas(), scalars_pandas_df_default_index + ) diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py new file mode 100644 index 00000000000..066e20bb12c --- /dev/null +++ b/tests/system/small/test_groupby.py @@ -0,0 +1,129 @@ +# 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 pandas as pd +import pytest + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x: x.sum(numeric_only=True)), + (lambda x: x.mean(numeric_only=True)), + (lambda x: x.min(numeric_only=True)), + (lambda x: x.max(numeric_only=True)), + (lambda x: x.std(numeric_only=True)), + (lambda x: x.var(numeric_only=True)), + ], + ids=[ + "sum", + "mean", + "min", + "max", + "std", + "var", + ], +) +def test_dataframe_groupby_numeric_aggregate( + scalars_df_index, scalars_pandas_df_index, operator +): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = operator(scalars_df_index[col_names].groupby("string_col")) + pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) + bf_result_computed = bf_result.compute() + # Pandas std function produces float64, not matching Float64 from bigframes + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x: x.count()), + (lambda x: x.any()), + (lambda x: x.all()), + ], + ids=[ + "count", + "any", + "all", + ], +) +def test_dataframe_groupby_aggregate( + scalars_df_index, scalars_pandas_df_index, operator +): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = operator(scalars_df_index[col_names].groupby("string_col")) + pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) + bf_result_computed = bf_result.compute() + + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + +@pytest.mark.parametrize( + ("as_index"), + [ + (True), + (False), + ], +) +def test_dataframe_groupby_multi_sum( + scalars_df_index, scalars_pandas_df_index, as_index +): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_series = ( + scalars_df_index[col_names] + .groupby(["bool_col", "int64_col"], as_index=as_index) + .sum(numeric_only=True) + ) + pd_series = ( + scalars_pandas_df_index[col_names] + .groupby(["bool_col", "int64_col"], as_index=as_index) + .sum(numeric_only=True) + ) + bf_result = bf_series.compute() + + if not as_index: + # BigQuery DataFrames default indices use nullable Int64 always + pd_series.index = pd_series.index.astype("Int64") + + pd.testing.assert_frame_equal( + pd_series, + bf_result, + ) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x: x.cumsum(numeric_only=True)), + (lambda x: x.cummax(numeric_only=True)), + (lambda x: x.cummin(numeric_only=True)), + (lambda x: x.cumprod()), + ], + ids=[ + "cumsum", + "cummax", + "cummin", + "cumprod", + ], +) +def test_dataframe_groupby_analytic( + scalars_df_index, scalars_pandas_df_index, operator +): + col_names = ["float64_col", "int64_col", "bool_col", "string_col"] + bf_result = operator(scalars_df_index[col_names].groupby("string_col")) + pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) + bf_result_computed = bf_result.compute() + + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py new file mode 100644 index 00000000000..312e86d7f23 --- /dev/null +++ b/tests/system/small/test_index.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. + +import numpy + +from tests.system.utils import assert_pandas_index_equal_ignore_index_type + + +def test_get_index(scalars_df_index, scalars_pandas_df_index): + index = scalars_df_index.index + bf_result = index.compute() + pd_result = scalars_pandas_df_index.index + + assert_pandas_index_equal_ignore_index_type(bf_result, pd_result) + + +def test_index_shape(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.index.shape + pd_result = scalars_pandas_df_index.index.shape + + assert bf_result == pd_result + + +def test_index_len(scalars_df_index, scalars_pandas_df_index): + bf_result = len(scalars_df_index.index) + pd_result = len(scalars_pandas_df_index.index) + + assert bf_result == pd_result + + +def test_index_array(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.__array__() + pd_result = scalars_pandas_df_index.__array__() + + numpy.array_equal(bf_result, pd_result) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py new file mode 100644 index 00000000000..bf2a2080de3 --- /dev/null +++ b/tests/system/small/test_multiindex.py @@ -0,0 +1,273 @@ +# 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 pandas +import pytest + +import bigframes.pandas as bpd + + +def test_set_multi_index(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index(["bool_col", "int64_too"]).compute() + pd_result = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_reset_multi_index(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.set_index(["bool_col", "int64_too"]).reset_index().compute() + ) + pd_result = scalars_pandas_df_index.set_index( + ["bool_col", "int64_too"] + ).reset_index() + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pandas.Int64Dtype()) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_binop_series_series_matching_multi_indices( + scalars_df_index, scalars_pandas_df_index +): + bf_left = scalars_df_index.set_index(["bool_col", "string_col"]) + bf_right = scalars_df_index.set_index(["bool_col", "string_col"]) + pd_left = scalars_pandas_df_index.set_index(["bool_col", "string_col"]) + pd_right = scalars_pandas_df_index.set_index(["bool_col", "string_col"]) + + bf_result = bf_left["int64_col"] + bf_right["int64_too"] + pd_result = pd_left["int64_col"] + pd_right["int64_too"] + + pandas.testing.assert_series_equal( + bf_result.sort_index().compute(), pd_result.sort_index() + ) + + +def test_binop_df_series_matching_multi_indices( + scalars_df_index, scalars_pandas_df_index +): + bf_left = scalars_df_index.set_index(["bool_col", "string_col"]) + bf_right = scalars_df_index.set_index(["bool_col", "string_col"]) + pd_left = scalars_pandas_df_index.set_index(["bool_col", "string_col"]) + pd_right = scalars_pandas_df_index.set_index(["bool_col", "string_col"]) + + bf_result = bf_left[["int64_col", "int64_too"]].add(bf_right["int64_too"], axis=0) + pd_result = pd_left[["int64_col", "int64_too"]].add(pd_right["int64_too"], axis=0) + + pandas.testing.assert_frame_equal( + bf_result.sort_index().compute(), pd_result.sort_index() + ) + + +def test_binop_multi_index_mono_index(scalars_df_index, scalars_pandas_df_index): + bf_left = scalars_df_index.set_index(["bool_col", "rowindex_2"]) + bf_right = scalars_df_index.set_index("rowindex_2") + pd_left = scalars_pandas_df_index.set_index(["bool_col", "rowindex_2"]) + pd_right = scalars_pandas_df_index.set_index("rowindex_2") + + bf_result = bf_left["int64_col"] + bf_right["int64_too"] + pd_result = pd_left["int64_col"] + pd_right["int64_too"] + + pandas.testing.assert_series_equal(bf_result.compute(), pd_result) + + +def test_binop_overlapping_multi_indices(scalars_df_index, scalars_pandas_df_index): + bf_left = scalars_df_index.set_index(["bool_col", "int64_too"]) + bf_right = scalars_df_index.set_index(["bool_col", "int64_col"]) + pd_left = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) + pd_right = scalars_pandas_df_index.set_index(["bool_col", "int64_col"]) + + bf_result = bf_left["int64_col"] + bf_right["int64_too"] + pd_result = pd_left["int64_col"] + pd_right["int64_too"] + + pandas.testing.assert_series_equal( + bf_result.sort_index().compute(), pd_result.sort_index() + ) + + +def test_concat_compatible_multi_indices(scalars_df_index, scalars_pandas_df_index): + if pandas.__version__.startswith("1."): + pytest.skip("Labels not preserved in pandas 1.x.") + bf_left = scalars_df_index.set_index(["bool_col", "int64_col"]) + bf_right = scalars_df_index.set_index(["bool_col", "int64_too"]) + pd_left = scalars_pandas_df_index.set_index(["bool_col", "int64_col"]) + pd_right = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) + + bf_result = bpd.concat([bf_left, bf_right]) + pd_result = pandas.concat([pd_left, pd_right]) + + pandas.testing.assert_frame_equal(bf_result.compute(), pd_result) + + +def test_concat_multi_indices_ignore_index(scalars_df_index, scalars_pandas_df_index): + bf_left = scalars_df_index.set_index(["bool_col", "int64_too"]) + bf_right = scalars_df_index.set_index(["bool_col", "int64_col"]) + pd_left = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) + pd_right = scalars_pandas_df_index.set_index(["bool_col", "int64_col"]) + + bf_result = bpd.concat([bf_left, bf_right], ignore_index=True) + pd_result = pandas.concat([pd_left, pd_right], ignore_index=True) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pandas.Int64Dtype()) + + pandas.testing.assert_frame_equal(bf_result.compute(), pd_result) + + +def test_multi_index_loc(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"]).loc[[2, 0]].compute() + ) + pd_result = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]).loc[[2, 0]] + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_multi_index_getitem_bool(scalars_df_index, scalars_pandas_df_index): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) + + bf_result = bf_frame[bf_frame["int64_col"] > 0].compute() + pd_result = pd_frame[pd_frame["int64_col"] > 0] + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("level"), + [ + (1), + ("int64_too"), + ([0, 2]), + ([2, "bool_col"]), + ], + ids=["level_num", "level_name", "list", "mixed_list"], +) +def test_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame.droplevel(level).compute() + pd_result = pd_frame.droplevel(level) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("order"), + [ + (1, 0, 2), + (["int64_col", "bool_col", "int64_too"]), + (["int64_col", "bool_col", 0]), + ], + ids=[ + "level_nums", + "level_names", + "num_names_mixed", + ], +) +def test_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_index, order): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame.reorder_levels(order).compute() + pd_result = pd_frame.reorder_levels(order) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_multi_index_series_groupby(scalars_df_index, scalars_pandas_df_index): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col"]) + bf_result = ( + bf_frame["float64_col"] + .groupby([bf_frame.int64_col % 2, "bool_col"]) + .mean() + .compute() + ) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) + pd_result = ( + pd_frame["float64_col"].groupby([pd_frame.int64_col % 2, "bool_col"]).mean() + ) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("level"), + [ + (1), + ([0]), + (["bool_col"]), + (["bool_col", "int64_too"]), + ], +) +def test_multi_index_series_groupby_level( + scalars_df_index, scalars_pandas_df_index, level +): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level) + .mean() + .compute() + ) + pd_result = ( + scalars_pandas_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level) + .mean() + ) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + +def test_multi_index_dataframe_groupby(scalars_df_index, scalars_pandas_df_index): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col"]) + bf_result = ( + bf_frame.groupby([bf_frame.int64_col % 2, "bool_col"]) + .mean(numeric_only=True) + .compute() + ) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) + pd_result = pd_frame.groupby([pd_frame.int64_col % 2, "bool_col"]).mean( + numeric_only=True + ) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("level"), + [ + (1), + ([0]), + (["bool_col"]), + (["bool_col", "int64_too"]), + ], +) +def test_multi_index_dataframe_groupby_level( + scalars_df_index, scalars_pandas_df_index, level +): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"]) + .groupby(level=level) + .mean(numeric_only=True) + .compute() + ) + pd_result = ( + scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) + .groupby(level=level) + .mean(numeric_only=True) + ) + + pandas.testing.assert_frame_equal(bf_result, pd_result) diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py new file mode 100644 index 00000000000..a612ad946fa --- /dev/null +++ b/tests/system/small/test_pandas.py @@ -0,0 +1,64 @@ +# 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 pandas as pd +import pytest + +import bigframes.pandas as bpd + + +def test_concat_dataframe(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = bpd.concat(11 * [scalars_df]) + bf_result = bf_result.compute() + pd_result = pd.concat(11 * [scalars_pandas_df]) + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_concat_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = bpd.concat( + [scalars_df.int64_col, scalars_df.int64_too, scalars_df.int64_col] + ) + bf_result = bf_result.compute() + pd_result = pd.concat( + [ + scalars_pandas_df.int64_col, + scalars_pandas_df.int64_too, + scalars_pandas_df.int64_col, + ] + ) + + pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("how",), + [ + ("inner",), + ("outer",), + ], +) +def test_concat_dataframe_mismatched_columns(scalars_dfs, how): + cols1 = ["int64_too", "int64_col", "float64_col"] + cols2 = ["int64_col", "string_col", "int64_too"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = bpd.concat([scalars_df[cols1], scalars_df[cols2]], join=how) + bf_result = bf_result.compute() + pd_result = pd.concat( + [scalars_pandas_df[cols1], scalars_pandas_df[cols2]], join=how + ) + + pd.testing.assert_frame_equal(bf_result, pd_result) diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py new file mode 100644 index 00000000000..261e0d5b14e --- /dev/null +++ b/tests/system/small/test_pandas_options.py @@ -0,0 +1,145 @@ +# 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 pytest + +import bigframes.pandas as bpd + + +@pytest.fixture(autouse=True) +def reset_default_session_and_location(): + bpd.reset_session() + bpd.options.bigquery.location = None + + +@pytest.mark.parametrize( + ("read_method", "query_prefix"), + [ + (bpd.read_gbq, None), + (bpd.read_gbq, "SELECT COUNT(1) FROM "), + (bpd.read_gbq_table, None), + (bpd.read_gbq_query, "SELECT COUNT(1) FROM "), + ], + ids=[ + "read_gbq-on-table-name", + "read_gbq-on-sql", + "read_gbq_table-on-table-name", + "read_gbq_query-on-sql", + ], +) +def test_read_gbq_start_sets_session_location( + test_data_tables_tokyo, + dataset_id_permanent_tokyo, + tokyo_location, + test_data_tables, + dataset_id_permanent, + read_method, + query_prefix, +): + # Form query as a table name or a SQL depending on the test scenario + query_tokyo = test_data_tables_tokyo["scalars"] + query = test_data_tables["scalars"] + if query_prefix: + query_tokyo = f"{query_prefix} {query_tokyo}" + query = f"{query_prefix} {query}" + + # Initially there is no location set in the bigquery options + assert not bpd.options.bigquery.location + + # Starting user journey with read_gbq* should work for a table in any + # location, in this case tokyo + df = read_method(query_tokyo) + assert df is not None + + # Now bigquery options location should be set to tokyo + assert bpd.options.bigquery.location == tokyo_location + + # Now read_gbq* from another location should fail + with pytest.raises( + google.api_core.exceptions.NotFound, + match=f"404 Not found: Dataset {dataset_id_permanent} was not found in location {tokyo_location}", + ): + read_method(query) + + # Reset global session to start over + bpd.reset_session() + + # There should still be the previous location set in the bigquery options + assert bpd.options.bigquery.location == tokyo_location + + # Starting over the user journey with read_gbq* should work for a table + # in another location, in this case US + df = read_method(query) + assert df is not None + + # Now bigquery options location should be set to US + assert bpd.options.bigquery.location == "US" + + # Now read_gbq* from another location should fail + with pytest.raises( + google.api_core.exceptions.NotFound, + match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + ): + read_method(query_tokyo) + + +@pytest.mark.parametrize( + ("read_method", "query_prefix"), + [ + (bpd.read_gbq, None), + (bpd.read_gbq, "SELECT COUNT(1) FROM "), + (bpd.read_gbq_table, None), + (bpd.read_gbq_query, "SELECT COUNT(1) FROM "), + ], + ids=[ + "read_gbq-on-table-name", + "read_gbq-on-sql", + "read_gbq_table-on-table-name", + "read_gbq_query-on-sql", + ], +) +def test_read_gbq_after_session_start_must_comply_with_default_location( + scalars_pandas_df_index, + test_data_tables, + test_data_tables_tokyo, + dataset_id_permanent_tokyo, + read_method, + query_prefix, +): + # Form query as a table name or a SQL depending on the test scenario + query_tokyo = test_data_tables_tokyo["scalars"] + query = test_data_tables["scalars"] + if query_prefix: + query_tokyo = f"{query_prefix} {query_tokyo}" + query = f"{query_prefix} {query}" + + # Initially there is no location set in the bigquery options + assert not bpd.options.bigquery.location + + # Starting user journey with anything other than read_gbq*, such as + # read_pandas would bind the session to default location US + df = bpd.read_pandas(scalars_pandas_df_index) + assert df is not None + + # Doing read_gbq* from a table in another location should fail + with pytest.raises( + google.api_core.exceptions.NotFound, + match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + ): + read_method(query_tokyo) + + # read_gbq* from a table in the default location should work + df = read_method(query) + assert df is not None diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py new file mode 100644 index 00000000000..6eeee3a3a43 --- /dev/null +++ b/tests/system/small/test_progress_bar.py @@ -0,0 +1,74 @@ +# 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 bigframes as bf +import bigframes.formatting_helpers as formatting_helpers + + +def test_progress_bar_dataframe( + penguins_df_default_index: bf.dataframe.DataFrame, capsys +): + bf.options.display.progress_bar = "notebook" + penguins_df_default_index.compute() + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = filter(None, lines) + assert penguins_df_default_index.query_job is not None + for line in lines: + assert html_check in line and open_job_check in line + + +def test_progress_bar_series(penguins_df_default_index: bf.dataframe.DataFrame, capsys): + bf.options.display.progress_bar = "notebook" + series = penguins_df_default_index["body_mass_g"].head(10) + series.compute() + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = filter(None, lines) + assert series.query_job is not None + for line in lines: + assert html_check in line and open_job_check in line + + +def test_progress_bar_scalar(penguins_df_default_index: bf.dataframe.DataFrame, capsys): + bf.options.display.progress_bar = "notebook" + penguins_df_default_index["body_mass_g"].head(10).mean() + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = filter(None, lines) + for line in lines: + assert html_check in line and open_job_check in line + + +def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame): + bf.options.display.progress_bar = "notebook" + penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = ( + False + ) + penguins_df_default_index.compute() + query_job_repr = formatting_helpers.repr_query_job( + penguins_df_default_index.query_job + ).value + string_checks = [ + "Job Id", + "Destination Table", + "Slot Time", + "Bytes Processed", + "Cache hit", + ] + for string in string_checks: + assert string in query_job_repr diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py new file mode 100644 index 00000000000..4c75b5d3d27 --- /dev/null +++ b/tests/system/small/test_remote_function.py @@ -0,0 +1,278 @@ +# 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 pandas as pd +import pytest + +import bigframes +from bigframes.remote_function import remote_function +from tests.system.utils import assert_pandas_df_equal_ignore_ordering + + +@pytest.fixture(scope="module") +def bq_cf_connection() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us --project_id=bigframes-dev bigframes-rf-conn + """ + return "bigframes-rf-conn" + + +@pytest.fixture(scope="module") +def session_with_bq_connection(bq_cf_connection) -> bigframes.Session: + return bigframes.Session( + bigframes.BigQueryOptions(remote_udf_connection=bq_cf_connection) + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_no_session_param( + bigquery_client, + bigqueryconnection_client, + scalars_dfs, + dataset_id_permanent, + bq_cf_connection, +): + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_session_param(session_with_bq_connection, scalars_dfs): + @remote_function( + [int], + int, + session=session_with_bq_connection, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_via_session_default(session_with_bq_connection, scalars_dfs): + # Session has bigquery connection initialized via context. Without an + # explicit dataset the default dataset from the session would be used. + # Without an explicit bigquery connection, the one present in Session set + # through the explicit BigQueryOptions would be used. Without an explicit `reuse` + # the default behavior of reuse=True will take effect. Please note that the + # udf is same as the one used in other tests in this file so the underlying + # cloud function would be common and quickly reused. + @session_with_bq_connection.remote_function([int], int) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_via_session_with_overrides( + session, scalars_dfs, dataset_id_permanent, bq_cf_connection +): + @session.remote_function( + [int], + int, + dataset_id_permanent, + bq_cf_connection, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_via_session_context_connection_setter( + scalars_dfs, dataset_id, bq_cf_connection +): + # Creating a session scoped only to this test as we would be setting a + # property in it + context = bigframes.BigQueryOptions() + context.remote_udf_connection = bq_cf_connection + session = bigframes.connect(context) + + # Without an explicit bigquery connection, the one present in Session, + # set via context setter would be used. Without an explicit `reuse` the + # default behavior of reuse=True will take effect. Please note that the + # udf is same as the one used in other tests in this file so the underlying + # cloud function would be common with reuse=True. Since we are using a + # unique dataset_id, even though the cloud function would be reused, the bq + # remote function would still be created, making use of the bq connection + # set in the BigQueryOptions above. + @session.remote_function([int], int, dataset=dataset_id) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_dataframe_applymap(session_with_bq_connection, scalars_dfs): + def add_one(x): + return x + 1 + + remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + + scalars_df, scalars_pandas_df = scalars_dfs + int64_cols = ["int64_col", "int64_too"] + + bf_int64_df = scalars_df[int64_cols] + bf_int64_df_filtered = bf_int64_df.dropna() + bf_result = bf_int64_df_filtered.applymap(remote_add_one).compute() + + pd_int64_df = scalars_pandas_df[int64_cols] + pd_int64_df_filtered = pd_int64_df.dropna() + pd_result = pd_int64_df_filtered.applymap(add_one) + # TODO(shobs): Figure why pandas .applymap() changes the dtype, i.e. + # pd_int64_df_filtered.dtype is Int64Dtype() + # pd_int64_df_filtered.applymap(lambda x: x).dtype is int64. + # For this test let's force the pandas dtype to be same as input. + for col in pd_result: + pd_result[col] = pd_result[col].astype(pd_int64_df_filtered[col].dtype) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_dataframe_applymap_na_ignore(session_with_bq_connection, scalars_dfs): + def add_one(x): + return x + 1 + + remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + + scalars_df, scalars_pandas_df = scalars_dfs + int64_cols = ["int64_col", "int64_too"] + + bf_int64_df = scalars_df[int64_cols] + bf_result = bf_int64_df.applymap(remote_add_one, na_action="ignore").compute() + + pd_int64_df = scalars_pandas_df[int64_cols] + pd_result = pd_int64_df.applymap(add_one, na_action="ignore") + # TODO(shobs): Figure why pandas .applymap() changes the dtype, i.e. + # pd_int64_df_filtered.dtype is Int64Dtype() + # pd_int64_df_filtered.applymap(lambda x: x).dtype is int64. + # For this test let's force the pandas dtype to be same as input. + for col in pd_result: + pd_result[col] = pd_result[col].astype(pd_int64_df[col].dtype) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) diff --git a/tests/system/small/test_scalar.py b/tests/system/small/test_scalar.py new file mode 100644 index 00000000000..8b9e1fd95a9 --- /dev/null +++ b/tests/system/small/test_scalar.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. + +import pandas + + +def test_repr(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_series = scalars_df[col_name] + pd_series = scalars_pandas_df[col_name].astype(pandas.Int64Dtype()) + bf_scalar = bf_series.sum() + pd_scalar = pd_series.sum() + assert repr(bf_scalar) == repr(pd_scalar) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py new file mode 100644 index 00000000000..751521bd754 --- /dev/null +++ b/tests/system/small/test_series.py @@ -0,0 +1,1934 @@ +# 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 math +import tempfile + +import geopandas as gpd # type: ignore +import numpy +import pandas as pd +import pyarrow as pa # type: ignore +import pytest + +import bigframes.series as series +from tests.system.utils import ( + assert_pandas_df_equal_ignore_ordering, + assert_series_equal_ignoring_order, +) + + +def test_series_construct_copy(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = series.Series( + scalars_df["int64_col"], name="test_series", dtype="Float64" + ).compute() + pd_result = pd.Series( + scalars_pandas_df["int64_col"], name="test_series", dtype="Float64" + ) + pd.testing.assert_series_equal(bf_result, pd_result) + + +def test_series_construct_pandas(scalars_dfs): + _, scalars_pandas_df = scalars_dfs + bf_result = series.Series( + scalars_pandas_df["int64_col"], name="test_series", dtype="Float64" + ) + pd_result = pd.Series( + scalars_pandas_df["int64_col"], name="test_series", dtype="Float64" + ) + assert bf_result.shape == pd_result.shape + pd.testing.assert_series_equal(bf_result.compute(), pd_result) + + +def test_series_construct_from_list(): + bf_result = series.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64").compute() + pd_result = pd.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64") + + # BigQuery DataFrame default indices use nullable Int64 always + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ["col_name", "expected_dtype"], + [ + ("bool_col", pd.BooleanDtype()), + # TODO(swast): Use a more efficient type. + ("bytes_col", numpy.dtype("object")), + ("date_col", pd.ArrowDtype(pa.date32())), + ("datetime_col", pd.ArrowDtype(pa.timestamp("us"))), + ("float64_col", pd.Float64Dtype()), + ("geography_col", gpd.array.GeometryDtype()), + ("int64_col", pd.Int64Dtype()), + # TODO(swast): Use a more efficient type. + ("numeric_col", numpy.dtype("object")), + ("int64_too", pd.Int64Dtype()), + ("string_col", pd.StringDtype(storage="pyarrow")), + ("time_col", pd.ArrowDtype(pa.time64("us"))), + ("timestamp_col", pd.ArrowDtype(pa.timestamp("us", tz="UTC"))), + ], +) +def test_get_column(scalars_dfs, col_name, expected_dtype): + scalars_df, scalars_pandas_df = scalars_dfs + series = scalars_df[col_name] + series_pandas = series.compute() + assert series_pandas.dtype == expected_dtype + assert series_pandas.shape[0] == scalars_pandas_df.shape[0] + + +def test_series_get_column_default(scalars_dfs): + scalars_df, _ = scalars_dfs + result = scalars_df.get(123123123123123, "default_val") + assert result == "default_val" + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_too",), + ), +) +def test_abs(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].abs().compute() + pd_result = scalars_pandas_df[col_name].abs() + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_fillna(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = scalars_df[col_name].fillna("Missing").compute() + pd_result = scalars_pandas_df[col_name].fillna("Missing") + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("string_col",), + ("int64_col",), + ), +) +def test_max(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].max() + pd_result = scalars_pandas_df[col_name].max() + assert pd_result == bf_result + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("string_col",), + ("int64_col",), + ), +) +def test_min(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].min() + pd_result = scalars_pandas_df[col_name].min() + assert pd_result == bf_result + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_col",), + ), +) +def test_std(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].std() + pd_result = scalars_pandas_df[col_name].std() + assert math.isclose(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_col",), + ), +) +def test_kurt(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].kurt() + pd_result = scalars_pandas_df[col_name].kurt() + assert math.isclose(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_col",), + ), +) +def test_var(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].var() + pd_result = scalars_pandas_df[col_name].var() + assert math.isclose(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("bool_col",), + ("int64_col",), + ), +) +def test_mode_stat(scalars_df_index, scalars_pandas_df_index, col_name): + bf_result = scalars_df_index[col_name].mode().compute() + pd_result = scalars_pandas_df_index[col_name].mode() + + ## Mode implicitly resets index, and bigframes default indices use nullable Int64 + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x, y: x + y), + (lambda x, y: x - y), + (lambda x, y: x * y), + (lambda x, y: x / y), + (lambda x, y: x // y), + (lambda x, y: x < y), + (lambda x, y: x > y), + (lambda x, y: x <= y), + (lambda x, y: x >= y), + ], + ids=[ + "add", + "subtract", + "multiply", + "divide", + "floordivide", + "less_than", + "greater_than", + "less_than_equal", + "greater_than_equal", + ], +) +@pytest.mark.parametrize(("other_scalar"), [-1, 0, 14, pd.NA]) +@pytest.mark.parametrize(("reverse_operands"), [True, False]) +def test_series_int_int_operators_scalar( + scalars_dfs, operator, other_scalar, reverse_operands +): + scalars_df, scalars_pandas_df = scalars_dfs + + maybe_reversed_op = (lambda x, y: operator(y, x)) if reverse_operands else operator + + bf_result = maybe_reversed_op(scalars_df["int64_col"], other_scalar).compute() + pd_result = maybe_reversed_op(scalars_pandas_df["int64_col"], other_scalar) + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x, y: x & y), + (lambda x, y: x | y), + ], + ids=[ + "and", + "or", + ], +) +@pytest.mark.parametrize(("other_scalar"), [True, False, pd.NA]) +@pytest.mark.parametrize(("reverse_operands"), [True, False]) +def test_series_bool_bool_operators_scalar( + scalars_dfs, operator, other_scalar, reverse_operands +): + scalars_df, scalars_pandas_df = scalars_dfs + + maybe_reversed_op = (lambda x, y: operator(y, x)) if reverse_operands else operator + + bf_result = maybe_reversed_op(scalars_df["bool_col"], other_scalar).compute() + pd_result = maybe_reversed_op(scalars_pandas_df["bool_col"], other_scalar) + + assert_series_equal_ignoring_order(pd_result.astype(pd.BooleanDtype()), bf_result) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x, y: x + y), + (lambda x, y: x - y), + (lambda x, y: x * y), + (lambda x, y: x / y), + (lambda x, y: x < y), + (lambda x, y: x > y), + (lambda x, y: x <= y), + (lambda x, y: x >= y), + (lambda x, y: x % y), + (lambda x, y: x // y), + (lambda x, y: x & y), + (lambda x, y: x | y), + ], + ids=[ + "add", + "subtract", + "multiply", + "divide", + "less_than", + "greater_than", + "less_than_equal", + "greater_than_equal", + "modulo", + "floordivide", + "bitwise_and", + "bitwise_or", + ], +) +def test_series_int_int_operators_series(scalars_dfs, operator): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = operator(scalars_df["int64_col"], scalars_df["int64_too"]).compute() + pd_result = operator(scalars_pandas_df["int64_col"], scalars_pandas_df["int64_too"]) + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("other",), + [ + (3,), + (-6.2,), + ], +) +def test_series_add_scalar(scalars_dfs, other): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (scalars_df["float64_col"] + other).compute() + pd_result = scalars_pandas_df["float64_col"] + other + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("left_col", "right_col"), + [ + ("float64_col", "float64_col"), + ("int64_col", "float64_col"), + ("int64_col", "int64_too"), + ], +) +def test_series_add_bigframes_series(scalars_dfs, left_col, right_col): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (scalars_df[left_col] + scalars_df[right_col]).compute() + pd_result = scalars_pandas_df[left_col] + scalars_pandas_df[right_col] + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("left_col", "right_col", "righter_col"), + [ + ("float64_col", "float64_col", "float64_col"), + ("int64_col", "int64_col", "int64_col"), + ], +) +def test_series_add_bigframes_series_nested( + scalars_dfs, left_col, right_col, righter_col +): + """Test that we can correctly add multiple times.""" + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = ( + (scalars_df[left_col] + scalars_df[right_col]) + scalars_df[righter_col] + ).compute() + pd_result = ( + scalars_pandas_df[left_col] + scalars_pandas_df[right_col] + ) + scalars_pandas_df[righter_col] + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_series_add_different_table_default_index( + scalars_df_default_index, + scalars_df_2_default_index, +): + bf_result = ( + scalars_df_default_index["float64_col"] + + scalars_df_2_default_index["float64_col"] + ).compute() + pd_result = ( + # Default index may not have a well defined order, but it should at + # least be consistent across compute() calls. + scalars_df_default_index["float64_col"].compute() + + scalars_df_2_default_index["float64_col"].compute() + ) + # TODO(swast): Can remove sort_index() when there's default ordering. + pd.testing.assert_series_equal(bf_result.sort_index(), pd_result.sort_index()) + + +def test_series_add_different_table_with_index( + scalars_df_index, scalars_df_2_index, scalars_pandas_df_index +): + scalars_pandas_df = scalars_pandas_df_index + bf_result = scalars_df_index["float64_col"] + scalars_df_2_index["int64_col"] + # When index values are unique, we can emulate with values from the same + # DataFrame. + pd_result = scalars_pandas_df["float64_col"] + scalars_pandas_df["int64_col"] + pd.testing.assert_series_equal(bf_result.compute(), pd_result) + + +def test_reset_index_drop(scalars_df_index, scalars_pandas_df_index): + scalars_pandas_df = scalars_pandas_df_index + bf_result = ( + scalars_df_index["float64_col"] + .sort_index(ascending=False) + .reset_index(drop=True) + ).iloc[::2] + pd_result = ( + scalars_pandas_df["float64_col"] + .sort_index(ascending=False) + .reset_index(drop=True) + ).iloc[::2] + + # BigQuery DataFrames default indices use nullable Int64 always + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_series_equal(bf_result.compute(), pd_result) + + +@pytest.mark.parametrize( + ("name",), + [ + ("some_name",), + (None,), + ], +) +def test_reset_index_no_drop(scalars_df_index, scalars_pandas_df_index, name): + scalars_pandas_df = scalars_pandas_df_index + kw_args = {"name": name} if name else {} + bf_result = ( + scalars_df_index["float64_col"] + .sort_index(ascending=False) + .reset_index(drop=False, **kw_args) + ) + pd_result = ( + scalars_pandas_df["float64_col"] + .sort_index(ascending=False) + .reset_index(drop=False, **kw_args) + ) + + # BigQuery DataFrames default indices use nullable Int64 always + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_frame_equal(bf_result.compute(), pd_result) + + +def test_series_add_pandas_series_not_implemented(scalars_dfs): + scalars_df, _ = scalars_dfs + with pytest.raises(NotImplementedError): + ( + scalars_df["float64_col"] + + pd.Series( + [1, 1, 1, 1], + ) + ).compute() + + +def test_copy(scalars_df_index, scalars_pandas_df_index): + col_name = "float64_col" + # Expect mutation on original not to effect_copy + bf_series = scalars_df_index[col_name].copy() + bf_copy = bf_series.copy() + bf_copy.loc[0] = 5.6 + bf_series.loc[0] = 3.4 + + pd_series = scalars_pandas_df_index[col_name].copy() + pd_copy = pd_series.copy() + pd_copy.loc[0] = 5.6 + pd_series.loc[0] = 3.4 + + assert bf_copy.compute().loc[0] != bf_series.compute().loc[0] + pd.testing.assert_series_equal(bf_copy.compute(), pd_copy) + + +def test_isnull(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "float64_col" + bf_series = scalars_df[col_name].isnull().compute() + pd_series = scalars_pandas_df[col_name].isnull() + + # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but + # the `pd_series.dtype` is `bool`. + assert_series_equal_ignoring_order(pd_series.astype(pd.BooleanDtype()), bf_series) + + +def test_notnull(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series = scalars_df[col_name].notnull().compute() + pd_series = scalars_pandas_df[col_name].notnull() + + # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but + # the `pd_series.dtype` is `bool`. + assert_series_equal_ignoring_order(pd_series.astype(pd.BooleanDtype()), bf_series) + + +def test_round(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "float64_col" + bf_result = scalars_df[col_name].round().compute() + pd_result = scalars_pandas_df[col_name].round() + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_eq_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_result = scalars_df[col_name].eq(0).compute() + pd_result = scalars_pandas_df[col_name].eq(0) + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_eq_wider_type_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_result = scalars_df[col_name].eq(1.0).compute() + pd_result = scalars_pandas_df[col_name].eq(1.0) + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_ne_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_result = (scalars_df[col_name] != 0).compute() + pd_result = scalars_pandas_df[col_name] != 0 + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_eq_int_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_result = (scalars_df[col_name] == 0).compute() + pd_result = scalars_pandas_df[col_name] == 0 + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("string_col",), + ("float64_col",), + ("int64_too",), + ), +) +def test_eq_same_type_series(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = (scalars_df[col_name] == scalars_df[col_name]).compute() + pd_result = scalars_pandas_df[col_name] == scalars_pandas_df[col_name] + + # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but + # the `pd_series.dtype` is `bool`. + assert_series_equal_ignoring_order(pd_result.astype(pd.BooleanDtype()), bf_result) + + +def test_loc_setitem_cell(scalars_df_index, scalars_pandas_df_index): + bf_original = scalars_df_index["string_col"] + bf_series = scalars_df_index["string_col"] + pd_original = scalars_pandas_df_index["string_col"] + pd_series = scalars_pandas_df_index["string_col"].copy() + bf_series.loc[2] = "This value isn't in the test data." + pd_series.loc[2] = "This value isn't in the test data." + bf_result = bf_series.compute() + pd_result = pd_series + pd.testing.assert_series_equal(bf_result, pd_result) + # Per Copy-on-Write semantics, other references to the original DataFrame + # should remain unchanged. + pd.testing.assert_series_equal(bf_original.compute(), pd_original) + + +def test_ne_obj_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = (scalars_df[col_name] != scalars_df[col_name]).compute() + pd_result = scalars_pandas_df[col_name] != scalars_pandas_df[col_name] + + # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but + # the `pd_series.dtype` is `bool`. + assert_series_equal_ignoring_order(pd_result.astype(pd.BooleanDtype()), bf_result) + + +def test_indexing_using_unselected_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = scalars_df[col_name][scalars_df["int64_too"].eq(0)].compute() + pd_result = scalars_pandas_df[col_name][scalars_pandas_df["int64_too"].eq(0)] + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_indexing_using_selected_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = scalars_df[col_name][ + scalars_df["string_col"].eq("Hello, World!") + ].compute() + pd_result = scalars_pandas_df[col_name][ + scalars_pandas_df["string_col"].eq("Hello, World!") + ] + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_nested_filter(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + string_col = scalars_df["string_col"] + int64_too = scalars_df["int64_too"] + bool_col = scalars_df["bool_col"] == bool( + True + ) # Convert from nullable bool to nonnullable bool usable as indexer + bf_result = string_col[int64_too == 0][~bool_col].compute() + + pd_string_col = scalars_pandas_df["string_col"] + pd_int64_too = scalars_pandas_df["int64_too"] + pd_bool_col = scalars_pandas_df["bool_col"] == bool( + True + ) # Convert from nullable bool to nonnullable bool usable as indexer + pd_result = pd_string_col[pd_int64_too == 0][~pd_bool_col] + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_binop_opposite_filters(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + int64_col1 = scalars_df["int64_col"] + int64_col2 = scalars_df["int64_col"] + bool_col = scalars_df["bool_col"] + bf_result = (int64_col1[bool_col] + int64_col2[bool_col.__invert__()]).compute() + + pd_int64_col1 = scalars_pandas_df["int64_col"] + pd_int64_col2 = scalars_pandas_df["int64_col"] + pd_bool_col = scalars_pandas_df["bool_col"] + pd_result = pd_int64_col1[pd_bool_col] + pd_int64_col2[pd_bool_col.__invert__()] + + assert_series_equal_ignoring_order( + bf_result, + pd_result, + ) + + +def test_binop_left_filtered(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + int64_col = scalars_df["int64_col"] + float64_col = scalars_df["float64_col"] + bool_col = scalars_df["bool_col"] + bf_result = (int64_col[bool_col] + float64_col).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_float64_col = scalars_pandas_df["float64_col"] + pd_bool_col = scalars_pandas_df["bool_col"] + pd_result = pd_int64_col[pd_bool_col] + pd_float64_col + + assert_series_equal_ignoring_order( + bf_result, + pd_result, + ) + + +def test_binop_right_filtered(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + int64_col = scalars_df["int64_col"] + float64_col = scalars_df["float64_col"] + bool_col = scalars_df["bool_col"] + bf_result = (float64_col + int64_col[bool_col]).compute() + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_float64_col = scalars_pandas_df["float64_col"] + pd_bool_col = scalars_pandas_df["bool_col"] + pd_result = pd_float64_col + pd_int64_col[pd_bool_col] + + assert_series_equal_ignoring_order( + bf_result, + pd_result, + ) + + +def test_mean(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].mean() + pd_result = scalars_pandas_df[col_name].mean() + assert math.isclose(pd_result, bf_result) + + +def test_repr(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + if scalars_pandas_df.index.name != "rowindex": + pytest.skip("Require index & ordering for consistent repr.") + + col_name = "int64_col" + bf_series = scalars_df[col_name] + pd_series = scalars_pandas_df[col_name] + assert repr(bf_series) == repr(pd_series) + + +def test_sum(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].sum() + pd_result = scalars_pandas_df[col_name].sum() + assert pd_result == bf_result + + +def test_product(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "float64_col" + bf_result = scalars_df[col_name].product() + pd_result = scalars_pandas_df[col_name].product() + assert math.isclose(pd_result, bf_result) + + +def test_count(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].count() + pd_result = scalars_pandas_df[col_name].count() + assert pd_result == bf_result + + +def test_nunique(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = (scalars_df[col_name] % 3).nunique() + pd_result = (scalars_pandas_df[col_name] % 3).nunique() + assert pd_result == bf_result + + +def test_all(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].all() + pd_result = scalars_pandas_df[col_name].all() + assert pd_result == bf_result + + +def test_any(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].any() + pd_result = scalars_pandas_df[col_name].any() + assert pd_result == bf_result + + +def test_groupby_sum(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = scalars_df[col_name].groupby(scalars_df["string_col"]).sum() + pd_series = ( + scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).sum() + ) + # TODO(swast): Update groupby to use index based on group by key(s). + bf_result = bf_series.compute() + assert_series_equal_ignoring_order( + pd_series, + bf_result, + check_exact=False, + ) + + +def test_groupby_std(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = scalars_df[col_name].groupby(scalars_df["string_col"]).std() + pd_series = ( + scalars_pandas_df[col_name] + .groupby(scalars_pandas_df["string_col"]) + .std() + .astype(pd.Float64Dtype()) + ) + bf_result = bf_series.compute() + assert_series_equal_ignoring_order( + pd_series, + bf_result, + check_exact=False, + ) + + +def test_groupby_var(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = scalars_df[col_name].groupby(scalars_df["string_col"]).var() + pd_series = ( + scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).var() + ) + bf_result = bf_series.compute() + assert_series_equal_ignoring_order( + pd_series, + bf_result, + check_exact=False, + ) + + +def test_groupby_level_sum(scalars_dfs): + # TODO(tbergeron): Use a non-unique index once that becomes possible in tests + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + if scalars_pandas_df.index.name != "rowindex": + pytest.skip("Require index for groupby level.") + + bf_series = scalars_df[col_name].groupby(level=0).sum() + pd_series = scalars_pandas_df[col_name].groupby(level=0).sum() + # TODO(swast): Update groupby to use index based on group by key(s). + pd.testing.assert_series_equal( + pd_series.sort_index(), + bf_series.compute().sort_index(), + ) + + +def test_groupby_level_list_sum(scalars_dfs): + # TODO(tbergeron): Use a non-unique index once that becomes possible in tests + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + if scalars_pandas_df.index.name != "rowindex": + pytest.skip("Require index for groupby level.") + + bf_series = scalars_df[col_name].groupby(level=["rowindex"]).sum() + pd_series = scalars_pandas_df[col_name].groupby(level=["rowindex"]).sum() + # TODO(swast): Update groupby to use index based on group by key(s). + pd.testing.assert_series_equal( + pd_series.sort_index(), + bf_series.compute().sort_index(), + ) + + +def test_groupby_mean(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = ( + scalars_df[col_name].groupby(scalars_df["string_col"], dropna=False).mean() + ) + pd_series = ( + scalars_pandas_df[col_name] + .groupby(scalars_pandas_df["string_col"], dropna=False) + .mean() + ) + # TODO(swast): Update groupby to use index based on group by key(s). + bf_result = bf_series.compute() + assert_series_equal_ignoring_order( + pd_series, + bf_result, + ) + + +def test_groupby_prod(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = scalars_df[col_name].groupby(scalars_df["int64_col"]).prod() + pd_series = ( + scalars_pandas_df[col_name].groupby(scalars_pandas_df["int64_col"]).prod() + ) + # TODO(swast): Update groupby to use index based on group by key(s). + bf_result = bf_series.compute() + assert_series_equal_ignoring_order( + pd_series, + bf_result, + ) + + +@pytest.mark.parametrize( + ("operator"), + [ + (lambda x: x.cumsum()), + (lambda x: x.cumcount()), + (lambda x: x.cummin()), + (lambda x: x.cummax()), + (lambda x: x.cumprod()), + (lambda x: x.diff()), + (lambda x: x.shift(2)), + (lambda x: x.shift(-2)), + ], + ids=[ + "cumsum", + "cumcount", + "cummin", + "cummax", + "cumprod", + "diff", + "shiftpostive", + "shiftnegative", + ], +) +def test_groupby_window_ops(scalars_df_index, scalars_pandas_df_index, operator): + col_name = "int64_col" + group_key = "int64_too" # has some duplicates values, good for grouping + bf_series = ( + operator(scalars_df_index[col_name].groupby(scalars_df_index[group_key])) + ).compute() + pd_series = operator( + scalars_pandas_df_index[col_name].groupby(scalars_pandas_df_index[group_key]) + ).astype(pd.Int64Dtype()) + pd.testing.assert_series_equal( + pd_series, + bf_series, + ) + + +def test_drop_label(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + bf_series = scalars_df_index[col_name].drop(1).compute() + pd_series = scalars_pandas_df_index[col_name].drop(1) + pd.testing.assert_series_equal( + pd_series, + bf_series, + ) + + +def test_drop_label_list(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + bf_series = scalars_df_index[col_name].drop([1, 3]).compute() + pd_series = scalars_pandas_df_index[col_name].drop([1, 3]) + pd.testing.assert_series_equal( + pd_series, + bf_series, + ) + + +@pytest.mark.parametrize( + ("col_name",), + [ + ("bool_col",), + ("int64_too",), + ], +) +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + (False,), + ], +) +def test_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, col_name): + bf_series = scalars_df_index[col_name].drop_duplicates(keep=keep).compute() + pd_series = scalars_pandas_df_index[col_name].drop_duplicates(keep=keep) + pd.testing.assert_series_equal( + pd_series, + bf_series, + ) + + +@pytest.mark.parametrize( + ("col_name",), + [ + ("bool_col",), + ("int64_too",), + ], +) +def test_unique(scalars_df_index, scalars_pandas_df_index, col_name): + bf_uniq = scalars_df_index[col_name].unique().to_numpy() + pd_uniq = scalars_pandas_df_index[col_name].unique() + numpy.array_equal(pd_uniq, bf_uniq) + + +@pytest.mark.parametrize( + ("col_name",), + [ + ("bool_col",), + ("int64_too",), + ], +) +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + (False,), + ], +) +def test_duplicated(scalars_df_index, scalars_pandas_df_index, keep, col_name): + bf_series = scalars_df_index[col_name].duplicated(keep=keep).compute() + pd_series = scalars_pandas_df_index[col_name].duplicated(keep=keep) + pd.testing.assert_series_equal(pd_series, bf_series, check_dtype=False) + + +def test_shape(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["string_col"].shape + pd_result = scalars_pandas_df["string_col"].shape + + assert pd_result == bf_result + + +def test_len(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = len(scalars_df["string_col"]) + pd_result = len(scalars_pandas_df["string_col"]) + + assert pd_result == bf_result + + +def test_size(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["string_col"].size + pd_result = scalars_pandas_df["string_col"].size + + assert pd_result == bf_result + + +def test_empty_false(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["string_col"].empty + pd_result = scalars_pandas_df["string_col"].empty + + assert pd_result == bf_result + + +def test_empty_true(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["string_col"][ + scalars_df["string_col"] == "won't find this" + ].empty + pd_result = scalars_pandas_df["string_col"][ + scalars_pandas_df["string_col"] == "won't find this" + ].empty + + assert pd_result == bf_result + + +def test_dtype(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["string_col"].dtype + pd_result = scalars_pandas_df["string_col"].dtype + + assert pd_result == bf_result + + +def test_dtypes(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["int64_col"].dtypes + pd_result = scalars_pandas_df["int64_col"].dtypes + + assert pd_result == bf_result + + +def test_head(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + if scalars_df.index.name is None: + pytest.skip("Require explicit index for offset ops.") + + bf_result = scalars_df["string_col"].head(2).compute() + pd_result = scalars_pandas_df["string_col"].head(2) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_tail(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + if scalars_df.index.name is None: + pytest.skip("Require explicit index for offset ops.") + + bf_result = scalars_df["string_col"].tail(2).compute() + pd_result = scalars_pandas_df["string_col"].tail(2) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_head_then_scalar_operation(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + if scalars_df.index.name is None: + pytest.skip("Require explicit index for offset ops.") + + bf_result = (scalars_df["float64_col"].head(1) + 4).compute() + pd_result = scalars_pandas_df["float64_col"].head(1) + 4 + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_head_then_series_operation(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + if scalars_df.index.name is None: + pytest.skip("Require explicit index for offset ops.") + + bf_result = ( + scalars_df["float64_col"].head(4) + scalars_df["float64_col"].head(2) + ).compute() + pd_result = scalars_pandas_df["float64_col"].head(4) + scalars_pandas_df[ + "float64_col" + ].head(2) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_shift(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + bf_result = scalars_df_index[col_name].shift().compute() + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index[col_name].shift().astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cumsum_int(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("1."): + pytest.skip("Series.cumsum NA mask are different in pandas 1.x.") + + col_name = "int64_col" + bf_result = scalars_df_index[col_name].cumsum().compute() + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index[col_name].cumsum().astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("na_option",), + [ + ("keep",), + ("top",), + ("bottom",), + ], +) +@pytest.mark.parametrize( + ("method",), + [ + ("average",), + ("min",), + ("max",), + ("first",), + ("dense",), + ], +) +@pytest.mark.skipif( + True, reason="Blocked by possible pandas rank() regression (b/283278923)" +) +def test_rank_with_nulls(scalars_df_index, scalars_pandas_df_index, na_option, method): + col_name = "bool_col" + bf_result = ( + scalars_df_index[col_name].rank(na_option=na_option, method=method).compute() + ) + pd_result = ( + scalars_pandas_df_index[col_name] + .rank(na_option=na_option, method=method) + .astype(pd.Float64Dtype()) + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + ("all",), + ], +) +def test_nlargest(scalars_df_index, scalars_pandas_df_index, keep): + col_name = "bool_col" + bf_result = scalars_df_index[col_name].nlargest(4, keep=keep).compute() + pd_result = scalars_pandas_df_index[col_name].nlargest(4, keep=keep) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + ("all",), + ], +) +def test_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): + col_name = "bool_col" + bf_result = scalars_df_index[col_name].nsmallest(2, keep=keep).compute() + pd_result = scalars_pandas_df_index[col_name].nsmallest(2, keep=keep) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_rank_ints(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_too" + bf_result = scalars_df_index[col_name].rank().compute() + pd_result = scalars_pandas_df_index[col_name].rank().astype(pd.Float64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cumsum_nested(scalars_df_index, scalars_pandas_df_index): + col_name = "float64_col" + bf_result = scalars_df_index[col_name].cumsum().cumsum().cumsum().compute() + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = ( + scalars_pandas_df_index[col_name] + .cumsum() + .cumsum() + .cumsum() + .astype(pd.Float64Dtype()) + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cumsum_int_filtered(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + + bf_col = scalars_df_index[col_name] + bf_result = bf_col[bf_col > -2].cumsum().compute() + + pd_col = scalars_pandas_df_index[col_name] + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = pd_col[pd_col > -2].cumsum().astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cumsum_float(scalars_df_index, scalars_pandas_df_index): + col_name = "float64_col" + bf_result = scalars_df_index[col_name].cumsum().compute() + # cumsum does not behave well on nullable floats in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index[col_name].cumsum().astype(pd.Float64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cummin_int(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + bf_result = scalars_df_index[col_name].cummin().compute() + pd_result = scalars_pandas_df_index[col_name].cummin() + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cummax_int(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_col" + bf_result = scalars_df_index[col_name].cummax().compute() + pd_result = scalars_pandas_df_index[col_name].cummax() + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_value_counts(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + + bf_result = scalars_df[col_name].value_counts().compute() + pd_result = scalars_pandas_df[col_name].value_counts() + + # Older pandas version may not have these values, bigframes tries to emulate 2.0+ + pd_result.name = "count" + pd_result.index.name = col_name + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_iloc_nested(scalars_df_index, scalars_pandas_df_index): + + bf_result = scalars_df_index["string_col"].iloc[1:].iloc[1:].compute() + pd_result = scalars_pandas_df_index["string_col"].iloc[1:].iloc[1:] + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("start", "stop", "step"), + [ + (1, None, None), + (None, 4, None), + (None, None, 2), + (None, 50000000000, 1), + (5, 4, None), + (3, None, 2), + (1, 7, 2), + (1, 7, 50000000000), + (-1, -7, -2), + (None, -7, -2), + (-1, None, -2), + (-7, -1, 2), + (-7, -1, None), + (-7, 7, None), + (7, -7, -2), + ], +) +def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, step): + bf_result = scalars_df_index["string_col"].iloc[start:stop:step].compute() + pd_result = scalars_pandas_df_index["string_col"].iloc[start:stop:step] + + # Pandas may assign non-object dtype to empty series and series index + if pd_result.empty: + pd_result = pd_result.astype("object") + pd_result.index = pd_result.index.astype("object") + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_add_prefix(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].add_prefix("prefix_").compute() + + pd_result = scalars_pandas_df_index["int64_too"].add_prefix("prefix_") + + # Index will be object type in pandas, string type in bigframes, but same values + pd.testing.assert_series_equal( + bf_result, + pd_result, + check_index_type=False, + ) + + +def test_series_add_suffix(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].add_suffix("_suffix").compute() + + pd_result = scalars_pandas_df_index["int64_too"].add_suffix("_suffix") + + # Index will be object type in pandas, string type in bigframes, but same values + pd.testing.assert_series_equal( + bf_result, + pd_result, + check_index_type=False, + ) + + +def test_where_with_series(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_col"] + .where(scalars_df_index["bool_col"], scalars_df_index["int64_too"]) + .compute() + ) + pd_result = scalars_pandas_df_index["int64_col"].where( + scalars_pandas_df_index["bool_col"], scalars_pandas_df_index["int64_too"] + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_where_with_different_indices(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_col"] + .iloc[::2] + .where( + scalars_df_index["bool_col"].iloc[2:], + scalars_df_index["int64_too"].iloc[:5], + ) + .compute() + ) + pd_result = ( + scalars_pandas_df_index["int64_col"] + .iloc[::2] + .where( + scalars_pandas_df_index["bool_col"].iloc[2:], + scalars_pandas_df_index["int64_too"].iloc[:5], + ) + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_where_with_default(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_col"].where(scalars_df_index["bool_col"]).compute() + ) + pd_result = scalars_pandas_df_index["int64_col"].where( + scalars_pandas_df_index["bool_col"] + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_clip(scalars_df_index, scalars_pandas_df_index): + col_bf = scalars_df_index["int64_col"] + lower_bf = scalars_df_index["int64_too"] - 1 + upper_bf = scalars_df_index["int64_too"] + 1 + bf_result = col_bf.clip(lower_bf, upper_bf).compute() + + col_pd = scalars_pandas_df_index["int64_col"] + lower_pd = scalars_pandas_df_index["int64_too"] - 1 + upper_pd = scalars_pandas_df_index["int64_too"] + 1 + pd_result = col_pd.clip(lower_pd, upper_pd) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_clip_filtered_two_sided(scalars_df_index, scalars_pandas_df_index): + col_bf = scalars_df_index["int64_col"].iloc[::2] + lower_bf = scalars_df_index["int64_too"].iloc[2:] - 1 + upper_bf = scalars_df_index["int64_too"].iloc[:5] + 1 + bf_result = col_bf.clip(lower_bf, upper_bf).compute() + + col_pd = scalars_pandas_df_index["int64_col"].iloc[::2] + lower_pd = scalars_pandas_df_index["int64_too"].iloc[2:] - 1 + upper_pd = scalars_pandas_df_index["int64_too"].iloc[:5] + 1 + pd_result = col_pd.clip(lower_pd, upper_pd) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_clip_filtered_one_sided(scalars_df_index, scalars_pandas_df_index): + col_bf = scalars_df_index["int64_col"].iloc[::2] + lower_bf = scalars_df_index["int64_too"].iloc[2:] - 1 + bf_result = col_bf.clip(lower_bf, None).compute() + + col_pd = scalars_pandas_df_index["int64_col"].iloc[::2] + lower_pd = scalars_pandas_df_index["int64_too"].iloc[2:] - 1 + pd_result = col_pd.clip(lower_pd, None) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_dot(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_too"] @ scalars_df["int64_too"] + + pd_result = scalars_pandas_df["int64_too"] @ scalars_pandas_df["int64_too"] + + assert bf_result == pd_result + + +@pytest.mark.parametrize( + ("left", "right", "inclusive"), + [ + (-234892, 55555, "left"), + (-234892, 55555, "both"), + (-234892, 55555, "neither"), + (-234892, 55555, "right"), + ], +) +def test_between(scalars_df_index, scalars_pandas_df_index, left, right, inclusive): + bf_result = scalars_df_index["int64_col"].between(left, right, inclusive).compute() + pd_result = scalars_pandas_df_index["int64_col"].between(left, right, inclusive) + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()), + ) + + +def test_to_frame(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df["int64_col"].to_frame().compute() + pd_result = scalars_pandas_df["int64_col"].to_frame() + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_to_json(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_col"].to_json() + pd_result = scalars_pandas_df_index["int64_col"].to_json() + + assert bf_result == pd_result + + +def test_to_csv(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_col"].to_csv() + pd_result = scalars_pandas_df_index["int64_col"].to_csv() + + assert bf_result == pd_result + + +def test_to_latex(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_col"].to_latex() + pd_result = scalars_pandas_df_index["int64_col"].to_latex() + + assert bf_result == pd_result + + +def test_to_dict(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_dict() + + pd_result = scalars_pandas_df_index["int64_too"].to_dict() + + assert bf_result == pd_result + + +def test_to_excel(scalars_df_index, scalars_pandas_df_index): + bf_result_file = tempfile.TemporaryFile() + pd_result_file = tempfile.TemporaryFile() + scalars_df_index["int64_too"].to_excel(bf_result_file) + scalars_pandas_df_index["int64_too"].to_excel(pd_result_file) + bf_result = bf_result_file.read() + pd_result = bf_result_file.read() + + assert bf_result == pd_result + + +def test_to_pickle(scalars_df_index, scalars_pandas_df_index): + bf_result_file = tempfile.TemporaryFile() + pd_result_file = tempfile.TemporaryFile() + scalars_df_index["int64_too"].to_pickle(bf_result_file) + scalars_pandas_df_index["int64_too"].to_pickle(pd_result_file) + bf_result = bf_result_file.read() + pd_result = bf_result_file.read() + + assert bf_result == pd_result + + +def test_to_string(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_string() + + pd_result = scalars_pandas_df_index["int64_too"].to_string() + + assert bf_result == pd_result + + +def test_to_list(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_list() + + pd_result = scalars_pandas_df_index["int64_too"].to_list() + + assert bf_result == pd_result + + +def test_to_numpy(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_numpy() + + pd_result = scalars_pandas_df_index["int64_too"].to_numpy() + + assert (bf_result == pd_result).all() + + +def test_to_xarray(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_xarray() + + pd_result = scalars_pandas_df_index["int64_too"].to_xarray() + + assert bf_result.equals(pd_result) + + +def test_to_markdown(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].to_markdown() + + pd_result = scalars_pandas_df_index["int64_too"].to_markdown() + + assert bf_result == pd_result + + +def test_series_values(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].values + + pd_result = scalars_pandas_df_index["int64_too"].values + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + pd.testing.assert_series_equal( + pd.Series(bf_result), pd.Series(pd_result), check_dtype=False + ) + + +def test_series___array__(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["float64_col"].__array__() + + pd_result = scalars_pandas_df_index["float64_col"].__array__() + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + numpy.array_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("ascending", "na_position"), + [ + (True, "first"), + (True, "last"), + (False, "first"), + (False, "last"), + ], +) +def test_sort_values(scalars_df_index, scalars_pandas_df_index, ascending, na_position): + # Test needs values to be unique + bf_result = ( + scalars_df_index["int64_col"] + .sort_values(ascending=ascending, na_position=na_position) + .compute() + ) + pd_result = scalars_pandas_df_index["int64_col"].sort_values( + ascending=ascending, na_position=na_position + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("ascending"), + [ + (True,), + (False,), + ], +) +def test_sort_index(scalars_df_index, scalars_pandas_df_index, ascending): + bf_result = scalars_df_index["int64_too"].sort_index(ascending=ascending).compute() + pd_result = scalars_pandas_df_index["int64_too"].sort_index(ascending=ascending) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_mask_default_value(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_col = scalars_df["int64_col"] + bf_col_masked = bf_col.mask(bf_col % 2 == 1) + bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).compute() + + pd_col = scalars_pandas_df["int64_col"] + pd_col_masked = pd_col.mask(pd_col % 2 == 1) + pd_result = pd_col.to_frame().assign(int64_col_masked=pd_col_masked) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_mask_custom_value(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_col = scalars_df["int64_col"] + bf_col_masked = bf_col.mask(bf_col % 2 == 1, -1) + bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).compute() + + pd_col = scalars_pandas_df["int64_col"] + pd_col_masked = pd_col.mask(pd_col % 2 == 1, -1) + pd_result = pd_col.to_frame().assign(int64_col_masked=pd_col_masked) + + # TODO(shobs): There is a pd.NA value in the original series, which is not + # odd so should be left as is, but it is being masked in pandas. + # Accidentally the bigframes bahavior matches, but it should be updated + # after the resolution of https://github.com/pandas-dev/pandas/issues/52955 + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("column", "to_type"), + [ + ("int64_col", "Float64"), + ("int64_col", "Int64"), # No-op + ("int64_col", pd.Float64Dtype()), + ("int64_col", "string[pyarrow]"), + ("int64_col", "boolean"), + ("bool_col", "Int64"), + ("bool_col", "string[pyarrow]"), + # pandas actually doesn't let folks convert to/from naive timestamp and + # raises a deprecation warning to use tz_localize/tz_convert instead, + # but BigQuery always stores values as UTC and doesn't have to deal + # with timezone conversions, so we'll allow it. + ("timestamp_col", pd.ArrowDtype(pa.timestamp("us"))), + ("datetime_col", pd.ArrowDtype(pa.timestamp("us", tz="UTC"))), + # TODO(bmil): fix Ibis bug: BigQuery backend rounds to nearest int + # ("float64_col", "Int64"), + # TODO(bmil): decide whether to fix Ibis bug: BigQuery backend + # formats floats with no decimal places if they have no fractional + # part, and does not switch to scientific notation for > 10^15 + # ("float64_col", "string[pyarrow]") + # TODO(bmil): add any other compatible conversions per + # https://cloud.google.com/bigquery/docs/reference/standard-sql/conversion_functions + ], +) +def test_astype(scalars_df_index, scalars_pandas_df_index, column, to_type): + bf_result = scalars_df_index[column].astype(to_type).compute() + pd_result = scalars_pandas_df_index[column].astype(to_type) + pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + "index", + [0, 5], +) +def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): + bf_result = scalars_df_index.string_col.iloc[index] + pd_result = scalars_pandas_df_index.string_col.iloc[index] + + assert bf_result == pd_result + + +def test_iloc_single_integer_out_of_bound_error( + scalars_df_index, scalars_pandas_df_index +): + with pytest.raises(IndexError, match="single positional indexer is out-of-bounds"): + scalars_df_index.string_col.iloc[99] + + +def test_loc_bool_series_explicit_index(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.loc[scalars_df_index.bool_col].compute() + pd_result = scalars_pandas_df_index.string_col.loc[scalars_pandas_df_index.bool_col] + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_loc_bool_series_default_index( + scalars_df_default_index, scalars_pandas_df_default_index +): + bf_result = scalars_df_default_index.string_col.loc[ + scalars_df_default_index.bool_col + ].compute() + pd_result = scalars_pandas_df_default_index.string_col.loc[ + scalars_pandas_df_default_index.bool_col + ] + + assert_pandas_df_equal_ignore_ordering( + bf_result.to_frame(), + pd_result.to_frame(), + ) + + +def test_argmin(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.argmin() + pd_result = scalars_pandas_df_index.string_col.argmin() + assert bf_result == pd_result + + +def test_argmax(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.int64_too.argmax() + pd_result = scalars_pandas_df_index.int64_too.argmax() + assert bf_result == pd_result + + +def test_getattr_not_implemented(scalars_df_index): + with pytest.raises(NotImplementedError): + scalars_df_index.string_col.asof() + + +def test_getattr_attribute_error(scalars_df_index): + with pytest.raises(AttributeError): + scalars_df_index.string_col.not_a_method() + + +def test_rename(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.rename("newname") + pd_result = scalars_pandas_df_index.string_col.rename("newname") + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_rename_axis(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.rename_axis("newindexname") + pd_result = scalars_pandas_df_index.string_col.rename_axis("newindexname") + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index): + index_list = scalars_pandas_df_index.string_col.iloc[[0, 1, 1, 5]].values + + scalars_df_index = scalars_df_index.set_index("string_col", drop=False) + scalars_pandas_df_index = scalars_pandas_df_index.set_index( + "string_col", drop=False + ) + + bf_result = scalars_df_index.string_col.loc[index_list] + pd_result = scalars_pandas_df_index.string_col.loc[index_list] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): + index_list = [3, 2, 1, 3, 2, 1] + + bf_result = scalars_df_index.bool_col.loc[index_list] + pd_result = scalars_pandas_df_index.bool_col.loc[index_list] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_iloc_list(scalars_df_index, scalars_pandas_df_index): + index_list = [0, 0, 0, 5, 4, 7] + + bf_result = scalars_df_index.string_col.iloc[index_list] + pd_result = scalars_pandas_df_index.string_col.iloc[index_list] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_iloc_list_nameless(scalars_df_index, scalars_pandas_df_index): + index_list = [0, 0, 0, 5, 4, 7] + + bf_series = scalars_df_index.string_col.rename(None) + bf_result = bf_series.iloc[index_list] + pd_series = scalars_pandas_df_index.string_col.rename(None) + pd_result = pd_series.iloc[index_list] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_list_nameless(scalars_df_index, scalars_pandas_df_index): + index_list = [0, 0, 0, 5, 4, 7] + + bf_series = scalars_df_index.string_col.rename(None) + bf_result = bf_series.loc[index_list] + + pd_series = scalars_pandas_df_index.string_col.rename(None) + pd_result = pd_series.loc[index_list] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): + pd_string_series = scalars_pandas_df_index.string_col.iloc[[0, 5, 1, 1, 5]] + bf_string_series = scalars_df_index.string_col.iloc[[0, 5, 1, 1, 5]] + + scalars_df_index = scalars_df_index.set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.set_index("string_col") + + bf_result = scalars_df_index.date_col.loc[bf_string_series] + pd_result = scalars_pandas_df_index.date_col.loc[pd_string_series] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_loc_bf_index_integer_index(scalars_df_index, scalars_pandas_df_index): + pd_index = scalars_pandas_df_index.iloc[[0, 5, 1, 1, 5]].index + bf_index = scalars_df_index.iloc[[0, 5, 1, 1, 5]].index + + bf_result = scalars_df_index.date_col.loc[bf_index] + pd_result = scalars_pandas_df_index.date_col.loc[pd_index] + + pd.testing.assert_series_equal( + bf_result.compute(), + pd_result, + ) + + +def test_series_bool_interpretation_error(scalars_df_index): + with pytest.raises(ValueError): + True if scalars_df_index["string_col"] else False diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py new file mode 100644 index 00000000000..31ab4dee121 --- /dev/null +++ b/tests/system/small/test_session.py @@ -0,0 +1,673 @@ +# 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 random +import tempfile +import typing +from typing import List + +import google.api_core.exceptions +import numpy as np +import pandas as pd +import pytest + +import bigframes +import bigframes.core.indexes.index +import bigframes.dataframe +import bigframes.dtypes +import bigframes.ml.linear_model + + +def test_read_gbq_tokyo( + session_tokyo: bigframes.Session, + scalars_table_tokyo: str, + scalars_pandas_df_index: pd.DataFrame, + tokyo_location: str, +): + df = session_tokyo.read_gbq(scalars_table_tokyo, index_col=["rowindex"]) + result = df.sort_index().compute() + expected = scalars_pandas_df_index + + _, query_job = df._block.expr.start_query() + assert query_job.location == tokyo_location + + pd.testing.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize( + ("query_or_table", "col_order"), + [ + pytest.param( + "{scalars_table_id}", ["bool_col", "int64_col"], id="two_cols_in_table" + ), + pytest.param( + """SELECT + t.float64_col * 2 AS my_floats, + CONCAT(t.string_col, "_2") AS my_strings, + t.int64_col > 0 AS my_bools, + FROM `{scalars_table_id}` AS t + """, + ["my_strings"], + id="one_cols_in_query", + ), + pytest.param( + "{scalars_table_id}", + ["unknown"], + marks=pytest.mark.xfail( + raises=ValueError, + reason="Column `unknown` not found in this table.", + ), + id="unknown_col", + ), + ], +) +def test_read_gbq_w_col_order( + session: bigframes.Session, + scalars_table_id: str, + query_or_table: str, + col_order: List[str], +): + df = session.read_gbq( + query_or_table.format(scalars_table_id=scalars_table_id), col_order=col_order + ) + assert df.columns.tolist() == col_order + + +@pytest.mark.parametrize( + ("query_or_table", "index_col"), + [ + pytest.param("{scalars_table_id}", ["bool_col", "int64_col"], id="multiindex"), + pytest.param( + """SELECT + t.float64_col * 2 AS my_floats, + CONCAT(t.string_col, "_2") AS my_strings, + t.int64_col > 0 AS my_bools, + FROM `{scalars_table_id}` AS t + """, + ["my_strings"], + id="string_index", + ), + pytest.param( + "{scalars_table_id}", + ["bool_col"], + id="non_unique_index", + ), + ], +) +def test_read_gbq_w_index_col( + session: bigframes.Session, + scalars_table_id: str, + query_or_table: str, + index_col: List[str], +): + df = session.read_gbq( + query_or_table.format(scalars_table_id=scalars_table_id), + index_col=index_col, + ) + assert list(df.index.names) == index_col + + +@pytest.mark.parametrize( + ("query_or_table", "max_results"), + [ + pytest.param("{scalars_table_id}", 2, id="two_rows_in_table"), + pytest.param( + """SELECT + t.float64_col * 2 AS my_floats, + CONCAT(t.string_col, "_2") AS my_strings, + t.int64_col > 0 AS my_bools, + FROM `{scalars_table_id}` AS t + """, + 2, + id="three_rows_in_query", + ), + pytest.param( + "{scalars_table_id}", + -1, + marks=pytest.mark.xfail( + raises=ValueError, + reason="`max_results` should be a positive number.", + ), + id="neg_rows", + ), + ], +) +def test_read_gbq_w_max_results( + session: bigframes.Session, + scalars_table_id: str, + query_or_table: str, + max_results: int, +): + df = session.read_gbq( + query_or_table.format(scalars_table_id=scalars_table_id), + max_results=max_results, + ) + bf_result = df.compute() + assert bf_result.shape[0] == max_results + + +def test_read_gbq_w_script(session, dataset_id: str): + ddl = f""" + CREATE TABLE `{dataset_id}.test_read_gbq_w_ddl` ( + `col_a` INT64, + `col_b` STRING + ); + + INSERT INTO `{dataset_id}.test_read_gbq_w_ddl` + VALUES (123, 'hello world'); + """ + df = session.read_gbq(ddl).to_pandas() + assert df["statement_type"][0] == "SCRIPT" + + +def test_read_gbq_model(session, penguins_linear_model_name): + model = session.read_gbq_model(penguins_linear_model_name) + assert isinstance(model, bigframes.ml.linear_model.LinearRegression) + + +def test_read_pandas(session, scalars_dfs): + _, scalars_pandas_df = scalars_dfs + + df = session.read_pandas(scalars_pandas_df) + assert df._block._expr._ordering is not None + + result = df.compute() + expected = scalars_pandas_df + + pd.testing.assert_frame_equal(result, expected) + + +def test_read_pandas_multi_index(session, scalars_pandas_df_multi_index): + df = session.read_pandas(scalars_pandas_df_multi_index) + result = df.compute() + pd.testing.assert_frame_equal(result, scalars_pandas_df_multi_index) + + +def test_read_pandas_rowid_exists_adds_suffix(session, scalars_pandas_df_default_index): + scalars_pandas_df_default_index["rowid"] = np.arange( + scalars_pandas_df_default_index.shape[0] + ) + + df = session.read_pandas(scalars_pandas_df_default_index) + assert df._block._expr._ordering.ordering_id == "rowid_2" + + +def test_read_pandas_tokyo( + session_tokyo: bigframes.Session, + scalars_pandas_df_index: pd.DataFrame, + tokyo_location: str, +): + df = session_tokyo.read_pandas(scalars_pandas_df_index) + result = df.compute() + expected = scalars_pandas_df_index + + _, query_job = df._block.expr.start_query() + assert query_job.location == tokyo_location + + pd.testing.assert_frame_equal(result, expected) + + +def test_read_csv_gcs_default_engine(session, scalars_dfs, gcs_folder): + scalars_df, _ = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index.csv" + else: + path = gcs_folder + "test_read_csv_gcs_default_engine_wo_index.csv" + scalars_df.to_csv(path, index=False) + dtype = scalars_df.dtypes.to_dict() + dtype.pop("geography_col") + df = session.read_csv( + path, + # Convert default pandas dtypes to match BigQuery DataFrames dtypes. + dtype=dtype, + ) + assert df._block._expr._ordering is not None + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # numeric and geometry types, so they're skipped here. + df = df.drop(columns=["bytes_col", "numeric_col", "geography_col"]) + scalars_df = scalars_df.drop(columns=["bytes_col", "numeric_col", "geography_col"]) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + +def test_read_csv_gcs_bq_engine(session, scalars_dfs, gcs_folder): + scalars_df, _ = scalars_dfs + if scalars_df.index.name is not None: + path = gcs_folder + "test_read_csv_gcs_bq_engine_w_index.csv" + else: + path = gcs_folder + "test_read_csv_gcs_bq_engine_wo_index.csv" + scalars_df.to_csv(path, index=False) + df = session.read_csv(path, engine="bigquery") + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # datetime, numeric and geometry types, so they're skipped here. + df = df.drop(columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"]) + scalars_df = scalars_df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + +@pytest.mark.parametrize( + "sep", + [ + pytest.param(",", id="default_sep"), + pytest.param("\t", id="custom_sep"), + ], +) +def test_read_csv_local_default_engine(session, scalars_dfs, sep): + scalars_df, scalars_pandas_df = scalars_dfs + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_default_engine.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df.to_csv(path, index=False, sep=sep) + dtype = scalars_df.dtypes.to_dict() + dtype.pop("geography_col") + df = session.read_csv( + path, + sep=sep, + # Convert default pandas dtypes to match BigQuery DataFrames dtypes. + dtype=dtype, + ) + assert df._block._expr._ordering is not None + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # numeric and geometry types, so they're skipped here. + df = df.drop(columns=["bytes_col", "numeric_col", "geography_col"]) + scalars_df = scalars_df.drop( + columns=["bytes_col", "numeric_col", "geography_col"] + ) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + +@pytest.mark.parametrize( + "sep", + [ + pytest.param(",", id="default_sep"), + pytest.param("\t", id="custom_sep"), + ], +) +def test_read_csv_local_bq_engine(session, scalars_dfs, sep): + scalars_df, scalars_pandas_df = scalars_dfs + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_bq_engine.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df.to_csv(path, index=False, sep=sep) + df = session.read_csv(path, engine="bigquery", sep=sep) + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # datetime, numeric and geometry types, so they're skipped here. + df = df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + scalars_df = scalars_df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + +def test_read_csv_localbuffer_bq_engine(session, scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_bq_engine.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df.to_csv(path, index=False) + with open(path, "rb") as buffer: + df = session.read_csv(buffer, engine="bigquery") + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # datetime, numeric and geometry types, so they're skipped here. + df = df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + scalars_df = scalars_df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + pytest.param( + {"engine": "bigquery", "names": []}, + "BigQuery engine does not support these arguments", + id="with_names", + ), + pytest.param( + {"engine": "bigquery", "dtype": {}}, + "BigQuery engine does not support these arguments", + id="with_dtype", + ), + pytest.param( + {"engine": "bigquery", "index_col": False}, + "BigQuery engine only supports a single column name for `index_col`.", + id="with_index_col_false", + ), + pytest.param( + {"engine": "bigquery", "index_col": 5}, + "BigQuery engine only supports a single column name for `index_col`.", + id="with_index_col_not_str", + ), + pytest.param( + {"engine": "bigquery", "usecols": [1, 2]}, + "BigQuery engine only supports an iterable of strings for `usecols`.", + id="with_usecols_invalid", + ), + pytest.param( + {"engine": "bigquery", "encoding": "ASCII"}, + "BigQuery engine only supports the following encodings", + id="with_encoding_invalid", + ), + ], +) +def test_read_csv_bq_engine_throws_not_implemented_error(session, kwargs, match): + with pytest.raises(NotImplementedError, match=match): + session.read_csv("", **kwargs) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + pytest.param( + {"chunksize": 5}, + "'chunksize' and 'iterator' arguments are not supported.", + id="with_chunksize", + ), + pytest.param( + {"iterator": True}, + "'chunksize' and 'iterator' arguments are not supported.", + id="with_iterator", + ), + ], +) +def test_read_csv_default_engine_throws_not_implemented_error( + session, + scalars_df_index, + gcs_folder, + kwargs, + match, +): + path = ( + gcs_folder + "test_read_csv_gcs_default_engine_throws_not_implemented_error.csv" + ) + scalars_df_index.to_csv(path) + with pytest.raises(NotImplementedError, match=match): + session.read_csv(path, **kwargs) + + +def test_read_csv_gcs_default_engine_w_header(session, scalars_df_index, gcs_folder): + path = gcs_folder + "test_read_csv_gcs_default_engine_w_header.csv" + scalars_df_index.to_csv(path) + + # Skips header=N rows, normally considers the N+1th row as the header, but overridden by + # passing the `names` argument. In this case, pandas will skip the N+1th row too, take + # the column names from `names`, and begin reading data from the N+2th row. + df = session.read_csv( + path, + header=2, + names=scalars_df_index.columns.to_list(), + ) + assert df.shape[0] == scalars_df_index.shape[0] - 2 + assert len(df.columns) == len(scalars_df_index.columns) + + +def test_read_csv_gcs_bq_engine_w_header(session, scalars_df_index, gcs_folder): + path = gcs_folder + "test_read_csv_gcs_bq_engine_w_header.csv" + scalars_df_index.to_csv(path, index=False) + + # Skip the header and the first 2 data rows. Without provided schema, the column names + # would be like `bool_field_0`, `string_field_1` and etc. + df = session.read_csv(path, header=2, engine="bigquery") + assert df.shape[0] == scalars_df_index.shape[0] - 2 + assert len(df.columns) == len(scalars_df_index.columns) + + +def test_read_csv_local_default_engine_w_header(session, scalars_pandas_df_index): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_default_engine_w_header.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df_index.to_csv(path, index=False) + + # Skips header=N rows. Normally row N+1 would be the header now, but overridden by + # passing the `names` argument. In this case, pandas will skip row N+1 too, infer + # the column names from `names`, and begin reading data from row N+2. + df = session.read_csv( + path, + header=2, + names=scalars_pandas_df_index.columns.to_list(), + ) + assert df.shape[0] == scalars_pandas_df_index.shape[0] - 2 + assert len(df.columns) == len(scalars_pandas_df_index.columns) + + +def test_read_csv_local_bq_engine_w_header(session, scalars_pandas_df_index): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_bq_engine_w_header.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df_index.to_csv(path, index=False) + + # Skip the header and the first 2 data rows. Without provided schema, the column names + # would be like `bool_field_0`, `string_field_1` and etc. + df = session.read_csv(path, header=2, engine="bigquery") + assert df.shape[0] == scalars_pandas_df_index.shape[0] - 2 + assert len(df.columns) == len(scalars_pandas_df_index.columns) + + +def test_read_csv_gcs_default_engine_w_index_col_name( + session, scalars_df_default_index, gcs_folder +): + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_name.csv" + scalars_df_default_index.to_csv(path) + + df = session.read_csv(path, index_col="rowindex") + scalars_df_default_index = scalars_df_default_index.set_index( + "rowindex" + ).sort_index() + pd.testing.assert_index_equal(df.columns, scalars_df_default_index.columns) + assert df.index.name == "rowindex" + + +def test_read_csv_gcs_default_engine_w_index_col_index( + session, scalars_df_default_index, gcs_folder +): + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_index.csv" + scalars_df_default_index.to_csv(path) + + index_col = scalars_df_default_index.columns.to_list().index("rowindex") + df = session.read_csv(path, index_col=index_col) + scalars_df_default_index = scalars_df_default_index.set_index( + "rowindex" + ).sort_index() + pd.testing.assert_index_equal(df.columns, scalars_df_default_index.columns) + assert df.index.name == "rowindex" + + +def test_read_csv_local_default_engine_w_index_col_name( + session, scalars_pandas_df_default_index +): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_default_engine_w_index_col_name" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df_default_index.to_csv(path, index=False) + + df = session.read_csv(path, index_col="rowindex") + scalars_pandas_df_default_index = scalars_pandas_df_default_index.set_index( + "rowindex" + ).sort_index() + pd.testing.assert_index_equal( + df.columns, scalars_pandas_df_default_index.columns + ) + assert df.index.name == "rowindex" + + +def test_read_csv_local_default_engine_w_index_col_index( + session, scalars_pandas_df_default_index +): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_default_engine_w_index_col_index" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df_default_index.to_csv(path, index=False) + + index_col = scalars_pandas_df_default_index.columns.to_list().index("rowindex") + df = session.read_csv(path, index_col=index_col) + scalars_pandas_df_default_index = scalars_pandas_df_default_index.set_index( + "rowindex" + ).sort_index() + pd.testing.assert_index_equal( + df.columns, scalars_pandas_df_default_index.columns + ) + assert df.index.name == "rowindex" + + +@pytest.mark.parametrize( + "engine", + [ + pytest.param("bigquery", id="bq_engine"), + pytest.param(None, id="default_engine"), + ], +) +def test_read_csv_gcs_w_usecols(session, scalars_df_index, gcs_folder, engine): + path = gcs_folder + "test_read_csv_gcs_w_usecols" + path = path + "_default_engine.csv" if engine is None else path + "_bq_engine.csv" + scalars_df_index.to_csv(path) + + # df should only have 1 column which is bool_col. + df = session.read_csv(path, usecols=["bool_col"], engine=engine) + assert len(df.columns) == 1 + + +@pytest.mark.parametrize( + "engine", + [ + pytest.param("bigquery", id="bq_engine"), + pytest.param(None, id="default_engine"), + ], +) +def test_read_csv_local_w_usecols(session, scalars_pandas_df_index, engine): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_w_usecols.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + scalars_pandas_df_index.to_csv(path, index=False) + + # df should only have 1 column which is bool_col. + df = session.read_csv(path, usecols=["bool_col"], engine=engine) + assert len(df.columns) == 1 + + +@pytest.mark.parametrize( + "engine", + [ + pytest.param("bigquery", id="bq_engine"), + pytest.param(None, id="default_engine"), + ], +) +def test_read_csv_local_w_encoding(session, penguins_pandas_df_default_index, engine): + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_local_w_encoding.csv" + # Using the pandas to_csv method because the BQ one does not support local write. + penguins_pandas_df_default_index.to_csv( + path, index=False, encoding="ISO-8859-1" + ) + + # File can only be read using the same character encoding as when written. + df = session.read_csv(path, engine=engine, encoding="ISO-8859-1") + + # TODO(chelsealin): If we serialize the index, can more easily compare values. + pd.testing.assert_index_equal( + df.columns, penguins_pandas_df_default_index.columns + ) + + assert df.shape[0] == penguins_pandas_df_default_index.shape[0] + + +def test_read_parquet_gcs(session: bigframes.Session, scalars_dfs, gcs_folder): + scalars_df, _ = scalars_dfs + # Include wildcard so that multiple files can be written/read if > 1 GB. + # https://cloud.google.com/bigquery/docs/exporting-data#exporting_data_into_one_or_more_files + path = gcs_folder + "test_read_parquet_gcs*.parquet" + df_in: bigframes.dataframe.DataFrame = scalars_df.copy() + # GEOGRAPHY not supported in parquet export. + df_in = df_in.drop(columns="geography_col") + # Make sure we can also serialize the order. + df_write = df_in.reset_index(drop=False) + df_write.index.name = f"ordering_id_{random.randrange(1_000_000)}" + df_write.to_parquet(path, index=True) + + df_out = ( + session.read_parquet(path) + # Restore order. + .set_index(df_write.index.name).sort_index() + # Restore index. + .set_index(typing.cast(str, df_in.index.name)) + ) + + # DATETIME gets loaded as TIMESTAMP in parquet. See: + # https://cloud.google.com/bigquery/docs/exporting-data#parquet_export_details + df_out = df_out.assign( + datetime_col=df_out["datetime_col"].astype("timestamp[us][pyarrow]") + ) + + # Make sure we actually have at least some values before comparing. + assert df_out.size != 0 + pd_df_in = df_in.to_pandas() + pd_df_out = df_out.to_pandas() + pd.testing.assert_frame_equal(pd_df_in, pd_df_out) + + +def test_session_id(session): + assert session._session_id is not None + + # BQ client always runs query within the opened session. + query_job = session.bqclient.query("SELECT 1") + assert query_job.session_info.session_id == session._session_id + + # TODO(chelsealin): Verify the session id can be binded with a load job. + + +def test_session_dataset_exists_and_configured(session: bigframes.Session): + dataset = session.bqclient.get_dataset(session._session_dataset_id) + assert dataset.default_table_expiration_ms == 24 * 60 * 60 * 1000 + + +@pytest.mark.flaky(retries=2) +def test_to_close_session(): + session = bigframes.Session() + assert session._session_id is not None + session.close() + assert session._session_id is None + + # Session has expired and is no longer available. + with pytest.raises(google.api_core.exceptions.BadRequest): + query_job = session.bqclient.query("SELECT 1") + query_job.result() # blocks until finished diff --git a/tests/system/small/test_window.py b/tests/system/small/test_window.py new file mode 100644 index 00000000000..34b56fcd730 --- /dev/null +++ b/tests/system/small/test_window.py @@ -0,0 +1,66 @@ +# 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 pandas as pd +import pytest + + +@pytest.mark.parametrize( + ("windowing"), + [ + (lambda x: x.expanding()), + (lambda x: x.rolling(3, min_periods=3)), + (lambda x: x.groupby(x % 2).rolling(3, min_periods=3)), + (lambda x: x.groupby(x % 3).expanding(min_periods=2)), + ], + ids=[ + "expanding", + "rolling", + "rollinggroupby", + "expandinggroupby", + ], +) +@pytest.mark.parametrize( + ("agg_op"), + [ + (lambda x: x.sum()), + (lambda x: x.min()), + (lambda x: x.max()), + (lambda x: x.mean()), + (lambda x: x.count()), + (lambda x: x.std()), + (lambda x: x.var()), + ], + ids=[ + "sum", + "min", + "max", + "mean", + "count", + "std", + "var", + ], +) +def test_window_agg_ops(scalars_df_index, scalars_pandas_df_index, windowing, agg_op): + col_name = "int64_too" + bf_series = agg_op(windowing(scalars_df_index[col_name])).compute() + pd_series = agg_op(windowing(scalars_pandas_df_index[col_name])) + + # Pandas always converts to float64, even for min/max/count, which is not desired + pd_series = pd_series.astype(bf_series.dtype) + + pd.testing.assert_series_equal( + pd_series, + bf_series, + ) diff --git a/tests/system/utils.py b/tests/system/utils.py new file mode 100644 index 00000000000..e2daf3b8bf0 --- /dev/null +++ b/tests/system/utils.py @@ -0,0 +1,141 @@ +# 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 base64 +import decimal + +import geopandas as gpd # type: ignore +import numpy as np +import pandas as pd +import pyarrow as pa # type: ignore + + +def assert_pandas_df_equal_ignore_ordering(df0, df1, **kwargs): + # Sort by a column to get consistent results. + if df0.index.name != "rowindex": + df0 = df0.sort_values( + list(df0.columns.drop("geography_col", errors="ignore")) + ).reset_index(drop=True) + df1 = df1.sort_values( + list(df1.columns.drop("geography_col", errors="ignore")) + ).reset_index(drop=True) + else: + df0 = df0.sort_index() + df1 = df1.sort_index() + + pd.testing.assert_frame_equal(df0, df1, **kwargs) + + +def assert_series_equal_ignoring_order(left: pd.Series, right: pd.Series, **kwargs): + if left.index.name is None: + left = left.sort_values().reset_index(drop=True) + right = right.sort_values().reset_index(drop=True) + else: + left = left.sort_index() + right = right.sort_index() + + pd.testing.assert_series_equal(left, right, **kwargs) + + +def _standardize_index(idx): + return pd.Index(list(idx), name=idx.name) + + +def assert_pandas_index_equal_ignore_index_type(idx0, idx1): + idx0 = _standardize_index(idx0) + idx1 = _standardize_index(idx1) + + pd.testing.assert_index_equal(idx0, idx1) + + +def convert_pandas_dtypes(df: pd.DataFrame, bytes_col: bool): + """Convert pandas dataframe dtypes compatible with bigframes dataframe.""" + + # TODO(chelsealin): updates the function to accept dtypes as input rather than + # hard-code the column names here. + + # Convert basic types columns + df["bool_col"] = df["bool_col"].astype(pd.BooleanDtype()) + df["int64_col"] = df["int64_col"].astype(pd.Int64Dtype()) + df["int64_too"] = df["int64_too"].astype(pd.Int64Dtype()) + df["float64_col"] = df["float64_col"].astype(pd.Float64Dtype()) + df["string_col"] = df["string_col"].astype(pd.StringDtype(storage="pyarrow")) + + if "rowindex" in df.columns: + df["rowindex"] = df["rowindex"].astype(pd.Int64Dtype()) + if "rowindex_2" in df.columns: + df["rowindex_2"] = df["rowindex_2"].astype(pd.Int64Dtype()) + + # Convert time types columns. The `astype` works for Pandas 2.0 but hits an assert + # error at Pandas 1.5. Hence, we have to convert to arrow table and convert back + # to pandas dataframe. + if not isinstance(df["date_col"].dtype, pd.ArrowDtype): + df["date_col"] = pd.to_datetime(df["date_col"], format="%Y-%m-%d") + arrow_table = pa.Table.from_pandas( + pd.DataFrame(df, columns=["date_col"]), + schema=pa.schema([("date_col", pa.date32())]), + ) + df["date_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)["date_col"] + + if not isinstance(df["datetime_col"].dtype, pd.ArrowDtype): + df["datetime_col"] = pd.to_datetime( + df["datetime_col"], format="%Y-%m-%d %H:%M:%S" + ) + arrow_table = pa.Table.from_pandas( + pd.DataFrame(df, columns=["datetime_col"]), + schema=pa.schema([("datetime_col", pa.timestamp("us"))]), + ) + df["datetime_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[ + "datetime_col" + ] + + if not isinstance(df["time_col"].dtype, pd.ArrowDtype): + df["time_col"] = pd.to_datetime(df["time_col"], format="%H:%M:%S.%f") + arrow_table = pa.Table.from_pandas( + pd.DataFrame(df, columns=["time_col"]), + schema=pa.schema([("time_col", pa.time64("us"))]), + ) + df["time_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)["time_col"] + + if not isinstance(df["timestamp_col"].dtype, pd.ArrowDtype): + df["timestamp_col"] = pd.to_datetime( + df["timestamp_col"], format="%Y-%m-%d %H:%M:%S.%f%Z" + ) + arrow_table = pa.Table.from_pandas( + pd.DataFrame(df, columns=["timestamp_col"]), + schema=pa.schema([("timestamp_col", pa.timestamp("us", tz="UTC"))]), + ) + df["timestamp_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[ + "timestamp_col" + ] + + # Convert geography types columns. + if "geography_col" in df.columns: + df["geography_col"] = df["geography_col"].astype( + pd.StringDtype(storage="pyarrow") + ) + df["geography_col"] = gpd.GeoSeries.from_wkt( + df["geography_col"].replace({np.nan: None}) + ) + + # Convert bytes types column. + if bytes_col: + df["bytes_col"] = df["bytes_col"].apply( + lambda value: base64.b64decode(value) if not pd.isnull(value) else value + ) + + # Convert numeric types column. + df["numeric_col"] = df["numeric_col"].apply( + lambda value: decimal.Decimal(str(value)) if value else None # type: ignore + ) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/unit/__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/tests/unit/_config/__init__.py b/tests/unit/_config/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/unit/_config/__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/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py new file mode 100644 index 00000000000..43b5663bf75 --- /dev/null +++ b/tests/unit/_config/test_bigquery_options.py @@ -0,0 +1,73 @@ +# 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 re + +import pytest + +import bigframes._config.bigquery_options as bigquery_options + + +@pytest.mark.parametrize( + ["attribute", "original_value", "new_value"], + [ + # For credentials, the match is by reference. + ("credentials", object(), object()), + ("location", "us-east1", "us-central1"), + ("project", "my-project", "my-other-project"), + ("remote_udf_connection", "path/to/connection/1", "path/to/connection/2"), + ], +) +def test_setter_raises_if_session_started(attribute, original_value, new_value): + options = bigquery_options.BigQueryOptions() + setattr(options, attribute, original_value) + assert getattr(options, attribute) is original_value + assert getattr(options, attribute) is not new_value + + options._session_started = True + expected_message = re.escape( + bigquery_options.SESSION_STARTED_MESSAGE.format(attribute=attribute) + ) + with pytest.raises(ValueError, match=expected_message): + setattr(options, attribute, new_value) + + assert getattr(options, attribute) is original_value + assert getattr(options, attribute) is not new_value + + +@pytest.mark.parametrize( + [ + "attribute", + ], + [ + (attribute,) + for attribute in [ + "credentials", + "location", + "project", + "remote_udf_connection", + ] + ], +) +def test_setter_if_session_started_but_setting_the_same_value(attribute): + options = bigquery_options.BigQueryOptions() + original_object = object() + setattr(options, attribute, original_object) + assert getattr(options, attribute) is original_object + + # This should work fine since we're setting the same value as before. + options._session_started = True + setattr(options, attribute, original_object) + + assert getattr(options, attribute) is original_object diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000000..dcf2d918a5e --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,223 @@ +# 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 math +from typing import Callable, Optional, Tuple, Union +from unittest import mock + +import google.api_core.exceptions +import google.auth +import google.cloud.bigquery as bigquery +import google.cloud.bigquery.table +import google.oauth2.credentials # type: ignore +import ibis.expr.types as ibis_types +import pandas +import pytest + +import bigframes +import bigframes.core +import bigframes.dataframe + +SCALARS_TABLE_ID = "project.dataset.scalars_table" + + +@pytest.fixture +def scalars_pandas_df_default_index() -> pandas.DataFrame: + # Note: as of 2023-02-07, using nullable dtypes with the ibis pandas + # backend requires running ibis at HEAD. See: + # https://github.com/ibis-project/ibis/pull/5345 + return pandas.DataFrame( + { + "rowindex": pandas.Series( + [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + ], + dtype="Int64", + ), + "bool_col": pandas.Series( + [ + True, + None, + False, + True, + None, + False, + True, + None, + False, + True, + ], + dtype="boolean", + ), + "int64_col": pandas.Series( + [ + 1, + 2, + 3, + None, + 0, + -1, + -2, + 2**63 - 1, + -(2**63), + None, + ], + dtype="Int64", + ), + "float64_col": pandas.Series( + [ + None, + 1, + math.pi, + math.e * 1e10, + 0, + float("nan"), + float("inf"), + float("-inf"), + -2.23e-308, + 1.8e308, + ], + dtype="Float64", + ), + "string_col": pandas.Series( + [ + "abc", + "XYZ", + "aBcDeFgHiJkLmNoPqRsTuVwXyZ", + "1_2-3+4=5~6*7/8&9%10#11@12$" "", + None, + "こんにちは", + "你好", + "வணக்கம்", + "שלום", + ], + dtype="string[pyarrow]", + ), + } + ) + + +# We parameterize the fixtures at this point with the real pandas +# dataframes and deferred bigframes dataframes as we have the following +# chain of dependencies: +# -> index/default_index parameterization +# -> pandas dataframe +# -> bqclient mock +# -> session +# -> bigframes dataframe +@pytest.fixture +def scalars_testdata_setup( + scalars_pandas_df_default_index, +) -> Tuple[ + pandas.DataFrame, Callable[[bigframes.Session], bigframes.dataframe.DataFrame] +]: + return ( + scalars_pandas_df_default_index.set_index("rowindex"), + lambda session: session.read_gbq(SCALARS_TABLE_ID, index_col=["rowindex"]), + ) + + +@pytest.fixture(autouse=True) +def mock_bigquery_client(monkeypatch, scalars_testdata_setup) -> bigquery.Client: + scalars_pandas_df, _ = scalars_testdata_setup + mock_client = mock.create_autospec(bigquery.Client) + # Constructor returns the mock itself, so this mock can be treated as the + # constructor or the instance. + mock_client.return_value = mock_client + mock_client.project = "default-project" + most_recent_table = None + + def mock_bigquery_client_get_table( + table_ref: Union[google.cloud.bigquery.table.TableReference, str] + ): + global most_recent_table + + if isinstance(table_ref, google.cloud.bigquery.table.TableReference): + table_name = table_ref.__str__() + else: + table_name = table_ref + + schema = [ + {"mode": "NULLABLE", "name": "rowindex", "type": "INTEGER"}, + { + "mode": "NULLABLE", + "name": "bigframes_ordering_id", + "type": "INTEGER", + }, + ] + + if table_name == SCALARS_TABLE_ID: + schema += [ + {"mode": "NULLABLE", "name": "bool_col", "type": "BOOL"}, + {"mode": "NULLABLE", "name": "int64_col", "type": "INTEGER"}, + {"mode": "NULLABLE", "name": "float64_col", "type": "FLOAT"}, + {"mode": "NULLABLE", "name": "string_col", "type": "STRING"}, + ] + else: + raise google.api_core.exceptions.NotFound("Not Found Table") + + most_recent_table = bigquery.Table(table_name, schema) # type: ignore + return most_recent_table # type: ignore + + def mock_query( + sql: str, + job_config: Optional[bigquery.QueryJobConfig] = None, + location: str = "US", + ) -> bigquery.QueryJob: + global most_recent_table + + def mock_result(max_results=None): + mock_rows = mock.create_autospec(google.cloud.bigquery.table.RowIterator) + mock_rows.total_rows = len(scalars_pandas_df.index) + mock_rows.schema = [ + bigquery.SchemaField(name=name, field_type="INT64") + for name in scalars_pandas_df.columns + ] + # Use scalars_pandas_df instead of ibis_expr.execute() to preserve dtypes. + mock_rows.to_dataframe.return_value = scalars_pandas_df.head(n=max_results) + return mock_rows + + mock_job = mock.create_autospec(bigquery.QueryJob) + mock_job.result = mock_result + return mock_job + + mock_client.get_table = mock_bigquery_client_get_table + mock_client.query.side_effect = mock_query + monkeypatch.setattr(bigquery, "Client", mock_client) + mock_client.reset_mock() + return mock_client + + +@pytest.fixture +def session() -> bigframes.Session: + return bigframes.Session( + context=bigframes.BigQueryOptions( + credentials=mock.create_autospec(google.oauth2.credentials.Credentials), + project="unit-test-project", + ) + ) + + +@pytest.fixture +def scalars_ibis_table(session) -> ibis_types.Table: + return session.ibis_client.table(SCALARS_TABLE_ID) diff --git a/tests/unit/ml/__init__.py b/tests/unit/ml/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/unit/ml/__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/tests/unit/ml/test_api_primitives.py b/tests/unit/ml/test_api_primitives.py new file mode 100644 index 00000000000..da77a180a8d --- /dev/null +++ b/tests/unit/ml/test_api_primitives.py @@ -0,0 +1,53 @@ +# 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 sklearn.decomposition as sklearn_decomposition # type: ignore +import sklearn.linear_model as sklearn_linear_model # type: ignore + +import bigframes.ml.decomposition +import bigframes.ml.linear_model + + +def test_base_estimator_repr(): + estimator = bigframes.ml.linear_model.LinearRegression() + assert estimator.__repr__() == "LinearRegression()" + + estimator = bigframes.ml.linear_model.LinearRegression(fit_intercept=False) + assert estimator.__repr__() == "LinearRegression(fit_intercept=False)" + + estimator = bigframes.ml.linear_model.LinearRegression(fit_intercept=True) + assert estimator.__repr__() == "LinearRegression()" + + estimator = bigframes.ml.decomposition.PCA(n_components=7) + assert estimator.__repr__() == "PCA(n_components=7)" + + +@pytest.mark.skipif(sklearn_linear_model is None, reason="requires sklearn") +def test_base_estimator_repr_matches_sklearn(): + estimator = bigframes.ml.linear_model.LinearRegression() + sklearn_estimator = sklearn_linear_model.LinearRegression() + assert estimator.__repr__() == sklearn_estimator.__repr__() + + estimator = bigframes.ml.linear_model.LinearRegression(fit_intercept=False) + sklearn_estimator = sklearn_linear_model.LinearRegression(fit_intercept=False) + assert estimator.__repr__() == sklearn_estimator.__repr__() + + estimator = bigframes.ml.linear_model.LinearRegression(fit_intercept=True) + sklearn_estimator = sklearn_linear_model.LinearRegression(fit_intercept=True) + assert estimator.__repr__() == sklearn_estimator.__repr__() + + estimator = bigframes.ml.decomposition.PCA(n_components=7) + sklearn_estimator = sklearn_decomposition.PCA(n_components=7) + assert estimator.__repr__() == sklearn_estimator.__repr__() diff --git a/tests/unit/ml/test_compose.py b/tests/unit/ml/test_compose.py new file mode 100644 index 00000000000..772a148c95f --- /dev/null +++ b/tests/unit/ml/test_compose.py @@ -0,0 +1,95 @@ +# 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 sklearn.compose as sklearn_compose # type: ignore +import sklearn.preprocessing as sklearn_preprocessing # type: ignore + +import bigframes.ml.compose +import bigframes.ml.preprocessing + + +def test_columntransformer_init_expectedtransforms(): + onehot_transformer = bigframes.ml.preprocessing.OneHotEncoder() + scaler_transformer = bigframes.ml.preprocessing.StandardScaler() + column_transformer = bigframes.ml.compose.ColumnTransformer( + [ + ("onehot", onehot_transformer, "species"), + ("scale", scaler_transformer, ["culmen_length_mm", "flipper_length_mm"]), + ] + ) + + assert column_transformer.transformers_ == [ + ("onehot", onehot_transformer, "species"), + ("scale", scaler_transformer, "culmen_length_mm"), + ("scale", scaler_transformer, "flipper_length_mm"), + ] + + +def test_columntransformer_repr(): + column_transformer = bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ) + + assert ( + column_transformer.__repr__() + == """ColumnTransformer(transformers=[('onehot', OneHotEncoder(), 'species'), + ('scale', StandardScaler(), + ['culmen_length_mm', 'flipper_length_mm'])])""" + ) + + +@pytest.mark.skipif(sklearn_compose is None, reason="requires sklearn") +def test_columntransformer_repr_matches_sklearn(): + bf_column_transformer = bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ) + sk_column_transformer = sklearn_compose.ColumnTransformer( + [ + ( + "onehot", + sklearn_preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + sklearn_preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ) + + assert bf_column_transformer.__repr__() == sk_column_transformer.__repr__() diff --git a/tests/unit/ml/test_pipeline.py b/tests/unit/ml/test_pipeline.py new file mode 100644 index 00000000000..27706a1a07f --- /dev/null +++ b/tests/unit/ml/test_pipeline.py @@ -0,0 +1,110 @@ +# 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 sklearn.compose as sklearn_compose # type: ignore +import sklearn.linear_model as sklearn_linear_model # type: ignore +import sklearn.pipeline as sklearn_pipeline # type: ignore +import sklearn.preprocessing as sklearn_preprocessing # type: ignore + +import bigframes.ml.compose +import bigframes.ml.linear_model +import bigframes.ml.pipeline +import bigframes.ml.preprocessing + + +def test_pipeline_repr(): + pipeline = bigframes.ml.pipeline.Pipeline( + [ + ( + "preproc", + bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ), + ), + ("linreg", bigframes.ml.linear_model.LinearRegression()), + ] + ) + + assert ( + pipeline.__repr__() + == """Pipeline(steps=[('preproc', + ColumnTransformer(transformers=[('onehot', OneHotEncoder(), + 'species'), + ('scale', StandardScaler(), + ['culmen_length_mm', + 'flipper_length_mm'])])), + ('linreg', LinearRegression())])""" + ) + + +@pytest.mark.skipif(sklearn_pipeline is None, reason="requires sklearn") +def test_pipeline_repr_matches_sklearn(): + bf_pipeline = bigframes.ml.pipeline.Pipeline( + [ + ( + "preproc", + bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ), + ), + ("linreg", bigframes.ml.linear_model.LinearRegression()), + ] + ) + sk_pipeline = sklearn_pipeline.Pipeline( + [ + ( + "preproc", + sklearn_compose.ColumnTransformer( + [ + ( + "onehot", + sklearn_preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + sklearn_preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ), + ), + ("linreg", sklearn_linear_model.LinearRegression()), + ] + ) + + assert bf_pipeline.__repr__() == sk_pipeline.__repr__() diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py new file mode 100644 index 00000000000..92cb52a6815 --- /dev/null +++ b/tests/unit/ml/test_sql.py @@ -0,0 +1,145 @@ +# 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 bigframes.ml.sql as ml_sql + + +def test_options_produces_correct_sql(): + sql = ml_sql.options(model_type="lin_reg", input_label_cols=["col_a"], l1_reg=0.6) + assert ( + sql + == """OPTIONS( + model_type="lin_reg", + input_label_cols=["col_a"], + l1_reg=0.6)""" + ) + + +def test_transform_produces_correct_sql(): + sql = ml_sql.transform( + "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", + "ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b", + ) + assert ( + sql + == """TRANSFORM( + ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a, + ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b)""" + ) + + +def test_standard_scaler_produces_correct_sql(): + sql = ml_sql.ml_standard_scaler("col_a", "scaled_col_a") + assert sql == "ML.STANDARD_SCALER(col_a) OVER() AS scaled_col_a" + + +def test_one_hot_encoder_produces_correct_sql(): + sql = ml_sql.ml_one_hot_encoder("col_a", "encoded_col_a") + assert sql == "ML.ONE_HOT_ENCODER(col_a) OVER() AS encoded_col_a" + + +def test_create_model_produces_correct_sql(): + sql = ml_sql.create_model( + model_name="my_dataset.my_model", + source_sql="my_source_sql", + options_sql="my_options_sql", + ) + assert ( + sql + == """CREATE MODEL `my_dataset.my_model` +my_options_sql +AS my_source_sql""" + ) + + +def test_create_model_transform_produces_correct_sql(): + sql = ml_sql.create_model( + model_name="my_dataset.my_model", + source_sql="my_source_sql", + options_sql="my_options_sql", + transform_sql="my_transform_sql", + ) + assert ( + sql + == """CREATE MODEL `my_dataset.my_model` +my_transform_sql +my_options_sql +AS my_source_sql""" + ) + + +def test_create_remote_model_produces_correct_sql(): + sql = ml_sql.create_remote_model( + model_name="my_dataset.my_model", + connection_name="my_project.us.my_connection", + options_sql="my_options_sql", + ) + assert ( + sql + == """CREATE MODEL `my_dataset.my_model` +REMOTE WITH CONNECTION `my_project.us.my_connection` +my_options_sql""" + ) + + +def test_alter_model_correct_sql(): + sql = ml_sql.alter_model( + model_name="my_dataset.my_model", + options_sql="my_options_sql", + ) + assert ( + sql + == """ALTER MODEL `my_dataset.my_model` +SET my_options_sql""" + ) + + +def test_ml_predict_produces_correct_sql(): + sql = ml_sql.ml_predict( + model_name="my_dataset.my_model", source_sql="SELECT * FROM my_table" + ) + assert ( + sql + == """SELECT * FROM ML.PREDICT(MODEL `my_dataset.my_model`, + (SELECT * FROM my_table))""" + ) + + +def test_ml_evaluate_produces_correct_sql(): + sql = ml_sql.ml_evaluate( + model_name="my_dataset.my_model", source_sql="SELECT * FROM my_table" + ) + assert ( + sql + == """SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`, + (SELECT * FROM my_table))""" + ) + + +def test_ml_evaluate_no_source_produces_correct_sql(): + sql = ml_sql.ml_evaluate(model_name="my_dataset.my_model") + assert sql == """SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`)""" + + +def test_ml_generate_text_produces_correct_sql(): + sql = ml_sql.ml_generate_text( + model_name="my_dataset.my_model", + source_sql="SELECT * FROM my_table", + struct_options="STRUCT(value AS item)", + ) + assert ( + sql + == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_dataset.my_model`, + (SELECT * FROM my_table), STRUCT(value AS item))""" + ) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py new file mode 100644 index 00000000000..25286189637 --- /dev/null +++ b/tests/unit/test_core.py @@ -0,0 +1,99 @@ +# 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. + +from unittest import mock + +import ibis +from ibis.expr.types import Column, Table + +from bigframes import core + + +def test_constructor_from_ibis_table_adds_all_columns( + session, scalars_ibis_table: Table +): + actual = core.ArrayValue(session=session, table=scalars_ibis_table) + assert actual._table is scalars_ibis_table + assert len(actual._columns) == len(scalars_ibis_table.columns) + + +def test_builder_doesnt_change_original(session): + mock_table = mock.create_autospec(Table) + mock_column = mock.create_autospec(Column) + original = core.ArrayValue(session=session, table=mock_table, columns=[mock_column]) + assert original._table is mock_table + assert len(original._columns) == 1 + assert original._columns[0] is mock_column + + # Create a new expression from a builder. + builder = original.builder() + new_table = mock.create_autospec(Table) + assert new_table is not mock_table + builder.table = new_table + new_column = mock.create_autospec(Column) + assert new_column is not mock_column + builder.columns.append(new_column) + actual = builder.build() + + # Expected values are present. + assert actual._table is new_table + assert len(actual._columns) == 2 + assert actual._columns[0] is mock_column + assert actual._columns[1] is new_column + # Don't modify the original. + assert original._table is mock_table + assert len(original._columns) == 1 + assert original._columns[0] is mock_column + + +def test_projection_doesnt_change_original(session): + mock_table = mock.create_autospec(Table) + mock_column = mock.create_autospec(Column) + original = core.ArrayValue(session=session, table=mock_table, columns=[mock_column]) + assert original._table is mock_table + assert len(original._columns) == 1 + assert original._columns[0] is mock_column + + # Create a new expression from a projection. + new_column_1 = mock.create_autospec(Column) + new_column_2 = mock.create_autospec(Column) + assert new_column_1 is not mock_column + assert new_column_2 is not mock_column + actual = original.projection([new_column_1, mock_column, new_column_2]) + + # Expected values are present. + assert actual._table is mock_table + assert len(actual._columns) == 3 + assert actual._columns[0] is new_column_1 + assert actual._columns[1] is mock_column + assert actual._columns[2] is new_column_2 + # Don't modify the original. + assert original._table is mock_table + assert len(original._columns) == 1 + assert original._columns[0] is mock_column + + +def test_to_ibis_expr_with_projection(session, scalars_ibis_table: Table): + expr = core.ArrayValue(session=session, table=scalars_ibis_table).projection( + [ + scalars_ibis_table["int64_col"], + ibis.literal(123456789).name("literals"), + scalars_ibis_table["string_col"], + ] + ) + actual = expr.to_ibis_expr() + assert len(actual.columns) == 3 + assert actual.columns[0] == "int64_col" + assert actual.columns[1] == "literals" + assert actual.columns[2] == "string_col" diff --git a/tests/unit/test_dtypes.py b/tests/unit/test_dtypes.py new file mode 100644 index 00000000000..ec49315f0fe --- /dev/null +++ b/tests/unit/test_dtypes.py @@ -0,0 +1,184 @@ +# 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 geopandas as gpd # type: ignore +import ibis +import ibis.expr.datatypes as ibis_dtypes +import numpy as np +import pandas as pd +import pyarrow as pa # type: ignore +import pytest + +import bigframes.dtypes + + +@pytest.mark.parametrize( + ["ibis_dtype", "bigframes_dtype"], + [ + # TODO(bmil): Add ARRAY, INTERVAL, STRUCT to cover all the standard + # BigQuery data types as they appear in Ibis: + # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types + (ibis_dtypes.Decimal(precision=76, scale=38, nullable=True), np.dtype("O")), + (ibis_dtypes.boolean, pd.BooleanDtype()), + (ibis_dtypes.binary, np.dtype("O")), + (ibis_dtypes.date, pd.ArrowDtype(pa.date32())), + (ibis_dtypes.Timestamp(), pd.ArrowDtype(pa.timestamp("us"))), + (ibis_dtypes.float64, pd.Float64Dtype()), + ( + ibis_dtypes.GeoSpatial(geotype="geography", srid=4326, nullable=True), + gpd.array.GeometryDtype(), + ), + (ibis_dtypes.int64, pd.Int64Dtype()), + (ibis_dtypes.json, np.dtype("O")), + (ibis_dtypes.Decimal(precision=38, scale=9, nullable=True), np.dtype("O")), + (ibis_dtypes.string, pd.StringDtype(storage="pyarrow")), + (ibis_dtypes.time, pd.ArrowDtype(pa.time64("us"))), + ( + ibis_dtypes.Timestamp(timezone="UTC"), + pd.ArrowDtype(pa.timestamp("us", tz="UTC")), # type: ignore + ), + ], + ids=[ + "bignumeric", + "bool", + "bytes", + "date", + "datetime", + "float", + "geography", + "int64", + "json", + "numeric", + "string", + "time", + "timestamp", + ], +) +def test_ibis_dtype_converts(ibis_dtype, bigframes_dtype): + """Test all the Ibis data types needed to read BigQuery tables""" + result = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_dtype) + assert result == bigframes_dtype + + +def test_ibis_timestamp_pst_raises_unexpected_datatype(): + """BigQuery timestamp only supports UTC time""" + with pytest.raises(ValueError, match="Unexpected Ibis data type"): + bigframes.dtypes.ibis_dtype_to_bigframes_dtype( + ibis_dtypes.Timestamp(timezone="PST") + ) + + +def test_ibis_float32_raises_unexpected_datatype(): + """Other Ibis types not read from BigQuery are not expected""" + with pytest.raises(ValueError, match="Unexpected Ibis data type"): + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_dtypes.float32) + + +@pytest.mark.parametrize( + ["bigframes_dtype", "ibis_dtype"], + [ + # This test covers all dtypes that BigQuery DataFrames can exactly map to Ibis + (pd.BooleanDtype(), ibis_dtypes.boolean), + (pd.ArrowDtype(pa.date32()), ibis_dtypes.date), + (pd.ArrowDtype(pa.timestamp("us")), ibis_dtypes.Timestamp()), + (pd.Float64Dtype(), ibis_dtypes.float64), + (pd.Int64Dtype(), ibis_dtypes.int64), + (pd.StringDtype(storage="pyarrow"), ibis_dtypes.string), + (pd.ArrowDtype(pa.time64("us")), ibis_dtypes.time), + ( + pd.ArrowDtype(pa.timestamp("us", tz="UTC")), # type: ignore + ibis_dtypes.Timestamp(timezone="UTC"), + ), + ], + ids=[ + "boolean", + "date", + "datetime", + "float", + "int", + "string", + "time", + "timestamp", + ], +) +def test_bigframes_dtype_converts(ibis_dtype, bigframes_dtype): + """Test all the Ibis data types needed to read BigQuery tables""" + result = bigframes.dtypes.bigframes_dtype_to_ibis_dtype(bigframes_dtype) + assert result == ibis_dtype + + +@pytest.mark.parametrize( + ["bigframes_dtype_str", "ibis_dtype"], + [ + # This test covers all dtypes that BigQuery DataFrames can exactly map to Ibis + ("boolean", ibis_dtypes.boolean), + ("date32[day][pyarrow]", ibis_dtypes.date), + ("timestamp[us][pyarrow]", ibis_dtypes.Timestamp()), + ("Float64", ibis_dtypes.float64), + ("Int64", ibis_dtypes.int64), + ("string[pyarrow]", ibis_dtypes.string), + ("time64[us][pyarrow]", ibis_dtypes.time), + ( + "timestamp[us, tz=UTC][pyarrow]", + ibis_dtypes.Timestamp(timezone="UTC"), + ), + # Special case - "string" is acceptable for "string[pyarrow]" + ("string", ibis_dtypes.string), + ], +) +def test_bigframes_string_dtype_converts(ibis_dtype, bigframes_dtype_str): + """Test all the Ibis data types needed to read BigQuery tables""" + result = bigframes.dtypes.bigframes_dtype_to_ibis_dtype(bigframes_dtype_str) + assert result == ibis_dtype + + +def test_unsupported_dtype_raises_unexpected_datatype(): + """Incompatible dtypes should fail when passed into BigQuery DataFrames""" + with pytest.raises(ValueError, match="Unexpected data type"): + bigframes.dtypes.bigframes_dtype_to_ibis_dtype(np.float32) + + +def test_unsupported_dtype_str_raises_unexpected_datatype(): + """Incompatible dtypes should fail when passed into BigQuery DataFrames""" + with pytest.raises(ValueError, match="Unexpected data type"): + bigframes.dtypes.bigframes_dtype_to_ibis_dtype("int64") + + +@pytest.mark.parametrize( + ["literal", "ibis_scalar"], + [ + (True, ibis.literal(True, ibis_dtypes.boolean)), + (5, ibis.literal(5, ibis_dtypes.int64)), + (-33.2, ibis.literal(-33.2, ibis_dtypes.float64)), + ], +) +def test_literal_to_ibis_scalar_converts(literal, ibis_scalar): + assert bigframes.dtypes.literal_to_ibis_scalar(literal).equals(ibis_scalar) + + +def test_literal_to_ibis_scalar_throws_on_incompatible_literal(): + with pytest.raises( + ValueError, + ): + bigframes.dtypes.literal_to_ibis_scalar({"mykey": "myval"}) + + +def test_remote_function_io_types_are_supported_bigframes_types(): + from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type + + from bigframes.remote_function import _supported_io_types as rf_supported_io_types + + for python_type in rf_supported_io_types: + ibis_type = python_type_to_bigquery_type(python_type) + assert ibis_type in bigframes.dtypes.IBIS_TO_BIGFRAMES diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py new file mode 100644 index 00000000000..3e8e16ee6ab --- /dev/null +++ b/tests/unit/test_pandas.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. + +import inspect +import re +import sys + +import pytest + +import bigframes.pandas +import bigframes.session + +leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE) + + +def all_session_methods(): + session_attributes = set( + attribute + for attribute in dir(bigframes.session.Session) + if not attribute.startswith("_") + ) + session_attributes.remove("close") + + for attribute in sorted(session_attributes): + session_method = getattr(bigframes.session.Session, attribute) + if not callable(session_method): + continue # pragma: no cover + yield attribute + + +@pytest.mark.parametrize( + [ + "method_name", + ], + [(method_name,) for method_name in all_session_methods()], +) +def test_method_matches_session(method_name: str): + if sys.version_info <= (3, 10): + pytest.skip( + "Need Python 3.10 to reconcile deferred annotations." + ) # pragma: no cover + + session_method = getattr(bigframes.session.Session, method_name) + session_doc = inspect.getdoc(session_method) + assert session_doc is not None, "docstrings are required" + + pandas_method = getattr(bigframes.pandas, method_name) + pandas_doc = inspect.getdoc(pandas_method) + assert pandas_doc is not None, "docstrings are required" + assert re.sub(leading_whitespace, "", pandas_doc) == re.sub( + leading_whitespace, "", session_doc + ) + + # Add `eval_str = True` so that deferred annotations are turned into their + # corresponding type objects. Need Python 3.10 for eval_str parameter. + session_signature = inspect.signature(session_method, eval_str=True) + pandas_signature = inspect.signature(pandas_method, eval_str=True) + assert [ + # Kind includes position, which will be an offset. + parameter.replace(kind=inspect.Parameter.POSITIONAL_ONLY) + for parameter in pandas_signature.parameters.values() + ] == [ + # Kind includes position, which will be an offset. + parameter.replace(kind=inspect.Parameter.POSITIONAL_ONLY) + for parameter in session_signature.parameters.values() + # Don't include the first parameter, which is `self: Session` + ][ + 1: + ] + assert pandas_signature.return_annotation == session_signature.return_annotation diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py new file mode 100644 index 00000000000..ab573c4c11f --- /dev/null +++ b/tests/unit/test_session.py @@ -0,0 +1,48 @@ +# 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 os +from unittest import mock + +import google.api_core.exceptions +import pytest + +import bigframes + + +@pytest.mark.parametrize("missing_parts_table_id", [(""), ("table")]) +def test_read_gbq_missing_parts(session, missing_parts_table_id): + with pytest.raises(ValueError): + session.read_gbq(missing_parts_table_id) + + +@pytest.mark.parametrize( + "not_found_table_id", + [("unknown.dataset.table"), ("project.unknown.table"), ("project.dataset.unknown")], +) +def test_read_gdb_not_found_tables(session, not_found_table_id): + with pytest.raises(google.api_core.exceptions.NotFound): + session.read_gbq(not_found_table_id) + + +@mock.patch.dict(os.environ, {}, clear=True) +def test_session_init_fails_with_no_project(): + with pytest.raises( + ValueError, match="Project must be set to initialize BigQuery client." + ): + bigframes.Session( + bigframes.BigQueryOptions( + credentials=mock.Mock(spec=google.auth.credentials.Credentials) + ) + ) diff --git a/third_party/__init__.py b/third_party/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/__init__.py b/third_party/bigframes_vendored/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/cpython/LICENSE b/third_party/bigframes_vendored/cpython/LICENSE new file mode 100644 index 00000000000..73935ad7468 --- /dev/null +++ b/third_party/bigframes_vendored/cpython/LICENSE @@ -0,0 +1,10 @@ +PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 + +1. This LICENSE AGREEMENT is between the Python Software Foundation ("PSF"), and the Individual or Organization ("Licensee") accessing and otherwise using this software ("Python") in source or binary form and its associated documentation. +2. Subject to the terms and conditions of this License Agreement, PSF hereby grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce, analyze, test, perform and/or display publicly, prepare derivative works, distribute, and otherwise use Python alone or in any derivative version, provided, however, that PSF's License Agreement and PSF's notice of copyright , i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006 Python Software Foundation All Rights Reserved" are retained in Python alone or in any derivative version prepared by Licensee. +3. In the event Licensee prepares a derivative work that is based on or incorporates Python or any part thereof, and wants to make the derivative work available to others as provided herein, then Licensee hereby agrees to include in any such work a brief summary of the changes made to Python. +4. PSF is making Python available to Licensee on an "AS IS" basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. +5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. +6. This License Agreement will automatically terminate upon a material breach of its terms and conditions. +7. Nothing in this License Agreement shall be deemed to create any relationship of agency, partnership, or joint venture between PSF and Licensee. This License Agreement does not grant permission to use PSF trademarks or trade name in a trademark sense to endorse or promote products or services of Licensee, or any third party. +8. By copying, installing or otherwise using Python, Licensee agrees to be bound by the terms and conditions of this License Agreement. diff --git a/third_party/bigframes_vendored/cpython/__init__.py b/third_party/bigframes_vendored/cpython/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/cpython/_pprint.py b/third_party/bigframes_vendored/cpython/_pprint.py new file mode 100644 index 00000000000..617c14df0d9 --- /dev/null +++ b/third_party/bigframes_vendored/cpython/_pprint.py @@ -0,0 +1,470 @@ +"""This module contains the _EstimatorPrettyPrinter class used in +BaseEstimator.__repr__ for pretty-printing estimators""" + +# Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, +# 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018 Python Software Foundation; +# All Rights Reserved + +# Authors: Fred L. Drake, Jr. (built-in CPython pprint module) +# Nicolas Hug (scikit-learn specific changes) +# Brenton Milne (bigframes specific changes) +# This version is based on the version at sklearn/utils/_pprint.py at 65d42c9 + +# License: PSF License version 2 (see below) + +# PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +# -------------------------------------------- + +# 1. This LICENSE AGREEMENT is between the Python Software Foundation ("PSF"), +# and the Individual or Organization ("Licensee") accessing and otherwise +# using this software ("Python") in source or binary form and its associated +# documentation. + +# 2. Subject to the terms and conditions of this License Agreement, PSF hereby +# grants Licensee a nonexclusive, royalty-free, world-wide license to +# reproduce, analyze, test, perform and/or display publicly, prepare +# derivative works, distribute, and otherwise use Python alone or in any +# derivative version, provided, however, that PSF's License Agreement and +# PSF's notice of copyright, i.e., "Copyright (c) 2001, 2002, 2003, 2004, +# 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, +# 2017, 2018 Python Software Foundation; All Rights Reserved" are retained in +# Python alone or in any derivative version prepared by Licensee. + +# 3. In the event Licensee prepares a derivative work that is based on or +# incorporates Python or any part thereof, and wants to make the derivative +# work available to others as provided herein, then Licensee hereby agrees to +# include in any such work a brief summary of the changes made to Python. + +# 4. PSF is making Python available to Licensee on an "AS IS" basis. PSF MAKES +# NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR IMPLIED. BY WAY OF EXAMPLE, BUT +# NOT LIMITATION, PSF MAKES NO AND DISCLAIMS ANY REPRESENTATION OR WARRANTY OF +# MERCHANTABILITY OR FITNESS FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF +# PYTHON WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +# 5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON FOR ANY +# INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS A RESULT OF +# MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, OR ANY DERIVATIVE +# THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +# 6. This License Agreement will automatically terminate upon a material +# breach of its terms and conditions. + +# 7. Nothing in this License Agreement shall be deemed to create any +# relationship of agency, partnership, or joint venture between PSF and +# Licensee. This License Agreement does not grant permission to use PSF +# trademarks or trade name in a trademark sense to endorse or promote products +# or services of Licensee, or any third party. + +# 8. By copying, installing or otherwise using Python, Licensee agrees to be +# bound by the terms and conditions of this License Agreement. + + +# Brief summary of changes to original code: +# - "compact" parameter is supported for dicts, not just lists or tuples +# - estimators have a custom handler, they're not just treated as objects +# - long sequences (lists, tuples, dict items) with more than N elements are +# shortened using ellipsis (', ...') at the end. +# +# Additional changes to adapt this to BigQuery DataFrame: +# - sklearn BaseEstimator replaced with BigQuery DataFrames BaseEstimator +# - removed global get_config, set _changed_only=True +# - replace is_scalar_nan with isinstance(x, numbers.Real) and math.isnan + +from collections import OrderedDict +import inspect +import math +import numbers +import pprint + +from bigframes.ml.base import BaseEstimator + + +class KeyValTuple(tuple): + """Dummy class for correctly rendering key-value tuples from dicts.""" + + def __repr__(self): + # needed for _dispatch[tuple.__repr__] not to be overridden + return super().__repr__() + + +class KeyValTupleParam(KeyValTuple): + """Dummy class for correctly rendering key-value tuples from parameters.""" + + pass + + +def _changed_params(estimator): + """Return dict (param_name: value) of parameters that were given to + estimator with non-default values.""" + + params = estimator.get_params(deep=False) + init_func = getattr(estimator.__init__, "deprecated_original", estimator.__init__) + init_params = inspect.signature(init_func).parameters + init_params = {name: param.default for name, param in init_params.items()} + + def has_changed(k, v): + if k not in init_params: # happens if k is part of a **kwargs + return True + if init_params[k] == inspect._empty: # k has no default value + return True + # try to avoid calling repr on nested estimators + if isinstance(v, BaseEstimator) and v.__class__ != init_params[k].__class__: + return True + # Use repr as a last resort. It may be expensive. + def is_scalar_nan(x): + return isinstance(x, numbers.Real) and math.isnan(x) + + if repr(v) != repr(init_params[k]) and not ( + is_scalar_nan(init_params[k]) and is_scalar_nan(v) + ): + return True + return False + + return {k: v for k, v in params.items() if has_changed(k, v)} + + +class _EstimatorPrettyPrinter(pprint.PrettyPrinter): + """Pretty Printer class for estimator objects. + + This extends the pprint.PrettyPrinter class, because: + - we need estimators to be printed with their parameters, e.g. + Estimator(param1=value1, ...) which is not supported by default. + - the 'compact' parameter of PrettyPrinter is ignored for dicts, which + may lead to very long representations that we want to avoid. + + Quick overview of pprint.PrettyPrinter (see also + https://stackoverflow.com/questions/49565047/pprint-with-hex-numbers): + + - the entry point is the _format() method which calls format() (overridden + here) + - format() directly calls _safe_repr() for a first try at rendering the + object + - _safe_repr formats the whole object recursively, only calling itself, + not caring about line length or anything + - back to _format(), if the output string is too long, _format() then calls + the appropriate _pprint_TYPE() method (e.g. _pprint_list()) depending on + the type of the object. This where the line length and the compact + parameters are taken into account. + - those _pprint_TYPE() methods will internally use the format() method for + rendering the nested objects of an object (e.g. the elements of a list) + + In the end, everything has to be implemented twice: in _safe_repr and in + the custom _pprint_TYPE methods. Unfortunately PrettyPrinter is really not + straightforward to extend (especially when we want a compact output), so + the code is a bit convoluted. + + This class overrides: + - format() to support the changed_only parameter + - _safe_repr to support printing of estimators (for when they fit on a + single line) + - _format_dict_items so that dict are correctly 'compacted' + - _format_items so that ellipsis is used on long lists and tuples + + When estimators cannot be printed on a single line, the builtin _format() + will call _pprint_estimator() because it was registered to do so (see + _dispatch[BaseEstimator.__repr__] = _pprint_estimator). + + both _format_dict_items() and _pprint_estimator() use the + _format_params_or_dict_items() method that will format parameters and + key-value pairs respecting the compact parameter. This method needs another + subroutine _pprint_key_val_tuple() used when a parameter or a key-value + pair is too long to fit on a single line. This subroutine is called in + _format() and is registered as well in the _dispatch dict (just like + _pprint_estimator). We had to create the two classes KeyValTuple and + KeyValTupleParam for this. + """ + + def __init__( + self, + indent=1, + width=80, + depth=None, + stream=None, + *, + compact=False, + indent_at_name=True, + n_max_elements_to_show=None, + ): + super().__init__(indent, width, depth, stream, compact=compact) + self._indent_at_name = indent_at_name + if self._indent_at_name: + self._indent_per_level = 1 # ignore indent param + # Max number of elements in a list, dict, tuple until we start using + # ellipsis. This also affects the number of arguments of an estimators + # (they are treated as dicts) + self.n_max_elements_to_show = n_max_elements_to_show + self._changed_only = True + + def format(self, object, context, maxlevels, level): + return _safe_repr( + object, context, maxlevels, level, changed_only=self._changed_only + ) + + def _pprint_estimator(self, object, stream, indent, allowance, context, level): + stream.write(object.__class__.__name__ + "(") + if self._indent_at_name: + indent += len(object.__class__.__name__) + + params = _changed_params(object) + + params = OrderedDict((name, val) for (name, val) in sorted(params.items())) + + self._format_params( + params.items(), stream, indent, allowance + 1, context, level + ) + stream.write(")") + + def _format_dict_items(self, items, stream, indent, allowance, context, level): + return self._format_params_or_dict_items( + items, stream, indent, allowance, context, level, is_dict=True + ) + + def _format_params(self, items, stream, indent, allowance, context, level): + return self._format_params_or_dict_items( + items, stream, indent, allowance, context, level, is_dict=False + ) + + def _format_params_or_dict_items( + self, object, stream, indent, allowance, context, level, is_dict + ): + """Format dict items or parameters respecting the compact=True + parameter. For some reason, the builtin rendering of dict items doesn't + respect compact=True and will use one line per key-value if all cannot + fit in a single line. + Dict items will be rendered as <'key': value> while params will be + rendered as . The implementation is mostly copy/pasting from + the builtin _format_items(). + This also adds ellipsis if the number of items is greater than + self.n_max_elements_to_show. + """ + write = stream.write + indent += self._indent_per_level + delimnl = ",\n" + " " * indent + delim = "" + width = max_width = self._width - indent + 1 + it = iter(object) + try: + next_ent = next(it) + except StopIteration: + return + last = False + n_items = 0 + while not last: + if n_items == self.n_max_elements_to_show: + write(", ...") + break + n_items += 1 + ent = next_ent + try: + next_ent = next(it) + except StopIteration: + last = True + max_width -= allowance + width -= allowance + if self._compact: + k, v = ent + krepr = self._repr(k, context, level) + vrepr = self._repr(v, context, level) + if not is_dict: + krepr = krepr.strip("'") + middle = ": " if is_dict else "=" + rep = krepr + middle + vrepr + w = len(rep) + 2 + if width < w: + width = max_width + if delim: + delim = delimnl + if width >= w: + width -= w + write(delim) + delim = ", " + write(rep) + continue + write(delim) + delim = delimnl + class_ = KeyValTuple if is_dict else KeyValTupleParam + self._format( + class_(ent), stream, indent, allowance if last else 1, context, level + ) + + def _format_items(self, items, stream, indent, allowance, context, level): + """Format the items of an iterable (list, tuple...). Same as the + built-in _format_items, with support for ellipsis if the number of + elements is greater than self.n_max_elements_to_show. + """ + write = stream.write + indent += self._indent_per_level + if self._indent_per_level > 1: + write((self._indent_per_level - 1) * " ") + delimnl = ",\n" + " " * indent + delim = "" + width = max_width = self._width - indent + 1 + it = iter(items) + try: + next_ent = next(it) + except StopIteration: + return + last = False + n_items = 0 + while not last: + if n_items == self.n_max_elements_to_show: + write(", ...") + break + n_items += 1 + ent = next_ent + try: + next_ent = next(it) + except StopIteration: + last = True + max_width -= allowance + width -= allowance + if self._compact: + rep = self._repr(ent, context, level) + w = len(rep) + 2 + if width < w: + width = max_width + if delim: + delim = delimnl + if width >= w: + width -= w + write(delim) + delim = ", " + write(rep) + continue + write(delim) + delim = delimnl + self._format(ent, stream, indent, allowance if last else 1, context, level) + + def _pprint_key_val_tuple(self, object, stream, indent, allowance, context, level): + """Pretty printing for key-value tuples from dict or parameters.""" + k, v = object + rep = self._repr(k, context, level) + if isinstance(object, KeyValTupleParam): + rep = rep.strip("'") + middle = "=" + else: + middle = ": " + stream.write(rep) + stream.write(middle) + self._format( + v, stream, indent + len(rep) + len(middle), allowance, context, level + ) + + # Note: need to copy _dispatch to prevent instances of the builtin + # PrettyPrinter class to call methods of _EstimatorPrettyPrinter (see issue + # 12906) + # mypy error: "Type[PrettyPrinter]" has no attribute "_dispatch" + _dispatch = pprint.PrettyPrinter._dispatch.copy() # type: ignore + _dispatch[BaseEstimator.__repr__] = _pprint_estimator + _dispatch[KeyValTuple.__repr__] = _pprint_key_val_tuple + + +def _safe_repr(object, context, maxlevels, level, changed_only=False): + """Same as the builtin _safe_repr, with added support for Estimator + objects.""" + typ = type(object) + + if typ in pprint._builtin_scalars: + return repr(object), True, False + + r = getattr(typ, "__repr__", None) + if issubclass(typ, dict) and r is dict.__repr__: + if not object: + return "{}", True, False + objid = id(object) + if maxlevels and level >= maxlevels: + return "{...}", False, objid in context + if objid in context: + return pprint._recursion(object), False, True + context[objid] = 1 + readable = True + recursive = False + components = [] + append = components.append + level += 1 + saferepr = _safe_repr + items = sorted(object.items(), key=pprint._safe_tuple) + for k, v in items: + krepr, kreadable, krecur = saferepr( + k, context, maxlevels, level, changed_only=changed_only + ) + vrepr, vreadable, vrecur = saferepr( + v, context, maxlevels, level, changed_only=changed_only + ) + append("%s: %s" % (krepr, vrepr)) + readable = readable and kreadable and vreadable + if krecur or vrecur: + recursive = True + del context[objid] + return "{%s}" % ", ".join(components), readable, recursive + + if (issubclass(typ, list) and r is list.__repr__) or ( + issubclass(typ, tuple) and r is tuple.__repr__ + ): + if issubclass(typ, list): + if not object: + return "[]", True, False + format = "[%s]" + elif len(object) == 1: + format = "(%s,)" + else: + if not object: + return "()", True, False + format = "(%s)" + objid = id(object) + if maxlevels and level >= maxlevels: + return format % "...", False, objid in context + if objid in context: + return pprint._recursion(object), False, True + context[objid] = 1 + readable = True + recursive = False + components = [] + append = components.append + level += 1 + for o in object: + orepr, oreadable, orecur = _safe_repr( + o, context, maxlevels, level, changed_only=changed_only + ) + append(orepr) + if not oreadable: + readable = False + if orecur: + recursive = True + del context[objid] + return format % ", ".join(components), readable, recursive + + if issubclass(typ, BaseEstimator): + objid = id(object) + if maxlevels and level >= maxlevels: + return "{...}", False, objid in context + if objid in context: + return pprint._recursion(object), False, True + context[objid] = 1 + readable = True + recursive = False + if changed_only: + params = _changed_params(object) + else: + params = object.get_params(deep=False) + components = [] + append = components.append + level += 1 + saferepr = _safe_repr + items = sorted(params.items(), key=pprint._safe_tuple) + for k, v in items: + krepr, kreadable, krecur = saferepr( + k, context, maxlevels, level, changed_only=changed_only + ) + vrepr, vreadable, vrecur = saferepr( + v, context, maxlevels, level, changed_only=changed_only + ) + append("%s=%s" % (krepr.strip("'"), vrepr)) + readable = readable and kreadable and vreadable + if krecur or vrecur: + recursive = True + del context[objid] + return ("%s(%s)" % (typ.__name__, ", ".join(components)), readable, recursive) + + rep = repr(object) + return rep, (rep and not rep.startswith("<")), False diff --git a/third_party/bigframes_vendored/pandas/AUTHORS.md b/third_party/bigframes_vendored/pandas/AUTHORS.md new file mode 100644 index 00000000000..84fcfe05e30 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/AUTHORS.md @@ -0,0 +1,56 @@ +About the Copyright Holders +=========================== + +* Copyright (c) 2008-2011 AQR Capital Management, LLC + + AQR Capital Management began pandas development in 2008. Development was + led by Wes McKinney. AQR released the source under this license in 2009. +* Copyright (c) 2011-2012, Lambda Foundry, Inc. + + Wes is now an employee of Lambda Foundry, and remains the pandas project + lead. +* Copyright (c) 2011-2012, PyData Development Team + + The PyData Development Team is the collection of developers of the PyData + project. This includes all of the PyData sub-projects, including pandas. The + core team that coordinates development on GitHub can be found here: + https://github.com/pydata. + +Full credits for pandas contributors can be found in the documentation. + +Our Copyright Policy +==================== + +PyData uses a shared copyright model. Each contributor maintains copyright +over their contributions to PyData. However, it is important to note that +these contributions are typically only changes to the repositories. Thus, +the PyData source code, in its entirety, is not the copyright of any single +person or institution. Instead, it is the collective copyright of the +entire PyData Development Team. If individual contributors want to maintain +a record of what changes/contributions they have specific copyright on, +they should indicate their copyright in the commit message of the change +when they commit the change to one of the PyData repositories. + +With this in mind, the following banner should be used in any source code +file to indicate the copyright and license terms: + +``` +#----------------------------------------------------------------------------- +# Copyright (c) 2012, PyData Development Team +# All rights reserved. +# +# Distributed under the terms of the BSD Simplified License. +# +# The full license is in the LICENSE file, distributed with this software. +#----------------------------------------------------------------------------- +``` + +Other licenses can be found in the LICENSES directory. + +License +======= + +pandas is distributed under a 3-clause ("Simplified" or "New") BSD +license. Parts of NumPy, SciPy, numpydoc, bottleneck, which all have +BSD-compatible licenses, are included. Their licenses follow the pandas +license. diff --git a/third_party/bigframes_vendored/pandas/LICENSE b/third_party/bigframes_vendored/pandas/LICENSE new file mode 100644 index 00000000000..cdfa749dc34 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/LICENSE @@ -0,0 +1,31 @@ +BSD 3-Clause License + +Copyright (c) 2008-2011, AQR Capital Management, LLC, Lambda Foundry, Inc. and PyData Development Team +All rights reserved. + +Copyright (c) 2011-2023, Open source contributors. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/third_party/bigframes_vendored/pandas/README.md b/third_party/bigframes_vendored/pandas/README.md new file mode 100644 index 00000000000..9f2bc800e84 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/README.md @@ -0,0 +1,174 @@ +
+
+
+ +----------------- + +# pandas: powerful Python data analysis toolkit +[![PyPI Latest Release](https://img.shields.io/pypi/v/pandas.svg)](https://pypi.org/project/pandas/) +[![Conda Latest Release](https://anaconda.org/conda-forge/pandas/badges/version.svg)](https://anaconda.org/anaconda/pandas/) +[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.3509134.svg)](https://doi.org/10.5281/zenodo.3509134) +[![Package Status](https://img.shields.io/pypi/status/pandas.svg)](https://pypi.org/project/pandas/) +[![License](https://img.shields.io/pypi/l/pandas.svg)](https://github.com/pandas-dev/pandas/blob/main/LICENSE) +[![Coverage](https://codecov.io/github/pandas-dev/pandas/coverage.svg?branch=main)](https://codecov.io/gh/pandas-dev/pandas) +[![Downloads](https://static.pepy.tech/personalized-badge/pandas?period=month&units=international_system&left_color=black&right_color=orange&left_text=PyPI%20downloads%20per%20month)](https://pepy.tech/project/pandas) +[![Slack](https://img.shields.io/badge/join_Slack-information-brightgreen.svg?logo=slack)](https://pandas.pydata.org/docs/dev/development/community.html?highlight=slack#community-slack) +[![Powered by NumFOCUS](https://img.shields.io/badge/powered%20by-NumFOCUS-orange.svg?style=flat&colorA=E1523D&colorB=007D8A)](https://numfocus.org) +[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) +[![Imports: isort](https://img.shields.io/badge/%20imports-isort-%231674b1?style=flat&labelColor=ef8336)](https://pycqa.github.io/isort/) + +## What is it? + +**pandas** is a Python package that provides fast, flexible, and expressive data +structures designed to make working with "relational" or "labeled" data both +easy and intuitive. It aims to be the fundamental high-level building block for +doing practical, **real world** data analysis in Python. Additionally, it has +the broader goal of becoming **the most powerful and flexible open source data +analysis / manipulation tool available in any language**. It is already well on +its way towards this goal. + +## Main Features +Here are just a few of the things that pandas does well: + + - Easy handling of [**missing data**][missing-data] (represented as + `NaN`, `NA`, or `NaT`) in floating point as well as non-floating point data + - Size mutability: columns can be [**inserted and + deleted**][insertion-deletion] from DataFrame and higher dimensional + objects + - Automatic and explicit [**data alignment**][alignment]: objects can + be explicitly aligned to a set of labels, or the user can simply + ignore the labels and let `Series`, `DataFrame`, etc. automatically + align the data for you in computations + - Powerful, flexible [**group by**][groupby] functionality to perform + split-apply-combine operations on data sets, for both aggregating + and transforming data + - Make it [**easy to convert**][conversion] ragged, + differently-indexed data in other Python and NumPy data structures + into DataFrame objects + - Intelligent label-based [**slicing**][slicing], [**fancy + indexing**][fancy-indexing], and [**subsetting**][subsetting] of + large data sets + - Intuitive [**merging**][merging] and [**joining**][joining] data + sets + - Flexible [**reshaping**][reshape] and [**pivoting**][pivot-table] of + data sets + - [**Hierarchical**][mi] labeling of axes (possible to have multiple + labels per tick) + - Robust IO tools for loading data from [**flat files**][flat-files] + (CSV and delimited), [**Excel files**][excel], [**databases**][db], + and saving/loading data from the ultrafast [**HDF5 format**][hdfstore] + - [**Time series**][timeseries]-specific functionality: date range + generation and frequency conversion, moving window statistics, + date shifting and lagging + + + [missing-data]: https://pandas.pydata.org/pandas-docs/stable/user_guide/missing_data.html + [insertion-deletion]: https://pandas.pydata.org/pandas-docs/stable/user_guide/dsintro.html#column-selection-addition-deletion + [alignment]: https://pandas.pydata.org/pandas-docs/stable/user_guide/dsintro.html?highlight=alignment#intro-to-data-structures + [groupby]: https://pandas.pydata.org/pandas-docs/stable/user_guide/groupby.html#group-by-split-apply-combine + [conversion]: https://pandas.pydata.org/pandas-docs/stable/user_guide/dsintro.html#dataframe + [slicing]: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#slicing-ranges + [fancy-indexing]: https://pandas.pydata.org/pandas-docs/stable/user_guide/advanced.html#advanced + [subsetting]: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#boolean-indexing + [merging]: https://pandas.pydata.org/pandas-docs/stable/user_guide/merging.html#database-style-dataframe-or-named-series-joining-merging + [joining]: https://pandas.pydata.org/pandas-docs/stable/user_guide/merging.html#joining-on-index + [reshape]: https://pandas.pydata.org/pandas-docs/stable/user_guide/reshaping.html + [pivot-table]: https://pandas.pydata.org/pandas-docs/stable/user_guide/reshaping.html + [mi]: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#hierarchical-indexing-multiindex + [flat-files]: https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#csv-text-files + [excel]: https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#excel-files + [db]: https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#sql-queries + [hdfstore]: https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#hdf5-pytables + [timeseries]: https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#time-series-date-functionality + +## Where to get it +The source code is currently hosted on GitHub at: +https://github.com/pandas-dev/pandas + +Binary installers for the latest released version are available at the [Python +Package Index (PyPI)](https://pypi.org/project/pandas) and on [Conda](https://docs.conda.io/en/latest/). + +```sh +# conda +conda install -c conda-forge pandas +``` + +```sh +# or PyPI +pip install pandas +``` + +The list of changes to pandas between each release can be found +[here](https://pandas.pydata.org/pandas-docs/stable/whatsnew/index.html). For full +details, see the commit logs at https://github.com/pandas-dev/pandas. + +## Dependencies +- [NumPy - Adds support for large, multi-dimensional arrays, matrices and high-level mathematical functions to operate on these arrays](https://www.numpy.org) +- [python-dateutil - Provides powerful extensions to the standard datetime module](https://dateutil.readthedocs.io/en/stable/index.html) +- [pytz - Brings the Olson tz database into Python which allows accurate and cross platform timezone calculations](https://github.com/stub42/pytz) + +See the [full installation instructions](https://pandas.pydata.org/pandas-docs/stable/install.html#dependencies) for minimum supported versions of required, recommended and optional dependencies. + +## Installation from sources +To install pandas from source you need [Cython](https://cython.org/) in addition to the normal +dependencies above. Cython can be installed from PyPI: + +```sh +pip install cython +``` + +In the `pandas` directory (same one where you found this file after +cloning the git repo), execute: + +```sh +python setup.py install +``` + +or for installing in [development mode](https://pip.pypa.io/en/latest/cli/pip_install/#install-editable): + + +```sh +python -m pip install -e . --no-build-isolation --no-use-pep517 +``` + +or alternatively + +```sh +python setup.py develop +``` + +See the full instructions for [installing from source](https://pandas.pydata.org/pandas-docs/stable/getting_started/install.html#installing-from-source). + +## License +[BSD 3](LICENSE) + +## Documentation +The official documentation is hosted on PyData.org: https://pandas.pydata.org/pandas-docs/stable + +## Background +Work on ``pandas`` started at [AQR](https://www.aqr.com/) (a quantitative hedge fund) in 2008 and +has been under active development since then. + +## Getting Help + +For usage questions, the best place to go to is [StackOverflow](https://stackoverflow.com/questions/tagged/pandas). +Further, general questions and discussions can also take place on the [pydata mailing list](https://groups.google.com/forum/?fromgroups#!forum/pydata). + +## Discussion and Development +Most development discussions take place on GitHub in this repo. Further, the [pandas-dev mailing list](https://mail.python.org/mailman/listinfo/pandas-dev) can also be used for specialized discussions or design issues, and a [Slack channel](https://pandas.pydata.org/docs/dev/development/community.html?highlight=slack#community-slack) is available for quick development related questions. + +## Contributing to pandas [![Open Source Helpers](https://www.codetriage.com/pandas-dev/pandas/badges/users.svg)](https://www.codetriage.com/pandas-dev/pandas) + +All contributions, bug reports, bug fixes, documentation improvements, enhancements, and ideas are welcome. + +A detailed overview on how to contribute can be found in the **[contributing guide](https://pandas.pydata.org/docs/dev/development/contributing.html)**. + +If you are simply looking to start working with the pandas codebase, navigate to the [GitHub "issues" tab](https://github.com/pandas-dev/pandas/issues) and start looking through interesting issues. There are a number of issues listed under [Docs](https://github.com/pandas-dev/pandas/issues?labels=Docs&sort=updated&state=open) and [good first issue](https://github.com/pandas-dev/pandas/issues?labels=good+first+issue&sort=updated&state=open) where you could start out. + +You can also triage issues which may include reproducing bug reports, or asking for vital information such as version numbers or reproduction instructions. If you would like to start triaging issues, one easy way to get started is to [subscribe to pandas on CodeTriage](https://www.codetriage.com/pandas-dev/pandas). + +Or maybe through using pandas you have an idea of your own or are looking for something in the documentation and thinking ‘this can be improved’...you can do something about it! + +Feel free to ask questions on the [mailing list](https://groups.google.com/forum/?fromgroups#!forum/pydata) or on [Slack](https://pandas.pydata.org/docs/dev/development/community.html?highlight=slack#community-slack). + +As contributors and maintainers to this project, you are expected to abide by pandas' code of conduct. More information can be found at: [Contributor Code of Conduct](https://github.com/pandas-dev/.github/blob/master/CODE_OF_CONDUCT.md) diff --git a/third_party/bigframes_vendored/pandas/__init__.py b/third_party/bigframes_vendored/pandas/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/config_init.py b/third_party/bigframes_vendored/pandas/core/config_init.py new file mode 100644 index 00000000000..10f5546f108 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/config_init.py @@ -0,0 +1,25 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/config_init.py +""" +This module is imported from the pandas package __init__.py file +in order to ensure that the core.config options registered here will +be available as soon as the user loads the package. if register_option +is invoked inside specific modules, they will not be registered until that +module is imported, which may or may not be a problem. + +If you need to make sure options are available even before a certain +module is imported, register them here rather than in the module. + +""" +from __future__ import annotations + +display_options_doc = """ +Encapsulates configuration for displaying objects. + +Attributes: + max_columns (int): + If max_columns is exceeded, switch to truncate view. + max_rows (int): + If max_rows is exceeded, switch to truncate view. + progress_bar Optional(str): + Determines if progress bars are shown during job runs. Use auto | notebook | terminal. Set to 'None' to remove progress bars +""" diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py new file mode 100644 index 00000000000..93b7c7683cf --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -0,0 +1,1124 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/frame.py +""" +DataFrame +--------- +An efficient 2D container for potentially mixed-type time series or other +labeled data series. + +Similar to its R counterpart, data.frame, except providing automatic data +alignment and a host of useful data manipulation methods having to do with the +labeling information +""" +from __future__ import annotations + +from typing import Iterable, Literal, Mapping, Optional, Sequence, Union + +from third_party.bigframes_vendored.pandas.core.generic import NDFrame + +# ----------------------------------------------------------------------- +# DataFrame class + + +class DataFrame(NDFrame): + """Two-dimensional, size-mutable, potentially heterogeneous tabular data. + + Data structure also contains labeled axes (rows and columns). + Arithmetic operations align on both row and column labels. Can be + thought of as a dict-like container for Series objects. The primary + pandas data structure. + """ + + @property + def shape(self) -> tuple[int, int]: + """Return a tuple representing the dimensionality of the DataFrame.""" + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # IO methods (to / from other formats) + + def to_gbq( + self, + destination_table: str, + *, + if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", + index: bool = True, + ) -> None: + """Write a DataFrame to a Google BigQuery table. + + Args: + destination_table : str + Name of table to be written, in the form ``dataset.tablename`` + or ``project.dataset.tablename``. + + if_exists : str, default 'fail' + Behavior when the destination table exists. Value can be one of: + + ``'fail'`` + If table exists raise pandas_gbq.gbq.TableCreationError. + ``'replace'`` + If table exists, drop it, recreate it, and insert data. + ``'append'`` + If table exists, insert data. Create if does not exist. + + index: whether write row names (index) or not. + """ + raise NotImplementedError("abstract method") + + def to_parquet( + self, + path: str, + *, + index: bool = True, + ) -> None: + """Write a DataFrame to the binary parquet format. + + This function writes the dataframe as a `parquet file + `_ to Google Cloud Storage. + + Args: + path: + Destination URI(s) of GCS files(s) to store the extracted dataframe + in format of ``gs:///``. + If the data size is more than 1GB, you must use a wildcard to export + the data into multiple files and the size of the files varies. + + index : bool, default None + If ``True``, include the dataframe's index(es) in the file output. + If ``False``, they will not be written to the file. + + Returns: + None. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Unsorted + + def assign(self, **kwargs) -> DataFrame: + r""" + Assign new columns to a DataFrame. + + Returns a new object with all original columns in addition to new ones. + Existing columns that are re-assigned will be overwritten. + + Args: + kwargs: + A dictionary of ``{str: values}``. The column names are + keywords. If the values (e.g. a Series, scalar, or array), they + are simply assigned to the column. + + Returns: + A new DataFrame with the new columns in addition to + all the existing columns. + + Notes: + Assigning multiple columns within the same ``assign`` is possible. + Later items in '\*\*kwargs' may refer to newly created or modified + columns in 'df'; items are computed and assigned into 'df' in + order. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Reindexing and alignment + + def drop( + self, + *, + columns: Union[str, Iterable[str]], + ) -> DataFrame | None: + """Drop specified labels from columns. + + Remove columns by directly specifying column names. + + Args: + columns: + Single label or list-like. + + Returns: + DataFrame without the removed column labels. + + Raises: + KeyError: + If any of the labels is not found in the selected axis. + """ + raise NotImplementedError("abstract method") + + def rename( + self, + *, + columns: Mapping, + ) -> DataFrame: + """Rename columns. + + Dict values must be unique (1-to-1). Labels not contained in a dict + will be left as-is. Extra labels listed don't throw an error. + + Args: + columns: + Dict-like from old column labels to new column labels. + + Returns: + DataFrame with the renamed axis labels. + + Raises: + KeyError: + If any of the labels is not found. + """ + raise NotImplementedError("abstract method") + + def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: + """ + Set the name of the axis for the index. + + Note: currently only accepts a single string parameter (the new name of the index) + + Args: + mapper : str + Value to set the axis name attribute. + + Returns: + DataFrame with the new index name + """ + raise NotImplementedError("abstract method") + + def set_index( + self, + keys, + *, + drop: bool = True, + ) -> DataFrame | None: + """ + Set the DataFrame index using existing columns. + + Set the DataFrame index (row labels) using one existing column. The + index can replace the existing index. + + Args: + keys: + A label. This parameter can be a single column key. + drop: + Delete columns to be used as the new index. + + Returns: + DataFrame: + Changed row labels. + """ + raise NotImplementedError("abstract method") + + def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: + """ + Rearrange index levels using input order. May not drop or duplicate levels. + + Args: + order: list of int or list of str + List representing new level order. Reference level by number + (position) or by key (label). + + Returns: + DataFrame + """ + raise NotImplementedError("abstract method") + + def droplevel(self, level): + """ + Return DataFrame with requested index / column level(s) removed. + + Args: + level: int, str, or list-like + If a string is given, must be the name of a level + If list-like, elements must be names or positional indexes + of levels. + Returns: + DataFrame with requested index / column level(s) removed. + """ + raise NotImplementedError("abstract method") + + def reset_index( + self, + *, + drop: bool = False, + ) -> DataFrame | None: + """Reset the index. + + Reset the index of the DataFrame, and use the default one instead. + + Args: + drop: + Do not try to insert index into dataframe columns. This resets + the index to the default integer index. + + Returns: + DataFrame with the new index. + """ + raise NotImplementedError("abstract method") + + def drop_duplicates( + self, + *, + keep="first", + ) -> DataFrame: + """ + Return DataFrame with duplicate rows removed. + + Considering certain columns is optional. Indexes, including time indexes + are ignored. + + Args: + subset : column label or sequence of labels, optional + Only consider certain columns for identifying duplicates, by + default use all of the columns. + keep : {'first', 'last', ``False``}, default 'first' + Determines which duplicates (if any) to keep. + + - 'first' : Drop duplicates except for the first occurrence. + - 'last' : Drop duplicates except for the last occurrence. + - ``False`` : Drop all duplicates. + + Returns: + DataFrame with duplicates removed + """ + raise NotImplementedError("abstract method") + + def duplicated(self, subset=None, keep="first"): + """ + Return boolean Series denoting duplicate rows. + + Considering certain columns is optional. + + Args: + subset : column label or sequence of labels, optional + Only consider certain columns for identifying duplicates, by + default use all of the columns. + keep : {'first', 'last', False}, default 'first' + Determines which duplicates (if any) to mark. + + - ``first`` : Mark duplicates as ``True`` except for the first occurrence. + - ``last`` : Mark duplicates as ``True`` except for the last occurrence. + - False : Mark all duplicates as ``True``. + + Returns: + Boolean series for each duplicated rows. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Reindex-based selection methods + + def dropna( + self, + ) -> DataFrame: + """Remove missing values. + + Returns: + DataFrame with NA entries dropped from it. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Sorting + + def sort_values( + self, + by: str | Sequence[str], + *, + ascending: bool | Sequence[bool] = True, + na_position="last", + ) -> DataFrame: + """Sort by the values along row axis. + + Args: + by: + Name or list of names to sort by. + ascending: + Sort ascending vs. descending. Specify list for multiple sort + orders. If this is a list of bools, must match the length of + the by. + na_position: + ``{'first', 'last'}``, default 'last' Puts NaNs at the beginning + if `first`; `last` puts NaNs at the end. + + Returns: + DataFrame with sorted values. + """ + raise NotImplementedError("abstract method") + + def sort_index( + self, + ) -> DataFrame: + """Sort object by labels (along an axis). + + Returns: + The original DataFrame sorted by the labels. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Arithmetic Methods + + def le(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'less than or equal to' of dataframe and other, element-wise (binary operator `<=`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + Args: + other: scalar, sequence, Series, or DataFrame + Any single or multiple element data structure, or list-like object. + axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame of bool. The result of the comparison. + + Notes: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + """ + raise NotImplementedError("abstract method") + + def lt(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'less than' of dataframe and other, element-wise (binary operator `<`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + Args: + other: scalar, sequence, Series, or DataFrame + Any single or multiple element data structure, or list-like object. + axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame of bool. The result of the comparison. + + Notes: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + """ + raise NotImplementedError("abstract method") + + def ge(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'greater than or equal to' of dataframe and other, element-wise (binary operator `>=`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + Args: + other: scalar, sequence, Series, or DataFrame + Any single or multiple element data structure, or list-like object. + axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame of bool. The result of the comparison. + + Notes: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + """ + raise NotImplementedError("abstract method") + + def gt(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'greater than' of dataframe and other, element-wise (binary operator `>`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + Args: + other: scalar, sequence, Series, or DataFrame + Any single or multiple element data structure, or list-like object. + axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame of bool. The result of the comparison. + + Notes: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + """ + raise NotImplementedError("abstract method") + + def add(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'addition' of dataframe and other, element-wise (binary operator `+`). + + Equivalent to ``dataframe + other``. With reverse version, `radd`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def sub(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'subtraction' of dataframe and other, element-wise (binary operator `-`). + + Equivalent to ``dataframe - other``. With reverse version, `rsub`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def rsub(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'subtraction' of dataframe and other, element-wise (binary operator `-`). + + Equivalent to ``other - dataframe``. With reverse version, `sub`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def mul(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'multiplication' of dataframe and other, element-wise (binary operator `*`). + + Equivalent to ``dataframe * other``. With reverse version, `rmul`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def truediv(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'floating division' of dataframe and other, element-wise (binary operator `/`). + + Equivalent to ``dataframe / other``. With reverse version, `rtruediv`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'floating division' of dataframe and other, element-wise (binary operator `/`). + + Equivalent to ``other / dataframe``. With reverse version, `truediv`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def floordiv(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'integer division' of dataframe and other, element-wise (binary operator `//`). + + Equivalent to ``dataframe // other``. With reverse version, `rfloordiv`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def rfloordiv(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'integer division' of dataframe and other, element-wise (binary operator `//`). + + Equivalent to ``other // dataframe``. With reverse version, `rfloordiv`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def mod(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'modulo' of dataframe and other, element-wise (binary operator `%`). + + Equivalent to ``dataframe % other``. With reverse version, `rmod`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + def rmod(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'modulo' of dataframe and other, element-wise (binary operator `%`). + + Equivalent to ``other % dataframe``. With reverse version, `mod`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + Args: + other: + Any single or multiple element data structure, or list-like object. + axis: + ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the + index (0 or 'index') or columns. (1 or 'columns'). For Series + input, axis to match Series index on. + + Returns: + DataFrame result of the arithmetic operation. + + Notes: + Mismatched indices will be unioned together. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Data reshaping + + def groupby( + self, + by: Union[str, Sequence[str]], + *, + level=None, + as_index: bool = True, + dropna: bool = True, + ): + """Group DataFrame by columns. + + A groupby operation involves some combination of splitting the + object, applying a function, and combining the results. This can be + used to group large amounts of data and compute operations on these + groups. + + Args: + by: + A label or list of labels may be passed to group by the columns + in ``self``. Notice that a tuple is interpreted as a (single) + key. + level : int, level name, or sequence of such, default None + If the axis is a MultiIndex (hierarchical), group by a particular + level or levels. Do not specify both ``by`` and ``level``. + as_index: + Default True. Return object with group labels as the index. + Only relevant for DataFrame input. ``as_index=False`` is + effectively "SQL-style" grouped output. This argument has no + effect on filtrations such as ``head()``, ``tail()``, ``nth()`` + and in transformations. + dropna: + Default True. If True, and if group keys contain NA values, NA + values together with row/column will be dropped. If False, NA + values will also be treated as the key in groups. + + Returns: + A groupby object that contains information about the groups. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Function application + + def map(self, func, na_action: Optional[str] = None) -> DataFrame: + """Apply a function to a Dataframe elementwise. + + .. note:: + + In pandas 2.1.0, DataFrame.applymap is deprecated and renamed to + DataFrame.map. + + This method applies a function that accepts and returns a scalar + to every element of a DataFrame. + + Args: + func: + Python function wrapped by ``remote_function`` decorator, + returns a single value from a single value. + na_action: + ``{None, 'ignore'}``, default None. If ‘ignore’, propagate NaN + values, without passing them to func. + + Returns: + Transformed DataFrame. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Merging / joining methods + + def join(self, other, *, how: str) -> DataFrame: + """Join columns of another DataFrame. + + Join columns with `other` DataFrame on index + + Args: + other: + DataFrame with an Index similar to the Index of this one. + how: + ``{'left', 'right', 'outer', 'inner'}, default 'left'`` + + How to handle the operation of the two objects. + + * left: use calling frame's index (or column if on is specified) + * right: use `other`'s index. + * outer: form union of calling frame's index (or column if on is + specified) with `other`'s index, and sort it. + lexicographically. + * inner: form intersection of calling frame's index (or column if + on is specified) with `other`'s index, preserving the order + of the calling's one. + + Returns: + A dataframe containing columns from both the caller and `other`. + """ + raise NotImplementedError("abstract method") + + def merge( + self, + right, + how: Literal[ + "inner", + "left", + "outer", + "right", + ] = "inner", + on: Optional[str] = None, + *, + left_on: Optional[str] = None, + right_on: Optional[str] = None, + sort: bool = False, + suffixes: tuple[str, str] = ("_x", "_y"), + ) -> DataFrame: + """Merge DataFrame objects with a database-style join. + + The join is done on columns or indexes. If joining columns on + columns, the DataFrame indexes *will be ignored*. Otherwise if joining indexes + on indexes or indexes on a column or columns, the index will be passed on. + When performing a cross merge, no column specifications to merge on are + allowed. + + .. warning:: + + If both key columns contain rows where the key is a null value, those + rows will be matched against each other. This is different from usual SQL + join behaviour and can lead to unexpected results. + + Args: + right: + Object to merge with. + how: + ``{'left', 'right', 'outer', 'inner'}, default 'inner'`` + Type of merge to be performed. + + * left: use only keys from left frame, similar to a SQL left outer join; + preserve key order. + * right: use only keys from right frame, similar to a SQL right outer join; + preserve key order. + * outer: use union of keys from both frames, similar to a SQL full outer + join; sort keys lexicographically. + * inner: use intersection of keys from both frames, similar to a SQL inner + join; preserve the order of the left keys. + + on: + Column join on. It must be found in both DataFrames. Either on or left_on + right_on + must be passed in. + left_on: + Column join on in the left DataFrame. Either on or left_on + right_on + must be passed in. + right_on: + Column join on in the right DataFrame. Either on or left_on + right_on + must be passed in. + sort: + Default False. Sort the join keys lexicographically in the + result DataFrame. If False, the order of the join keys depends + on the join type (how keyword). + suffixes: + Default ``("_x", "_y")``. A length-2 sequence where each + element is optionally a string indicating the suffix to add to + overlapping column names in `left` and `right` respectively. + Pass a value of `None` instead of a string to indicate that the + column name from `left` or `right` should be left as-is, with + no suffix. At least one of the values must not be None. + + Returns: + A DataFrame of the two merged objects. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # ndarray-like stats methods + + def any(self, *, bool_only: bool = False): + """ + Return whether any element is True, potentially over an axis. + + Returns False unless there is at least one element within a series or + along a Dataframe axis that is True or equivalent (e.g. non-zero or + non-empty). + + Args: + bool_only: + Include only boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def all(self, *, bool_only: bool = False): + """ + Return whether all elements are True, potentially over an axis. + + Returns True unless there at least one element within a series or + along a Dataframe axis that is False or equivalent (e.g. zero or + empty). + + Args: + bool_only: + Include only boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def prod(self, *, numeric_only: bool = False): + """ + Return the product of the values over the requested axis. + + Args: + numeric_only: + Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def min(self, *, numeric_only: bool = False): + """Return the minimum of the values over the requested axis. + + If you want the *index* of the minimum, use ``idxmin``. This is the + equivalent of the ``numpy.ndarray`` method ``argmin``. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def max(self, *, numeric_only: bool = False): + """Return the maximum of the values over the requested axis. + + If you want the *index* of the maximum, use ``idxmax``. This is + the equivalent of the ``numpy.ndarray`` method ``argmax``. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def sum(self, *, numeric_only: bool = False): + """Return the sum of the values over the requested axis. + + This is equivalent to the method ``numpy.sum``. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def mean(self, *, numeric_only: bool = False): + """Return the mean of the values over the requested axis. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def var(self, *, numeric_only: bool = False): + """Return unbiased variance over requested axis. + + Normalized by N-1 by default. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def std(self, *, numeric_only: bool = False): + """Return sample standard deviation over requested axis. + + Normalized by N-1 by default. + + Args: + numeric_only: + Default False. Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def count(self, *, numeric_only: bool = False): + """ + Count non-NA cells for each column or row. + + The values `None`, `NaN`, `NaT`, and optionally `numpy.inf` (depending + on `pandas.options.mode.use_inf_as_na`) are considered NA. + + Args: + numeric_only : bool, default False + Include only `float`, `int` or `boolean` data. + + Returns: + For each column/row the number of non-NA/null entries. + If `level` is specified returns a `DataFrame`. + """ + raise NotImplementedError("abstract method") + + def nunique(self): + """ + Count number of distinct elements in specified axis. + + Return Series with number of distinct elements. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + def cummin(self) -> DataFrame: + """Return cumulative minimum over a DataFrame axis. + + Returns a DataFrame of the same size containing the cumulative minimum. + + Returns: + Return cumulative minimum of DataFrame. + """ + raise NotImplementedError("abstract method") + + def cummax(self) -> DataFrame: + """Return cumulative maximum over a DataFrame axis. + + Returns a DataFrame of the same size containing the cumulative maximum. + + Returns: + Return cumulative maximum of DataFrame. + """ + raise NotImplementedError("abstract method") + + def cumsum(self) -> DataFrame: + """Return cumulative sum over a DataFrame axis. + + Returns a DataFrame of the same size containing the cumulative sum. + + Returns: + Return cumulative sum of DataFrame. + """ + raise NotImplementedError("abstract method") + + def cumprod(self) -> DataFrame: + """Return cumulative product over a DataFrame axis. + + Returns a DataFrame of the same size containing the cumulative product. + + Returns: + Return cumulative product of DataFrame. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Add index and columns + + @property + def index(self): + """The index (row labels) of the DataFrame. + + The index of a DataFrame is a series of labels that identify each row. + The labels can be integers, strings, or any other hashable type. The + index is used for label-based access and alignment, and can be accessed + or modified using this attribute. + + Returns: + The index labels of the DataFrame. + """ + raise NotImplementedError("abstract method") + + @property + def columns(self): + "The column labels of the DataFrame." + raise NotImplementedError("abstract method") + + def value_counts( + self, + subset=None, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + dropna: bool = True, + ): + """ + Return a Series containing counts of unique rows in the DataFrame. + + Args: + subset : label or list of labels, optional + Columns to use when counting unique combinations. + normalize : bool, default False + Return proportions rather than frequencies. + sort : bool, default True + Sort by frequencies. + ascending : bool, default False + Sort in ascending order. + dropna : bool, default True + Don’t include counts of rows that contain NA values. + + Returns: + Series + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py new file mode 100644 index 00000000000..b6ee19b46ac --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -0,0 +1,438 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/generic.py +from __future__ import annotations + +from typing import Literal, Optional + +from third_party.bigframes_vendored.pandas.core import indexing + + +class NDFrame(indexing.IndexingMixin): + """ + N-dimensional analogue of DataFrame. Store multi-dimensional in a + size-mutable, labeled data structure + """ + + # ---------------------------------------------------------------------- + # Axis + + @property + def ndim(self) -> int: + """Return an int representing the number of axes / array dimensions. + + Return 1 if Series. Otherwise return 2 if DataFrame. + """ + raise NotImplementedError("abstract method") + + @property + def size(self) -> int: + """Return an int representing the number of elements in this object. + + Return the number of rows if Series. Otherwise return the number of + rows times number of columns if DataFrame. + """ + raise NotImplementedError("abstract method") + + # ------------------------------------------------------------------------- + # Unary Methods + + def abs(self): + """Return a Series/DataFrame with absolute numeric value of each element. + + This function only applies to elements that are all numeric. + + Returns: + Series/DataFrame containing the absolute value of each element. + """ + raise NotImplementedError("abstract method") + + def astype(self, dtype): + """ + Cast a pandas object to a specified dtype ``dtype``. + + Parameters + ---------- + dtype : str, data type, Series or Mapping of column name -> data type + Use a str, numpy.dtype, pandas.ExtensionDtype or Python type to + cast entire pandas object to the same type. Alternatively, use a + mapping, e.g. {col: dtype, ...}, where col is a column label and dtype is + a numpy.dtype or Python type to cast one or more of the DataFrame's + columns to column-specific types. + + Returns + ------- + same type as caller + + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Iteration + + @property + def empty(self) -> bool: + """Indicator whether Series/DataFrame is empty. + + True if Series/DataFrame is entirely empty (no items), meaning any of the + axes are of length 0. + + Returns: + If Series/DataFrame is empty, return True, if not return False. + + Note: + If Series/DataFrame contains only NA values, it is still not + considered empty. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # I/O Methods + + def to_json( + self, + path_or_buf: str, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + *, + index: bool = True, + lines: bool = False, + ) -> str | None: + """Convert the object to a JSON string, written to GCS. + + Note NaN's and None will be converted to null and datetime objects + will be converted to UNIX timestamps. + + Args: + path_or_buf: + A destination URI of GCS files(s) to store the extracted dataframe + in format of ``gs:///``. + + If the data size is more than 1GB, you must use a wildcard to + export the data into multiple files and the size of the files + varies. + + None, file-like objects or local file paths not yet supported. + orient: + Indication of expected JSON string format. + + .. note:: + + In BigQuery DataFrame, only `orient='records'` is supported so far. + + * Series: + + - default is 'index' + - allowed values are: {{'split', 'records', 'index', 'table'}}. + + * DataFrame: + + - default is 'columns' + - allowed values are: {{'split', 'records', 'index', 'columns', + 'values', 'table'}}. + + * The format of the JSON string: + + - 'split' : dict like {{'index' -> [index], 'columns' -> [columns], + 'data' -> [values]}} + - 'records' : list like [{{column -> value}}, ... , {{column -> value}}] + - 'index' : dict like {{index -> {{column -> value}}}} + - 'columns' : dict like {{column -> {{index -> value}}}} + - 'values' : just the values array + - 'table' : dict like {{'schema': {{schema}}, 'data': {{data}}}} + + Describing the data, where data component is like ``orient='records'``. + + lines: + If 'orient' is 'records' write out line-delimited json format. Will + throw ValueError if incorrect 'orient' since others are not + list-like. + + .. note:: + + BigQuery DataFrames only supports ``lines=True`` so far. + + index: + If True, write row names (index). + + Returns: + None. String output not yet supported. + """ + raise NotImplementedError("abstract method") + + def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: + """Write object to a comma-separated values (csv) file on GCS. + + Args: + path_or_buf: + A destination URI of GCS files(s) to store the extracted dataframe + in format of ``gs:///``. + + If the data size is more than 1GB, you must use a wildcard to + export the data into multiple files and the size of the files + varies. + + None, file-like objects or local file paths not yet supported. + + index: + If True, write row names (index). + + Returns: + None. String output not yet supported. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Unsorted + + def get(self, key, default=None): + """ + Get item from object for given key (ex: DataFrame column). + + Returns default value if not found. + + Args: + key: object + + Returns: + same type as items contained in object + """ + try: + return self[key] + except (KeyError, ValueError, IndexError): + return default + + def add_prefix(self, prefix: str, axis: int | str | None = None): + """Prefix labels with string `prefix`. + + For Series, the row labels are prefixed. + For DataFrame, the column labels are prefixed. + + Args: + prefix: + The string to add before each label. + axis: + ``{{0 or 'index', 1 or 'columns', None}}``, default None. Axis + to add prefix on + + Returns: + New Series or DataFrame with updated labels. + """ + raise NotImplementedError("abstract method") + + def add_suffix(self, suffix: str, axis: int | str | None = None): + """Suffix labels with string `suffix`. + + For Series, the row labels are suffixed. + For DataFrame, the column labels are suffixed. + + Args: + suffix: + The string to add after each label. + axis: + ``{{0 or 'index', 1 or 'columns', None}}``, default None. Axis + to add suffix on + + Returns: + New Series or DataFrame with updated labels. + """ + raise NotImplementedError("abstract method") + + def head(self, n: int = 5): + """Return the first `n` rows. + + This function returns the first `n` rows for the object based + on position. It is useful for quickly testing if your object + has the right type of data in it. + + **Not yet supported** For negative values of `n`, this function returns + all rows except the last `|n|` rows, equivalent to ``df[:n]``. + + If n is larger than the number of rows, this function returns all rows. + + Args: + n: + Default 5. Number of rows to select. + + Returns: + The first `n` rows of the caller object. + """ + raise NotImplementedError("abstract method") + + def tail(self, n: int = 5): + """Return the last `n` rows. + + This function returns last `n` rows from the object based on + position. It is useful for quickly verifying data, for example, + after sorting or appending rows. + + For negative values of `n`, this function returns all rows except + the first `|n|` rows, equivalent to ``df[|n|:]``. + + If n is larger than the number of rows, this function returns all rows. + + Args: + n: int, default 5. Number of rows to select. + + Returns: + The last `n` rows of the caller object. + """ + raise NotImplementedError("abstract method") + + def sample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + *, + random_state: Optional[int] = None, + ): + """Return a random sample of items from an axis of object. + + You can use `random_state` for reproducibility. + + Args: + n: + Number of items from axis to return. Cannot be used with `frac`. + Default = 1 if `frac` = None. + frac: + Fraction of axis items to return. Cannot be used with `n`. + random_state: + Seed for random number generator. + + Returns: + A new object of same type as caller containing `n` items randomly + sampled from the caller object. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Internal Interface Methods + + @property + def dtypes(self): + """Return the dtypes in the DataFrame. + + This returns a Series with the data type of each column. + The result's index is the original DataFrame's columns. Columns + with mixed types aren't supported yet in BigQuery DataFrames. + + Returns: + A *pandas* Series with the data type of each column. + """ + raise NotImplementedError("abstract method") + + def copy(self): + """Make a copy of this object's indices and data. + + A new object will be created with a copy of the calling object's data + and indices. Modifications to the data or indices of the copy will not + be reflected in the original object. + + Returns: + Object type matches caller. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # Action Methods + + def isna(self) -> NDFrame: + """Detect missing values. + + Return a boolean same-sized object indicating if the values are NA. + NA values get mapped to True values. Everything else gets mapped to + False values. Characters such as empty strings ``''`` or + :attr:`numpy.inf` are not considered NA values. + + Returns: + Mask of bool values for each element that indicates whether an + element is an NA value. + """ + raise NotImplementedError("abstract method") + + isnull = isna + + def notna(self) -> NDFrame: + """Detect existing (non-missing) values. + + Return a boolean same-sized object indicating if the values are not NA. + Non-missing values get mapped to True. Characters such as empty + strings ``''`` or :attr:`numpy.inf` are not considered NA values. + NA values get mapped to False values. + + Returns: + Mask of bool values for each element that indicates whether an + element is not an NA value. + """ + raise NotImplementedError("abstract method") + + notnull = notna + + def shift( + self, + periods: int = 1, + ) -> NDFrame: + """Shift index by desired number of periods. + + Shifts the index without realigning the data. + + Args: + periods: + Number of periods to shift. Can be positive or negative. + + Returns: + Copy of input object, shifted. + """ + raise NotImplementedError("abstract method") + + def rank( + self, + axis=0, + method: str = "average", + numeric_only: bool = False, + na_option: str = "keep", + ascending: bool = True, + ): + """ + Compute numerical data ranks (1 through n) along axis. + + By default, equal values are assigned a rank that is the average of the + ranks of those values. + + Parameters + ---------- + method : {'average', 'min', 'max', 'first', 'dense'}, default 'average' + How to rank the group of records that have the same value (i.e. ties): + + * average: average rank of the group + * min: lowest rank in the group + * max: highest rank in the group + * first: ranks assigned in order they appear in the array + * dense: like 'min', but rank always increases by 1 between groups. + + numeric_only : bool, default False + For DataFrame objects, rank only numeric columns if set to True. + + na_option : {'keep', 'top', 'bottom'}, default 'keep' + How to rank NaN values: + + * keep: assign NaN rank to NaN values + * top: assign lowest rank to NaN values + * bottom: assign highest rank to NaN values + + ascending : bool, default True + Whether or not the elements should be ranked in ascending order. + + Returns + ------- + same type as caller + Return a Series or DataFrame with data ranks as values. + """ + raise NotImplementedError("abstract method") + + def __nonzero__(self): + raise ValueError( + f"The truth value of a {type(self).__name__} is ambiguous. " + "Use a.empty, a.bool(), a.item(), a.any() or a.all()." + ) + + __bool__ = __nonzero__ diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py new file mode 100644 index 00000000000..5b5a9f206b3 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -0,0 +1,318 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/groupby/groupby.py +""" +Provide the groupby split-apply-combine paradigm. Define the GroupBy +class providing the base-class of operations. + +The SeriesGroupBy and DataFrameGroupBy sub-class +(defined in pandas.core.groupby.generic) +expose these user-facing objects to provide specific functionality. +""" +from __future__ import annotations + + +class GroupBy: + """ + Class for grouping and aggregating relational data. + """ + + def any(self): + """ + Return True if any value in the group is truthful, else False. + + + Returns + ------- + Series or DataFrame + DataFrame or Series of boolean values, where a value is True if any element + is True within its respective group, False otherwise. + """ + raise NotImplementedError("abstract property") + + def all(self): + """ + Return True if all values in the group are truthful, else False. + + Returns + ------- + Series or DataFrame + DataFrame or Series of boolean values, where a value is True if all elements + are True within its respective group, False otherwise. + """ + raise NotImplementedError("abstract property") + + def count(self): + """ + Compute count of group, excluding missing values. + + Returns + ------- + Series or DataFrame + Count of values within each group. + """ + raise NotImplementedError("abstract property") + + def mean( + self, + numeric_only: bool = False, + ): + """ + Compute mean of groups, excluding missing values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. + + Returns + ------- + pandas.Series or pandas.DataFrame + """ + raise NotImplementedError("abstract property") + + def std( + self, + *, + numeric_only: bool = False, + ): + """ + Compute standard deviation of groups, excluding missing values. + + For multiple groupings, the result index will be a MultiIndex. + + Parameters + ---------- + numeric_only : bool, default False + Include only `float`, `int` or `boolean` data. + + Returns + ------- + Series or DataFrame + Standard deviation of values within each group. + """ + raise NotImplementedError("abstract property") + + def var( + self, + *, + numeric_only: bool = False, + ): + """ + Compute variance of groups, excluding missing values. + + For multiple groupings, the result index will be a MultiIndex. + + Parameters + ---------- + numeric_only : bool, default False + Include only `float`, `int` or `boolean` data. + + Returns + ------- + Series or DataFrame + Variance of values within each group. + """ + raise NotImplementedError("abstract property") + + def sum( + self, + numeric_only: bool = False, + min_count: int = 0, + ): + """ + Compute sum of group values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. + min_count : int, default 0 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + Returns + ------- + Series or DataFrame + Computed sum of values within each group. + """ + raise NotImplementedError("abstract property") + + def prod(self, numeric_only: bool = False, min_count: int = 0): + """ + Compute prod of group values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. + min_count : int, default 0 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + Returns + ------- + Series or DataFrame + Computed prod of values within each group. + """ + raise NotImplementedError("abstract property") + + def min( + self, + numeric_only: bool = False, + min_count: int = -1, + ): + """ + Compute min of group values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. + min_count : int, default 0 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + Returns + ------- + Series or DataFrame + Computed min of values within each group. + """ + raise NotImplementedError("abstract property") + + def max( + self, + numeric_only: bool = False, + min_count: int = -1, + ): + """ + Compute max of group values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. + min_count : int, default 0 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + Returns + ------- + Series or DataFrame + Computed max of values within each group. + """ + raise NotImplementedError("abstract property") + + def cumcount(self, ascending: bool = True): + """ + Number each item in each group from 0 to the length of that group - 1. + + Parameters + ---------- + ascending : bool, default True + If False, number in reverse, from length of group - 1 to 0. + + Returns + ------- + Series + Sequence number of each element within each group. + """ + raise NotImplementedError("abstract property") + + def cumprod(self, *args, **kwargs): + """ + Cumulative product for each group. + + Returns + ------- + Series or DataFrame + """ + raise NotImplementedError("abstract property") + + def cumsum(self, *args, **kwargs): + """ + Cumulative sum for each group. + + Returns + ------- + Series or DataFrame + """ + raise NotImplementedError("abstract property") + + def cummin(self, *args, numeric_only: bool = False, **kwargs): + """ + Cumulative min for each group. + + Returns + ------- + Series or DataFrame + """ + raise NotImplementedError("abstract property") + + def cummax(self, *args, numeric_only: bool = False, **kwargs): + """ + Cumulative max for each group. + + Returns + ------- + Series or DataFrame + """ + raise NotImplementedError("abstract property") + + def diff(self): + """ + First discrete difference of element. + Calculates the difference of each element compared with another + element in the group (default is element in previous row). + + Returns + ------- + Series or DataFrame + First differences. + """ + raise NotImplementedError("abstract property") + + def shift(self, periods: int = 1): + """ + Shift each group by periods observations. + If freq is passed, the index will be increased using the periods and the freq. + + Parameters + ---------- + periods : int, default 1 + Number of periods to shift. + Returns + ------- + Series or DataFrame + Object shifted within each group. + """ + raise NotImplementedError("abstract property") + + def rolling(self, *args, **kwargs): + """ + Return a rolling grouper, providing rolling functionality per group. + + Args: + min_periods : int, default None + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. + + For a window that is specified by an offset, + ``min_periods`` will default to 1. + + For a window that is specified by an integer, ``min_periods`` will default + to the size of the window. + + Returns: + Return a new grouper with our rolling appended. + """ + raise NotImplementedError("abstract property") + + def expanding(self, *args, **kwargs): + """ + Return an expanding grouper, providing expanding + functionality per group. + """ + raise NotImplementedError("abstract property") + + +class SeriesGroupBy(GroupBy): + pass + + +class DataFrameGroupBy(GroupBy): + pass diff --git a/third_party/bigframes_vendored/pandas/core/indexes/__init__.py b/third_party/bigframes_vendored/pandas/core/indexes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/indexes/accessor.py b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py new file mode 100644 index 00000000000..5d431a8a949 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py @@ -0,0 +1,94 @@ +class DatetimeProperties: + """ + Accessor object for datetimelike properties of the Series values. + """ + + @property + def day(self): + """The day of the datetime.""" + + raise NotImplementedError("abstract method") + + @property + def dayofweek(self): + """The day of the week with Monday=0, Sunday=6. + + Return the day of the week. It is assumed the week starts on + Monday, which is denoted by 0 and ends on Sunday which is denoted + by 6. This method is available on both Series with datetime + values (using the `dt` accessor) or DatetimeIndex. + + Returns: + Series or Index + Containing integers indicating the day number. + """ + + raise NotImplementedError("abstract method") + + @property + def date(self): + """Returns numpy array of python :class:`datetime.date` objects. + + Namely, the date part of Timestamps without time and + timezone information. + + warning: + This method returns a Series whereas pandas returns + a numpy array. + """ + + raise NotImplementedError("abstract method") + + @property + def hour(self): + """The hours of the datetime.""" + + raise NotImplementedError("abstract method") + + @property + def minute(self): + """The minutes of the datetime.""" + + raise NotImplementedError("abstract method") + + @property + def month(self): + """The month as January=1, December=12.""" + + raise NotImplementedError("abstract method") + + @property + def second(self): + """The seconds of the datetime.""" + + raise NotImplementedError("abstract method") + + @property + def time(self): + """Returns numpy array of :class:`datetime.time` objects. + + The time part of the Timestamps. + + warning: + This method returns a Series whereas pandas returns + a numpy array. + """ + + raise NotImplementedError("abstract method") + + @property + def quarter(self): + """The quarter of the date. + + warning: + This method returns a Series whereas pandas returns + a numpy array. + """ + + raise NotImplementedError("abstract method") + + @property + def year(self): + """The year of the datetime.""" + + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py new file mode 100644 index 00000000000..bdd05a01007 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -0,0 +1,20 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexes/base.py + + +class Index: + """Immutable sequence used for indexing and alignment. + + The basic object storing axis labels for all objects. + """ + + @property + def name(self): + """Return Index name.""" + raise NotImplementedError("abstract method") + + @property + def shape(self): + """ + Return a tuple of the shape of the underlying data. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/indexing.py b/third_party/bigframes_vendored/pandas/core/indexing.py new file mode 100644 index 00000000000..87fb34b32a3 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/indexing.py @@ -0,0 +1,72 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexing.py + + +class IndexingMixin: + """ + Mixin for adding .loc/.iloc/.at/.iat to Dataframes and Series. + """ + + @property + def iloc(self): + """Purely integer-location based indexing for selection by position. + + ``.iloc[]`` is primarily integer position based (from ``0`` to + ``length-1`` of the axis), but may also be used with a boolean + array. + + Allowed inputs are: + + - **Not supported yet** An integer, e.g. ``5``. + - **Not supported yet** A list or array of integers, e.g. ``[4, 3, 0]``. + - A slice object with ints, e.g. ``1:7``. + - **Not supported yet** A boolean array. + - **Not supported yet** A ``callable`` function with one argument (the + calling Series or DataFrame) and that returns valid output for + indexing (one of the above). This is useful in method chains, when you + don't have a reference to the calling object, but would like to base + your selection on some value. + - **Not supported yet** A tuple of row and column indexes. The tuple + elements consist of one of the above inputs, e.g. ``(0, 1)``. + + ``.iloc`` will raise ``IndexError`` if a requested indexer is + out-of-bounds, except *slice* indexers which allow out-of-bounds + indexing (this conforms with python/numpy *slice* semantics). + """ + raise NotImplementedError("abstract methdod") + + @property + def loc(self): + """Access a group of rows and columns by label(s) or a boolean array. + + ``.loc[]`` is primarily label based, but may also be used with a + boolean array. + + Allowed inputs are: + + - **Not supported yet** A single label, e.g. ``5`` or ``'a'``, (note + that ``5`` is interpreted as a *label* of the index, and **never** as + an integer position along the index). + - **Not supported yet** A list or array of labels, e.g. ``['a', 'b', 'c']``. + - **Not supported yet** A slice object with labels, e.g. ``'a':'f'``. + + .. warning:: Note that contrary to usual python slices, **both** the + start and the stop are included + + - **Not supported yet** A boolean array of the same length as the axis being sliced, + e.g. ``[True, False, True]``. + - An alignable boolean Series. The index of the key will be aligned before + masking. + - **Not supported yet** An alignable Index. The Index of the returned + selection will be the input. + - **Not supported yet** A ``callable`` function with one argument (the + calling Series or DataFrame) + and that returns valid output for indexing (one of the above) + + Raises: + KeyError: If any items are not found. + IndexingError: + If an indexed key is passed and its index is unalignable to the + frame index. + + """ + raise NotImplementedError("abstract methdod") diff --git a/third_party/bigframes_vendored/pandas/core/reshape/concat.py b/third_party/bigframes_vendored/pandas/core/reshape/concat.py new file mode 100644 index 00000000000..98b43c7dfba --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/reshape/concat.py @@ -0,0 +1,135 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/concat.py +""" +Concat routines. +""" +from __future__ import annotations + + +def concat( + objs, + *, + join: str = "outer", + ignore_index: bool = False, +): + """ + Concatenate BigQuery DataFrames objects along a particular axis. + + Allows optional set logic along the other axes. + + Can also add a layer of hierarchical indexing on the concatenation axis, + which may be useful if the labels are the same (or overlapping) on + the passed axis number. + + Parameters + ---------- + objs: + Objects to concatenate. Any None objects will be dropped silently unless + they are all None in which case a ValueError will be raised. + join: {'inner', 'outer'}, default 'outer' + How to handle indexes on other axis (or axes). + ignore_index : bool, default False + If True, do not use the index values along the concatenation axis. The + resulting axis will be labeled 0, ..., n - 1. This is useful if you are + concatenating objects where the concatenation axis does not have + meaningful indexing information. Note the index values on the other + axes are still respected in the join. + + Returns + ------- + object, type of objs + When concatenating all ``Series`` along the index (axis=0), a + ``Series`` is returned. When ``objs`` contains at least one + ``DataFrame``, a ``DataFrame`` is returned. + + Notes + ----- + It is not recommended to build DataFrames by adding single rows in a + for loop. Build a list of rows and make a DataFrame in a single concat. + + Examples + -------- + Combine two ``Series``. + + >>> import bigframes.pandas as pd + >>> pd.options.display.progress_bar = None + >>> s1 = pd.Series(['a', 'b']) + >>> s2 = pd.Series(['c', 'd']) + >>> pd.concat([s1, s2]) + 0 a + 1 b + 0 c + 1 d + dtype: string + + Clear the existing index and reset it in the result + by setting the ``ignore_index`` option to ``True``. + + >>> pd.concat([s1, s2], ignore_index=True) + 0 a + 1 b + 2 c + 3 d + dtype: string + + Combine two ``DataFrame`` objects with identical columns. + + >>> df1 = pd.DataFrame([['a', 1], ['b', 2]], + ... columns=['letter', 'number']) + >>> df1 + letter number + 0 a 1 + 1 b 2 + + [2 rows x 2 columns] + >>> df2 = pd.DataFrame([['c', 3], ['d', 4]], + ... columns=['letter', 'number']) + >>> df2 + letter number + 0 c 3 + 1 d 4 + + [2 rows x 2 columns] + >>> pd.concat([df1, df2]) + letter number + 0 a 1 + 1 b 2 + 0 c 3 + 1 d 4 + + [4 rows x 2 columns] + + Combine ``DataFrame`` objects with overlapping columns + and return everything. Columns outside the intersection will + be filled with ``NaN`` values. + + >>> df3 = pd.DataFrame([['c', 3, 'cat'], ['d', 4, 'dog']], + ... columns=['letter', 'number', 'animal']) + >>> df3 + letter number animal + 0 c 3 cat + 1 d 4 dog + + [2 rows x 3 columns] + >>> pd.concat([df1, df3]) + letter number animal + 0 a 1 + 1 b 2 + 0 c 3 cat + 1 d 4 dog + + [4 rows x 3 columns] + + Combine ``DataFrame`` objects with overlapping columns + and return only those that are shared by passing ``inner`` to + the ``join`` keyword argument. + + >>> pd.concat([df1, df3], join="inner") + letter number + 0 a 1 + 1 b 2 + 0 c 3 + 1 d 4 + + [4 rows x 2 columns] + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py new file mode 100644 index 00000000000..a254fcf842d --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -0,0 +1,1553 @@ +""" +Data structure for 1-dimensional cross-sectional and time series data +""" +from __future__ import annotations + +from typing import Hashable, IO, Literal, Mapping, Sequence, TYPE_CHECKING + +import numpy as np +from pandas._libs import lib +from pandas._typing import Axis, FilePath, NaPosition, WriteBuffer +import pandas.io.formats.format as fmt + +from third_party.bigframes_vendored.pandas.core.generic import NDFrame + +if TYPE_CHECKING: + from third_party.bigframes_vendored.pandas.core.frame import DataFrame + from third_party.bigframes_vendored.pandas.core.groupby import SeriesGroupBy + + +class Series(NDFrame): # type: ignore[misc] + @property + def dt(self): + """ + Accessor object for datetimelike properties of the Series values. + """ + raise NotImplementedError("abstract property") + + @property + def index(self): + """The index (axis labels) of the Series.""" + raise NotImplementedError("abstract property") + + @property + def shape(self): + """Return a tuple of the shape of the underlying data.""" + raise NotImplementedError("abstract property") + + @property + def dtype(self): + """ + Return the dtype object of the underlying data. + """ + raise NotImplementedError("abstract property") + + @property + def dtypes(self): + """ + Return the dtype object of the underlying data. + """ + raise NotImplementedError("abstract property") + + @property + def name(self) -> Hashable: + """ + Return the name of the Series. + + The name of a Series becomes its index or column name if it is used + to form a DataFrame. It is also used whenever displaying the Series + using the interpreter. + + Returns + ------- + label (hashable object) + The name of the Series, also the column name if part of a DataFrame. + """ + raise NotImplementedError("abstract property") + + def reset_index( + self, + *, + drop: bool = False, + name=lib.no_default, + ) -> DataFrame | Series | None: + """ + Generate a new DataFrame or Series with the index reset. + + This is useful when the index needs to be treated as a column, or + when the index is meaningless and needs to be reset to the default + before another operation. + + Parameters + ---------- + drop : bool, default False + Just reset the index, without inserting it as a column in + the new DataFrame. + name : object, optional + The name to use for the column containing the original Series + values. Uses ``self.name`` by default. This argument is ignored + when `drop` is True. + + Returns + ------- + Series or DataFrame or None + When `drop` is False (the default), a DataFrame is returned. + The newly created columns will come first in the DataFrame, + followed by the original Series values. + When `drop` is True, a `Series` is returned. + In either case, if ``inplace=True``, no value is returned. + + """ + raise NotImplementedError("abstract method") + + def __repr__(self) -> str: + """ + Return a string representation for a particular Series. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # IO methods (to / from other formats) + + def to_string( + self, + buf: FilePath | WriteBuffer[str] | None = None, + na_rep: str = "NaN", + float_format: str | None = None, + header: bool = True, + index: bool = True, + length: bool = False, + dtype: bool = False, + name: bool = False, + max_rows: int | None = None, + min_rows: int | None = None, + ) -> str | None: + """ + Render a string representation of the Series. + + Parameters + ---------- + buf : StringIO-like, optional + Buffer to write to. + na_rep : str, optional + String representation of NaN to use, default 'NaN'. + float_format : one-parameter function, optional + Formatter function to apply to columns' elements if they are + floats, default None. + header : bool, default True + Add the Series header (index name). + index : bool, optional + Add index (row) labels, default True. + length : bool, default False + Add the Series length. + dtype : bool, default False + Add the Series dtype. + name : bool, default False + Add the Series name if not None. + max_rows : int, optional + Maximum number of rows to show before truncating. If None, show + all. + min_rows : int, optional + The number of rows to display in a truncated repr (when number + of rows is above `max_rows`). + + Returns + ------- + str or None + String representation of Series if ``buf=None``, otherwise None. + """ + formatter = fmt.SeriesFormatter( + self, + name=name, + length=length, + header=header, + index=index, + dtype=dtype, + na_rep=na_rep, + float_format=float_format, + min_rows=min_rows, + max_rows=max_rows, + ) + result = formatter.to_string() + + # catch contract violations + raise NotImplementedError("abstract method") + + def to_markdown( + self, + buf: IO[str] | None = None, + mode: str = "wt", + index: bool = True, + **kwargs, + ) -> str | None: + """ + Print {klass} in Markdown-friendly format. + + Parameters + ---------- + buf : str, Path or StringIO-like, optional, default None + Buffer to write to. If None, the output is returned as a string. + mode : str, optional + Mode in which file is opened, "wt" by default. + index : bool, optional, default True + Add index (row) labels. + Returns + ------- + str + {klass} in Markdown-friendly format. + """ + raise NotImplementedError("abstract method") + + def to_dict(self, into: type[dict] = dict) -> Mapping: + """ + Convert Series to {label -> value} dict or dict-like object. + + Parameters + ---------- + into : class, default dict + The collections.abc.Mapping subclass to use as the return + object. Can be the actual class or an empty + instance of the mapping type you want. If you want a + collections.defaultdict, you must pass it initialized. + + Returns + ------- + collections.abc.Mapping + Key-value representation of Series. + """ + raise NotImplementedError("abstract method") + + def to_frame(self) -> DataFrame: + """ + Convert Series to DataFrame. + + Returns + ------- + DataFrame + DataFrame representation of Series. + """ + raise NotImplementedError("abstract method") + + def to_excel(self, excel_writer, sheet_name): + """ + Write Series to an Excel sheet. + + To write a single Series to an Excel .xlsx file it is only necessary to + specify a target file name. To write to multiple sheets it is necessary to + create an `ExcelWriter` object with a target file name, and specify a sheet + in the file to write to. + + Multiple sheets may be written to by specifying unique `sheet_name`. + With all data written to the file it is necessary to save the changes. + Note that creating an `ExcelWriter` object with a file name that already + exists will result in the contents of the existing file being erased. + + Parameters + ---------- + excel_writer : path-like, file-like, or ExcelWriter object + File path or existing ExcelWriter. + sheet_name : str, default 'Sheet1' + Name of sheet which will contain DataFrame. + """ + raise NotImplementedError("abstract method") + + def to_latex(self, buf=None, columns=None, header=True, index=True, **kwargs): + """ + Render object to a LaTeX tabular, longtable, or nested table. + + Parameters + ---------- + buf : str, Path or StringIO-like, optional, default None + Buffer to write to. If None, the output is returned as a string. + columns : list of label, optional + The subset of columns to write. Writes all columns by default. + header : bool or list of str, default True + Write out the column names. If a list of strings is given, + it is assumed to be aliases for the column names. + index : bool, default True + Write row names (index). + + + Returns + ------- + str or None + If buf is None, returns the result as a string. Otherwise returns None. + """ + raise NotImplementedError("abstract method") + + def tolist(self) -> list: + """ + Return a list of the values. + + These are each a scalar type, which is a Python scalar + (for str, int, float) or a pandas scalar + (for Timestamp/Timedelta/Interval/Period) + + Returns + ------- + list + """ + raise NotImplementedError("abstract method") + + to_list = tolist + + def to_numpy(self, dtype, copy=False, na_value=None): + """ + A NumPy ndarray representing the values in this Series or Index. + + Parameters + ---------- + dtype : str or numpy.dtype, optional + The dtype to pass to :meth:`numpy.asarray`. + copy : bool, default False + Whether to ensure that the returned value is not a view on + another array. Note that ``copy=False`` does not *ensure* that + ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that + a copy is made, even if not strictly necessary. + na_value : Any, optional + The value to use for missing values. The default value depends + on `dtype` and the type of the array. + **kwargs + Additional keywords passed through to the ``to_numpy`` method + of the underlying array (for extension arrays). + + Returns + ------- + numpy.ndarray + """ + raise NotImplementedError("abstract method") + + def to_pickle(self, path, **kwargs): + """ + Pickle (serialize) object to file. + + Parameters + ---------- + path : str, path object, or file-like object + String, path object (implementing ``os.PathLike[str]``), or file-like + object implementing a binary ``write()`` function. File path where + the pickled object will be stored. + """ + raise NotImplementedError("abstract method") + + def to_xarray(self): + """ + Return an xarray object from the pandas object. + + Returns + ------- + xarray.DataArray or xarray.Dataset + Data in the pandas structure converted to Dataset if the object is + a DataFrame, or a DataArray if the object is a Series. + """ + raise NotImplementedError("abstract method") + + def to_json( + self, + path_or_buf=None, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + **kwarg, + ) -> str | None: + """ + Convert the object to a JSON string. + + Note NaN's and None will be converted to null and datetime objects + will be converted to UNIX timestamps. + + Parameters + ---------- + path_or_buf : str, path object, file-like object, or None, default None + String, path object (implementing os.PathLike[str]), or file-like + object implementing a write() function. If None, the result is + returned as a string. + orient: + Indication of expected JSON string format. + + * Series: + + - default is 'index' + - allowed values are: {{'split', 'records', 'index', 'table'}}. + + * DataFrame: + + - default is 'columns' + - allowed values are: {{'split', 'records', 'index', 'columns', + 'values', 'table'}}. + + * The format of the JSON string: + + - 'split' : dict like {{'index' -> [index], 'columns' -> [columns], + 'data' -> [values]}} + - 'records' : list like [{{column -> value}}, ... , {{column -> value}}] + - 'index' : dict like {{index -> {{column -> value}}}} + - 'columns' : dict like {{column -> {{index -> value}}}} + - 'values' : just the values array + - 'table' : dict like {{'schema': {{schema}}, 'data': {{data}}}} + + Describing the data, where data component is like ``orient='records'``. + + Returns + ------- + None or str + If path_or_buf is None, returns the resulting json format as a + string. Otherwise returns None. + """ + raise NotImplementedError("abstract method") + + def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: + """ + Write object to a comma-separated values (csv) file. + + Parameters + ---------- + path_or_buf : str, path object, file-like object, or None, default None + String, path object (implementing os.PathLike[str]), or file-like + object implementing a write() function. If None, the result is + returned as a string. If a non-binary file object is passed, it should + be opened with `newline=''`, disabling universal newlines. If a binary + file object is passed, `mode` might need to contain a `'b'`. + + Returns + ------- + None or str + If path_or_buf is None, returns the resulting csv format as a + string. Otherwise returns None. + """ + raise NotImplementedError("abstract method") + + def count(self): + """ + Return number of non-NA/null observations in the Series. + + Returns + ------- + int or Series (if level specified) + Number of non-null values in the Series. + """ + raise NotImplementedError("abstract method") + + def nunique(self) -> int: + """ + Return number of unique elements in the object. + + Excludes NA values by default. + + Returns + ------- + int + """ + raise NotImplementedError("abstract method") + + def mode(self) -> Series: + """ + Return the mode(s) of the Series. + + The mode is the value that appears most often. There can be multiple modes. + + Always returns Series even if only one value is returned. + + Returns + ------- + Series + Modes of the Series in sorted order. + """ + raise NotImplementedError("abstract method") + + def drop_duplicates( + self, + *, + keep="first", + ) -> Series | None: + """ + Return Series with duplicate values removed. + + Args: + keep : {'first', 'last', ``False``}, default 'first' + Method to handle dropping duplicates: + + - 'first' : Drop duplicates except for the first occurrence. + - 'last' : Drop duplicates except for the last occurrence. + - ``False`` : Drop all duplicates. + + Returns: + Series with duplicates dropped or None if ``inplace=True``. + """ + raise NotImplementedError("abstract method") + + def duplicated(self, keep="first") -> Series: + """ + Indicate duplicate Series values. + + Duplicated values are indicated as ``True`` values in the resulting + Series. Either all duplicates, all except the first or all except the + last occurrence of duplicates can be indicated. + + Args: + keep : {'first', 'last', False}, default 'first' + Method to handle dropping duplicates: + + - 'first' : Mark duplicates as ``True`` except for the first + occurrence. + - 'last' : Mark duplicates as ``True`` except for the last + occurrence. + - ``False`` : Mark all duplicates as ``True``. + + Returns: + Series indicating whether each value has occurred in the + preceding values. + """ + raise NotImplementedError("abstract method") + + def round(self, decimals: int = 0) -> Series: + """ + Round each value in a Series to the given number of decimals. + + Parameters + ---------- + decimals : int, default 0 + Number of decimal places to round to. If decimals is negative, + it specifies the number of positions to the left of the decimal point. + + Returns + ------- + Series + Rounded values of the Series. + """ + raise NotImplementedError("abstract method") + + def diff(self) -> Series: + """ + First discrete difference of element. + + Calculates the difference of a {klass} element compared with another + element in the {klass} (default is element in previous row). + + Returns + ------- + {klass} + First differences of the Series. + """ + raise NotImplementedError("abstract method") + + def dot(self, other) -> Series | np.ndarray: + """ + Compute the dot product between the Series and the columns of other. + + This method computes the dot product between the Series and another + one, or the Series and each columns of a DataFrame, or the Series and + each columns of an array. + + It can also be called using `self @ other` in Python >= 3.5. + + Parameters + ---------- + other : Series, DataFrame or array-like + The other object to compute the dot product with its columns. + + Returns + ------- + scalar, Series or numpy.ndarray + Return the dot product of the Series and other if other is a + Series, the Series of the dot product of Series and each rows of + other if other is a DataFrame or a numpy.ndarray between the Series + and each columns of the numpy array. + + Notes + ----- + The Series and other has to share the same index if other is a Series + or a DataFrame. + + """ + raise NotImplementedError("abstract method") + + def __matmul__(self, other): + """ + Matrix multiplication using binary `@` operator in Python>=3.5. + """ + raise NotImplementedError("abstract method") + + def __rmatmul__(self, other): + """ + Matrix multiplication using binary `@` operator in Python>=3.5. + """ + raise NotImplementedError("abstract method") + + def sort_values( + self, + *, + axis: Axis = 0, + ascending: bool | int | Sequence[bool] | Sequence[int] = True, + na_position: str = "last", + ) -> Series | None: + """ + Sort by the values. + + Sort a Series in ascending or descending order by some + criterion. + + Parameters + ---------- + axis : {0 or 'index'} + Unused. Parameter needed for compatibility with DataFrame. + ascending : bool or list of bools, default True + If True, sort values in ascending order, otherwise descending. + na_position : {'first' or 'last'}, default 'last' + Argument 'first' puts NaNs at the beginning, 'last' puts NaNs at + the end. + + Returns + ------- + Series or None + Series ordered by values or None if ``inplace=True``. + """ + raise NotImplementedError("abstract method") + + def sort_index( + self, + *, + axis: Axis = 0, + ascending: bool | Sequence[bool] = True, + na_position: NaPosition = "last", + ) -> Series | None: + """ + Sort Series by index labels. + + Returns a new Series sorted by label if `inplace` argument is + ``False``, otherwise updates the original series and returns None. + + Parameters + ---------- + axis : {0 or 'index'} + Unused. Parameter needed for compatibility with DataFrame. + ascending : bool or list-like of bools, default True + Sort ascending vs. descending. When the index is a MultiIndex the + sort direction can be controlled for each level individually. + na_position : {'first', 'last'}, default 'last' + If 'first' puts NaNs at the beginning, 'last' puts NaNs at the end. + Not implemented for MultiIndex. + + Returns + ------- + Series or None + The original Series sorted by the labels or None if ``inplace=True``. + + """ + + raise NotImplementedError("abstract method") + + def nlargest( + self, n: int = 5, keep: Literal["first", "last", "all"] = "first" + ) -> Series: + """ + Return the largest `n` elements. + + Parameters + ---------- + n : int, default 5 + Return this many descending sorted values. + keep : {'first', 'last', 'all'}, default 'first' + When there are duplicate values that cannot all fit in a + Series of `n` elements: + + - ``first`` : return the first `n` occurrences in order + of appearance. + - ``last`` : return the last `n` occurrences in reverse + order of appearance. + - ``all`` : keep all occurrences. This can result in a Series of + size larger than `n`. + + Returns + ------- + Series + The `n` largest values in the Series, sorted in decreasing order. + """ + raise NotImplementedError("abstract method") + + def nsmallest(self, n: int = 5, keep: str = "first") -> Series: + """ + Return the smallest `n` elements. + + Parameters + ---------- + n : int, default 5 + Return this many ascending sorted values. + keep : {'first', 'last', 'all'}, default 'first' + When there are duplicate values that cannot all fit in a + Series of `n` elements: + + - ``first`` : return the first `n` occurrences in order + of appearance. + - ``last`` : return the last `n` occurrences in reverse + order of appearance. + - ``all`` : keep all occurrences. This can result in a Series of + size larger than `n`. + + Returns + ------- + Series + The `n` smallest values in the Series, sorted in increasing order. + """ + raise NotImplementedError("abstract method") + + # ---------------------------------------------------------------------- + # function application + + def apply( + self, + func, + ) -> DataFrame | Series: + """ + Invoke function on values of Series. + + Can be ufunc (a NumPy function that applies to the entire Series) + or a Python function that only works on single values. + + Parameters + ---------- + func : function + Python function or NumPy ufunc to apply. + Returns + ------- + Series or DataFrame + If func returns a Series object the result will be a DataFrame. + """ + raise NotImplementedError("abstract method") + + def groupby( + self, + by=None, + axis: Axis = 0, + level=None, + as_index: bool = True, + *, + dropna: bool = True, + ) -> SeriesGroupBy: + """Group Series using a mapper or by a Series of columns. + + A groupby operation involves some combination of splitting the + object, applying a function, and combining the results. This can be + used to group large amounts of data and compute operations on these + groups. + + Parameters + ---------- + by : mapping, function, label, pd.Grouper or list of such + Used to determine the groups for the groupby. + If ``by`` is a function, it's called on each value of the object's + index. If a dict or Series is passed, the Series or dict VALUES + will be used to determine the groups (the Series' values are first + aligned; see ``.align()`` method). If a list or ndarray of length + equal to the selected axis is passed (see the `groupby user guide + `_), + the values are used as-is to determine the groups. A label or list + of labels may be passed to group by the columns in ``self``. + Notice that a tuple is interpreted as a (single) key. + axis : {0 or 'index', 1 or 'columns'}, default 0 + Split along rows (0) or columns (1). For `Series` this parameter + is unused and defaults to 0. + level : int, level name, or sequence of such, default None + If the axis is a MultiIndex (hierarchical), group by a particular + level or levels. Do not specify both ``by`` and ``level``. + as_index : bool, default True + Return object with group labels as the + index. Only relevant for DataFrame input. as_index=False is + effectively "SQL-style" grouped output. This argument has no effect + on filtrations (see the `filtrations in the user guide + `_), + such as ``head()``, ``tail()``, ``nth()`` and in transformations + (see the `transformations in the user guide + `_). + dropna : bool, default True + If True, and if group keys contain NA values, NA values together + with row/column will be dropped. + If False, NA values will also be treated as the key in groups. + + Returns + ------- + SeriesGroupBy + Returns a groupby object that contains information about the groups. + """ + raise NotImplementedError("abstract method") + + def drop( + self, + labels=None, + ) -> Series | None: + """ + Return Series with specified index labels removed. + + Remove elements of a Series based on specifying the index labels. + When using a multi-index, labels on different levels can be removed + by specifying the level. + + Parameters + ---------- + labels : single label or list-like + Index labels to drop. + + Returns + ------- + Series or None + Series with specified index labels removed or None if ``inplace=True``. + + Raises + ------ + KeyError + If none of the labels are found in the index. + """ + raise NotImplementedError("abstract method") + + def reorder_levels(self, order: Sequence) -> Series: + """ + Rearrange index levels using input order. + + May not drop or duplicate levels. + + Args: + order: list of int representing new level order + Reference level by number or key. + + Returns: + type of caller (new object) + """ + raise NotImplementedError("abstract method") + + def droplevel(self, level): + """ + Return Series with requested index / column level(s) removed. + + Args: + level: int, str, or list-like + If a string is given, must be the name of a level + If list-like, elements must be names or positional indexes + of levels. + + Returns: + Series with requested index / column level(s) removed. + """ + raise NotImplementedError("abstract method") + + def fillna( + self, + value=None, + ) -> Series | None: + raise NotImplementedError("abstract method") + + def between( + self, + left, + right, + inclusive: Literal["both", "neither", "left", "right"] = "both", + ) -> Series: + """ + Return boolean Series equivalent to left <= series <= right. + + This function returns a boolean vector containing `True` wherever the + corresponding Series element is between the boundary values `left` and + `right`. NA values are treated as `False`. + + Parameters + ---------- + left : scalar or list-like + Left boundary. + right : scalar or list-like + Right boundary. + inclusive : {"both", "neither", "left", "right"} + Include boundaries. Whether to set each bound as closed or open. + + Returns + ------- + Series + Series representing whether each element is between left and + right (inclusive). + + """ + raise NotImplementedError("abstract method") + + def cumsum(self): + """ + Return cumulative sum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative + sum. + + Parameters + ---------- + axis : {{0 or 'index', 1 or 'columns'}}, default 0 + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. + + Returns + ------- + scalar or Series + Return cumulative sum of scalar or Series. + """ + raise NotImplementedError("abstract method") + + def cummax(self): + """ + Return cumulative maximum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative + maximum. + + Parameters + ---------- + axis : {{0 or 'index', 1 or 'columns'}}, default 0 + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. + + Returns + ------- + scalar or Series + Return cumulative maximum of scalar or Series. + """ + raise NotImplementedError("abstract method") + + def cummin(self): + """ + Return cumulative minimum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative + minimum. + + Parameters + ---------- + axis : {{0 or 'index', 1 or 'columns'}}, default 0 + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. + skipna : bool, default True + Exclude NA/null values. If an entire row/column is NA, the result + will be NA. + *args, **kwargs + Additional keywords have no effect but might be accepted for + compatibility with NumPy. + + Returns + ------- + scalar or Series + Return cumulative minimum of scalar or Series. + """ + raise NotImplementedError("abstract method") + + def eq(self, other) -> Series: + """Return Equal of series and other, element-wise (binary operator eq). + + Equivalent to ``other == series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def ne(self, other) -> Series: + """Return Not equal of series and other, element-wise (binary operator ne). + + Equivalent to ``other != series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the comparison. + + """ + raise NotImplementedError("abstract method") + + def le(self, other) -> Series: + """Get 'less than or equal to' of series and other, element-wise (binary operator `<=`). + + Equivalent to ``series <= other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the comparison. + + """ + raise NotImplementedError("abstract method") + + def lt(self, other) -> Series: + """Get 'less than' of series and other, element-wise (binary operator `<`). + + Equivalent to ``series < other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the comparison. + + """ + raise NotImplementedError("abstract method") + + def ge(self, other) -> Series: + """Get 'greater than or equal to' of series and other, element-wise (binary operator `>=`). + + Equivalent to ``series >= other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the comparison. + + """ + raise NotImplementedError("abstract method") + + def gt(self, other) -> Series: + """Get 'less than or equal to' of series and other, element-wise (binary operator `<=`). + + Equivalent to ``series <= other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def add(self, other) -> Series: + """Return Addition of series and other, element-wise (binary operator add). + + Equivalent to ``series + other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def sub( + self, + other, + ) -> Series: + """Return Subtraction of series and other, element-wise (binary operator sub). + + Equivalent to ``series - other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def rsub(self, other) -> Series: + """Return Subtraction of series and other, element-wise (binary operator rsub). + + Equivalent to ``other - series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def mul(self, other) -> Series: + """Return Multiplication of series and other, element-wise (binary operator mul). + + Equivalent to ``other * series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def truediv(self, other) -> Series: + """Return Floating division of series and other, element-wise (binary operator truediv). + + Equivalent to ``series / other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def rtruediv(self, other) -> Series: + """Return Floating division of series and other, element-wise (binary operator rtruediv). + + Equivalent to ``other / series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def floordiv(self, other) -> Series: + """Return Integer division of series and other, element-wise (binary operator floordiv). + + Equivalent to ``series // other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def rfloordiv(self, other) -> Series: + """Return Integer division of series and other, element-wise (binary operator rfloordiv). + + Equivalent to ``other // series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def mod(self, other) -> Series: + """Return Modulo of series and other, element-wise (binary operator mod). + + Equivalent to ``series % other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def rmod(self, other) -> Series: + """Get Modulo of series and other, element-wise (binary operator `rmod`). + + Equivalent to ``other % series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other: Series, or scalar value + + Returns: + Series. The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def all( + self, + ): + """ + Return whether all elements are True, potentially over an axis. + + Returns True unless there at least one element within a series or along a + Dataframe axis that is False or equivalent (e.g. zero or empty). + + + Returns + ------- + scalar or Series + If level is specified, then, Series is returned; otherwise, scalar + is returned. + """ + raise NotImplementedError("abstract method") + + def any( + self, + ): + """ + Return whether any element is True, potentially over an axis. + + Returns False unless there is at least one element within a series or along + a Dataframe axis that is True or equivalent (e.g. non-zero or non-empty). + + + Returns + ------- + scalar or Series + If level is specified, then, Series is returned; otherwise, scalar + is returned. + """ + raise NotImplementedError("abstract method") + + def max( + self, + ): + """ + Return the maximum of the values over the requested axis. + + If you want the index of the maximum, use ``idxmax``. This is the equivalent + of the ``numpy.ndarray`` method ``argmax``. + + + Returns + ------- + scalar or scalar + """ + raise NotImplementedError("abstract method") + + def min( + self, + ): + """ + Return the maximum of the values over the requested axis. + + If you want the index of the minimum, use ``idxmin``. This is the equivalent + of the ``numpy.ndarray`` method ``argmin``. + + + Returns + ------- + scalar or scalar + """ + raise NotImplementedError("abstract method") + + def std( + self, + ): + """ + Return sample standard deviation over requested axis. + + Normalized by N-1 by default. This can be changed using the ddof argument. + + + Returns + ------- + scalar or Series (if level specified) + """ + raise NotImplementedError("abstract method") + + def var( + self, + ): + """ + Return unbiased variance over requested axis. + + Normalized by N-1 by default. This can be changed using the ddof argument. + + + Returns + ------- + scalar or Series (if level specified) + """ + raise NotImplementedError("abstract method") + + def sum(self): + """Return the sum of the values over the requested axis. + + This is equivalent to the method ``numpy.sum``. + + Returns: + scalar + """ + raise NotImplementedError("abstract method") + + def mean(self): + """Return the mean of the values over the requested axis. + + Returns: + scalar + """ + raise NotImplementedError("abstract method") + + def prod(self): + """Return the product of the values over the requested axis. + + Returns: + scalar + """ + raise NotImplementedError("abstract method") + + def kurt(self): + """Return unbiased kurtosis over requested axis. + + Kurtosis obtained using Fisher’s definition of kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + + Returns + ------- + scalar or scalar + """ + raise NotImplementedError("abstract method") + + def where(self, cond, other): + """Replace values where the condition is False. + + Parameters + ---------- + cond: bool Series/DataFrame, array-like, or callable + Where cond is True, keep the original value. Where False, replace + with corresponding value from other. If cond is callable, it is + computed on the Series/DataFrame and should return boolean + Series/DataFrame or array. The callable must not change input + Series/DataFrame (though pandas doesn’t check it). + other: scalar, Series/DataFrame, or callable + Entries where cond is False are replaced with corresponding value + from other. If other is callable, it is computed on the + Series/DataFrame and should return scalar or Series/DataFrame. + The callable must not change input Series/DataFrame (though pandas + doesn’t check it). If not specified, entries will be filled with + the corresponding NULL value (np.nan for numpy dtypes, pd.NA for + extension dtypes). + + Returns + ------- + Series + """ + raise NotImplementedError("abstract method") + + def mask(self, cond, other): + """Replace values where the condition is True. + + Parameters + ---------- + cond: bool Series/DataFrame, array-like, or callable + Where cond is False, keep the original value. Where True, replace + with corresponding value from other. If cond is callable, it is + computed on the Series/DataFrame and should return boolean + Series/DataFrame or array. The callable must not change input + Series/DataFrame (though pandas doesn’t check it). + other: scalar, Series/DataFrame, or callable + Entries where cond is True are replaced with corresponding value + from other. If other is callable, it is computed on the + Series/DataFrame and should return scalar or Series/DataFrame. + The callable must not change input Series/DataFrame (though pandas + doesn’t check it). If not specified, entries will be filled with + the corresponding NULL value (np.nan for numpy dtypes, pd.NA for + extension dtypes). + + Returns + ------- + Series + """ + raise NotImplementedError("abstract method") + + def clip(self): + """Trim values at input threshold(s). + + Assigns values outside boundary to boundary values. Thresholds can be + singular values or array like, and in the latter case the clipping is + performed element-wise in the specified axis. + + Parameters + ---------- + + lower: float or array-like, default None + Minimum threshold value. All values below this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. + + upper: float or array-like, default None + Maximum threshold value. All values above this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. + + Returns + ------- + Series + """ + raise NotImplementedError("abstract method") + + def argmax(self): + """ + Return int position of the smallest value in the Series. + + If the minimum is achieved in multiple locations, the first row position is returned. + + Returns + ------- + Series + Row position of the maximum value. + """ + raise NotImplementedError("abstract method") + + def argmin(self): + """ + Return int position of the largest value in the Series. + + If the maximum is achieved in multiple locations, the first row position is returned. + + Returns + ------- + Series + Row position of the minimum value. + """ + raise NotImplementedError("abstract method") + + def rename(self, index, **kwargs) -> Series | None: + """ + Alter Series index labels or name. + + Function / dict values must be unique (1-to-1). Labels not contained in + a dict / Series will be left as-is. Extra labels listed don't throw an + error. + + Alternatively, change ``Series.name`` with a scalar value. + + Parameters + ---------- + index : scalar, hashable sequence, dict-like or function optional + Functions or dict-like are transformations to apply to + the index. + Scalar or hashable sequence-like will alter the ``Series.name`` + attribute. + + Returns + ------- + Series + Series with index labels + + """ + raise NotImplementedError("abstract method") + + def rename_axis(self, mapper, **kwargs): + """ + Set the name of the axis for the index or columns. + + Parameters + ---------- + mapper : scalar, list-like, optional + Value to set the axis name attribute. + + Returns + ------- + Series + """ + raise NotImplementedError("abstract method") + + def rolling( + self, + window, + min_periods: int | None = None, + ): + """ + Provide rolling window calculations. + + Parameters + ---------- + window : int, timedelta, str, offset, or BaseIndexer subclass + Size of the moving window. + + If an integer, the fixed number of observations used for + each window. + + If a timedelta, str, or offset, the time period of each window. Each + window will be a variable sized based on the observations included in + the time-period. This is only valid for datetimelike indexes. + To learn more about the offsets & frequency strings, please see `this link + `__. + + If a BaseIndexer subclass, the window boundaries + based on the defined ``get_window_bounds`` method. Additional rolling + keyword arguments, namely ``min_periods``, ``center``, ``closed`` and + ``step`` will be passed to ``get_window_bounds``. + + min_periods : int, default None + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. + + For a window that is specified by an offset, ``min_periods`` will default to 1. + + For a window that is specified by an integer, ``min_periods`` will default + to the size of the window. + + Returns + ------- + ``Window`` subclass if a ``win_type`` is passed + + ``Rolling`` subclass if ``win_type`` is not passed + """ + raise NotImplementedError("abstract method") + + def expanding(self, min_periods=1): + """ + Provide expanding window calculations. + + Parameters + ---------- + min_periods : int, default 1 + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. + + Returns + ------- + ``Expanding`` subclass + """ + raise NotImplementedError("abstract method") + + def value_counts( + self, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + *, + dropna: bool = True, + ) -> Series: + """ + Return a Series containing counts of unique values. + + The resulting object will be in descending order so that the + first element is the most frequently-occurring element. + Excludes NA values by default. + + Args: + normalize : bool, default False + If True then the object returned will contain the relative + frequencies of the unique values. + sort : bool, default True + Sort by frequencies. + ascending : bool, default False + Sort in ascending order. + dropna : bool, default True + Don't include counts of NaN. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + + @property + def str(self): + """ + Vectorized string functions for Series and Index. + + NAs stay NA unless handled otherwise by a particular method. Patterned + after Python’s string methods, with some inspiration from R’s stringr package. + """ + raise NotImplementedError("abstract property") diff --git a/third_party/bigframes_vendored/pandas/core/strings/accessor.py b/third_party/bigframes_vendored/pandas/core/strings/accessor.py new file mode 100644 index 00000000000..81b9bd3d146 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/strings/accessor.py @@ -0,0 +1,187 @@ +class StringMethods: + """ + Vectorized string functions for Series and Index. + + NAs stay NA unless handled otherwise by a particular method. + Patterned after Python's string methods, with some inspiration from + R's stringr package. + """ + + def find(self, sub, start: int = 0, end=None): + """Return lowest indexes in each strings in the Series/Index. + + Each of returned indexes corresponds to the position where the + substring is fully contained between [start:end]. Return -1 on + failure. Equivalent to standard :meth:`str.find`. + + Args: + sub: + Substring being searched. + start: + Left edge index. + end: + Right edge index. + + Returns: + Series or Index of int. + """ + + raise NotImplementedError("abstract method") + + def len(self): + """Compute the length of each element in the Series/Index. + + The element may be a sequence (such as a string, tuple or list) or a collection + (such as a dictionary). + + Returns: + Series or Index of int + A Series or Index of integer values indicating the length of each + element in the Series or Index. + """ + + raise NotImplementedError("abstract method") + + def lower(self): + """Convert strings in the Series/Index to lowercase. + + Equivalent to :meth:`str.lower`. + + Returns: + Series or Index of object + """ + + raise NotImplementedError("abstract method") + + def slice(self, start=None, stop=None): + """Slice substrings from each element in the Series or Index. + + Args: + start : int, optional + Start position for slice operation. + stop : int, optional + Stop position for slice operation. + step : int, optional + Step size for slice operation. + + Returns: + Series or Index of object + Series or Index from sliced substring from original string object. + """ + + raise NotImplementedError("abstract method") + + def strip(self): + """Remove leading and trailing characters. + + Strip whitespaces (including newlines) or a set of specified characters + from each string in the Series/Index from left and right sides. + Replaces any non-strings in Series with NaNs. + Equivalent to :meth:`str.strip`. + + Returns: + Series or Index of object + """ + + raise NotImplementedError("abstract method") + + def upper(self): + """Convert strings in the Series/Index to uppercase. + + Equivalent to :meth:`str.upper`. + + Returns: + Series or Index of object + """ + + raise NotImplementedError("abstract method") + + def isnumeric(self): + """Check whether all characters in each string are numeric. + + This is equivalent to running the Python string method + :meth:`str.isnumeric` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + Series or Index of bool + Series or Index of boolean values with the same length as the original + Series/Index. + """ + + raise NotImplementedError("abstract method") + + def rstrip(self): + """Remove trailing characters. + + Strip whitespaces (including newlines) or a set of specified characters + from each string in the Series/Index from right side. + Replaces any non-strings in Series with NaNs. + Equivalent to :meth:`str.rstrip`. + + Returns: + Series or Index of object + """ + + raise NotImplementedError("abstract method") + + def lstrip(self): + """Remove leading characters. + + Strip whitespaces (including newlines) or a set of specified characters + from each string in the Series/Index from left side. + Replaces any non-strings in Series with NaNs. + Equivalent to :meth:`str.lstrip`. + + Returns: + Series or Index of object` + """ + + raise NotImplementedError("abstract method") + + def repeat(self, repeats: int): + """Duplicate each string in the Series or Index. + + Args: + repeats : int or sequence of int + Same value for all (int) or different value per (sequence). + + Returns: + Series or pandas.Index + Series or Index of repeated string objects specified by + input parameter repeats. + """ + + raise NotImplementedError("abstract method") + + def capitalize(self): + """Convert strings in the Series/Index to be capitalized. + + Equivalent to :meth:`str.capitalize`. + + Returns: + Series or Index of object + """ + + raise NotImplementedError("abstract method") + + def cat(self, others, *, join): + """Concatenate strings in the Series/Index with given separator. + + If `others` is specified, this function concatenates the Series/Index + and elements of `others` element-wise. + + Args: + others : Series + + join : {'left', 'outer'}, default 'left' + Determines the join-style between the calling Series and any + Series in `others` (objects without an index need + to match the length of the calling Series). To disable + alignment, use `.values` on any Series/Index/DataFrame in `others`. + + Returns: + Series + """ + + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/window/__init__.py b/third_party/bigframes_vendored/pandas/core/window/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/window/rolling.py b/third_party/bigframes_vendored/pandas/core/window/rolling.py new file mode 100644 index 00000000000..7a9239b70cf --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/window/rolling.py @@ -0,0 +1,37 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/window/rolling.py +""" +Provide a generic structure to support window functions, +similar to how we have a Groupby object. +""" + + +class Window: + """Provide window calculations.""" + + def count(self): + """Calculate the window count of non-NULL observations.""" + raise NotImplementedError("abstract method") + + def sum(self): + """Calculate the weighted window sum.""" + raise NotImplementedError("abstract method") + + def mean(self): + """Calculate the weighted window mean.""" + raise NotImplementedError("abstract method") + + def var(self): + """Calculate the weighted window variance.""" + raise NotImplementedError("abstract method") + + def std(self): + """Calculate the weighted window standard deviation.""" + raise NotImplementedError("abstract method") + + def max(self): + """Calculate the weighted window maximum.""" + raise NotImplementedError("abstract method") + + def min(self): + """Calculate the weighted window minimum.""" + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/__init__.py b/third_party/bigframes_vendored/pandas/io/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/io/common.py b/third_party/bigframes_vendored/pandas/io/common.py new file mode 100644 index 00000000000..057ed96721d --- /dev/null +++ b/third_party/bigframes_vendored/pandas/io/common.py @@ -0,0 +1,44 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/io/common.py +"""Common IO api utilities""" +from __future__ import annotations + +from collections import defaultdict +from typing import DefaultDict, Hashable, Sequence + + +def dedup_names( + names: Sequence[Hashable], + is_potential_multiindex: bool, +) -> Sequence[Hashable]: + """ + Rename column names if duplicates exist. + + Currently the renaming is done by appending a period and an autonumeric, + but a custom pattern may be supported in the future. + + Examples + -------- + >>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False) + ['x', 'y', 'x.1', 'x.2'] + """ + names = list(names) # so we can index + counts: DefaultDict[Hashable, int] = defaultdict(int) + + for i, col in enumerate(names): + cur_count = counts[col] + + while cur_count > 0: + counts[col] = cur_count + 1 + + if is_potential_multiindex: + # for mypy + assert isinstance(col, tuple) + col = col[:-1] + (f"{col[-1]}.{cur_count}",) + else: + col = f"{col}.{cur_count}" + cur_count = counts[col] + + names[i] = col + counts[col] = cur_count + 1 + + return names diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py new file mode 100644 index 00000000000..4f4b1b8199f --- /dev/null +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -0,0 +1,37 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/io/gbq.py +""" Google BigQuery support """ + +from __future__ import annotations + +from typing import Iterable, Optional + + +class GBQIOMixin: + def read_gbq( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + ): + """Loads DataFrame from Google BigQuery. + + Args: + query: + A SQL string to be executed or a BigQuery table to be read. The + table must be specified in the format of + `project.dataset.tablename` or `dataset.tablename`. + index_col: + Name of result column(s) to use for index in results DataFrame. + col_order: + List of BigQuery column names in the desired order for results + DataFrame. + max_results: + If set, limit the maximum number of rows to fetch from the + query results. + + Returns: + A DataFrame representing results of the query or table. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/parquet.py b/third_party/bigframes_vendored/pandas/io/parquet.py new file mode 100644 index 00000000000..ee2e9a65f22 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/io/parquet.py @@ -0,0 +1,22 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/io/parquet.py +""" parquet compat """ +from __future__ import annotations + + +class ParquetIOMixin: + def read_parquet( + self, + path: str, + ): + r"""Load a parquet object from the file path (local or GCS), returning a DataFrame. + + Args: + path: + Local or GCS path to parquet file. + + Note: + This method will not guarantee the same ordering as the file. + Instead, set a serialized index column as the index and sort by + that in the resulting DataFrame. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/parsers/__init__.py b/third_party/bigframes_vendored/pandas/io/parsers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/io/parsers/readers.py b/third_party/bigframes_vendored/pandas/io/parsers/readers.py new file mode 100644 index 00000000000..e5ed39b54aa --- /dev/null +++ b/third_party/bigframes_vendored/pandas/io/parsers/readers.py @@ -0,0 +1,110 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/io/parsers/readers.py +""" +Module contains tools for processing files into DataFrames or other objects + +GH#48849 provides a convenient way of deprecating keyword arguments +""" +from __future__ import annotations + +from typing import Any, Dict, Literal, MutableSequence, Optional, Sequence, Tuple, Union + +import numpy as np + + +class ReaderIOMixin: + def read_csv( + self, + filepath_or_buffer: str, + *, + sep: Optional[str] = ",", + header: Optional[int] = 0, + names: Optional[ + Union[MutableSequence[Any], np.ndarray[Any, Any], Tuple[Any, ...], range] + ] = None, + index_col: Optional[ + Union[int, str, Sequence[Union[str, int]], Literal[False]] + ] = None, + usecols=None, + dtype: Optional[Dict] = None, + engine: Optional[ + Literal["c", "python", "pyarrow", "python-fwf", "bigquery"] + ] = None, + encoding: Optional[str] = None, + **kwargs, + ): + r"""Loads DataFrame from comma-separated values (csv) file locally or from GCS. + + The CSV file data will be persisted as a temporary BigQuery table, which can be + automatically recycled after the Session is closed. + + Note: using `engine="bigquery"` will not guarantee the same ordering as the + file. Instead, set a serialized index column as the index and sort by + that in the resulting DataFrame. + + Args: + filepath_or_buffer: a string path including GCS and local file. + + sep: the separator for fields in a CSV file. For the BigQuery engine, the separator + can be any ISO-8859-1 single-byte character. To use a character in the range + 128-255, you must encode the character as UTF-8. Both engines support + `sep="\t"` to specify tab character as separator. Default engine supports + having any number of spaces as separator by specifying `sep="\s+"`. Separators + longer than 1 character are interpreted as regular expressions by the default + engine. BigQuery engine only supports single character separators. + + header: row number to use as the column names. + - ``None``: Instructs autodetect that there are no headers and data should be + read starting from the first row. + - ``0``: If using `engine="bigquery"`, Autodetect tries to detect headers in the + first row. If they are not detected, the row is read as data. Otherwise data + is read starting from the second row. When using default engine, pandas assumes + the first row contains column names unless the `names` argument is specified. + If `names` is provided, then the first row is ignored, second row is read as + data, and column names are inferred from `names`. + - ``N > 0``: If using `engine="bigquery"`, Autodetect skips N rows and tries + to detect headers in row N+1. If headers are not detected, row N+1 is just + skipped. Otherwise row N+1 is used to extract column names for the detected + schema. When using default engine, pandas will skip N rows and assumes row N+1 + contains column names unless the `names` argument is specified. If `names` is + provided, row N+1 will be ignored, row N+2 will be read as data, and column + names are inferred from `names`. + + names: a list of column names to use. If the file contains a header row and you + want to pass this parameter, then `header=0` should be passed as well so the + first (header) row is ignored. Only to be used with default engine. + + index_col: column(s) to use as the row labels of the DataFrame, either given as + string name or column index. `index_col=False` can be used with the default + engine only to enforce that the first column is not used as the index. Using + column index instead of column name is only supported with the default engine. + The BigQuery engine only supports having a single column name as the `index_col`. + Neither engine supports having a multi-column index. + + usecols: list of column names to use. The BigQuery engine only supports having a list + of string column names. Column indices and callable functions are only supported + with the default engine. Using the default engine, the column names in `usecols` + can be defined to correspond to column names provided with the `names` parameter + (ignoring the document's header row of column names). The order of the column + indices/names in `usecols` is ignored with the default engine. The order of the + column names provided with the BigQuery engine will be consistent in the resulting + dataframe. If using a callable function with the default engine, only column names + that evaluate to True by the callable function will be in the resulting dataframe. + + dtype: data type for data or columns. Only to be used with default engine. + + engine: type of engine to use. If `engine="bigquery"` is specified, then BigQuery's + load API will be used. Otherwise, the engine will be passed to `pandas.read_csv`. + + encoding: the character encoding of the data. The default encoding is `UTF-8` for both + engines. The default engine acceps a wide range of encodings. Refer to Python + documentation for a comprehensive list, + https://docs.python.org/3/library/codecs.html#standard-encodings + The BigQuery engine only supports `UTF-8` and `ISO-8859-1`. + + **kwargs: keyword arguments. + + + Returns: + A BigQuery DataFrames. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/pandas/_typing.py b/third_party/bigframes_vendored/pandas/pandas/_typing.py new file mode 100644 index 00000000000..e665339fc83 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/pandas/_typing.py @@ -0,0 +1,425 @@ +# Copied from https://github.com/pandas-dev/pandas/blob/main/pandas/_typing.py +from __future__ import annotations + +from datetime import datetime, timedelta, tzinfo +from os import PathLike +import sys +from typing import ( + Any, + Callable, + Dict, + Hashable, + Iterator, + List, + Literal, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, +) +from typing import Type as type_t +from typing import TYPE_CHECKING, TypeVar, Union + +import numpy as np + +# To prevent import cycles place any internal imports in the branch below +# and use a string literal forward reference to it in subsequent types +# https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles +if TYPE_CHECKING: + import numpy.typing as npt + from pandas import Interval + from pandas._libs import NaTType, Period, Timedelta, Timestamp + from pandas._libs.tslibs import BaseOffset + from pandas.arrays import DatetimeArray, TimedeltaArray + from pandas.core.arrays.base import ExtensionArray + from pandas.core.dtypes.dtypes import ExtensionDtype + from pandas.core.frame import DataFrame + from pandas.core.generic import NDFrame + from pandas.core.groupby.generic import DataFrameGroupBy, GroupBy, SeriesGroupBy + from pandas.core.indexes.base import Index + from pandas.core.internals import ( + ArrayManager, + BlockManager, + SingleArrayManager, + SingleBlockManager, + ) + from pandas.core.resample import Resampler + from pandas.core.series import Series + from pandas.core.window.rolling import BaseWindow + from pandas.io.formats.format import EngFormatter + + ScalarLike_co = Union[ + int, + float, + complex, + str, + bytes, + np.generic, + ] + + # numpy compatible types + NumpyValueArrayLike = Union[ScalarLike_co, npt.ArrayLike] + # Name "npt._ArrayLikeInt_co" is not defined [name-defined] + NumpySorter = Optional[npt._ArrayLikeInt_co] # type: ignore[name-defined] + + if sys.version_info >= (3, 10): + from typing import TypeGuard + else: + from typing_extensions import TypeGuard # pyright: reportUnusedImport = false + + if sys.version_info >= (3, 11): + from typing import Self + else: + from typing_extensions import Self # pyright: reportUnusedImport = false +else: + npt: Any = None + Self: Any = None + TypeGuard: Any = None + +HashableT = TypeVar("HashableT", bound=Hashable) + +# array-like + +ArrayLike = Union["ExtensionArray", np.ndarray] +AnyArrayLike = Union[ArrayLike, "Index", "Series"] +TimeArrayLike = Union["DatetimeArray", "TimedeltaArray"] + +# list-like + +# Cannot use `Sequence` because a string is a sequence, and we don't want to +# accept that. Could refine if https://github.com/python/typing/issues/256 is +# resolved to differentiate between Sequence[str] and str +ListLike = Union[AnyArrayLike, List, range] + +# scalars + +PythonScalar = Union[str, float, bool] +DatetimeLikeScalar = Union["Period", "Timestamp", "Timedelta"] +PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"] +Scalar = Union[PythonScalar, PandasScalar, np.datetime64, np.timedelta64, datetime] +IntStrT = TypeVar("IntStrT", int, str) + + +# timestamp and timedelta convertible types + +TimestampConvertibleTypes = Union[ + "Timestamp", datetime, np.datetime64, np.int64, float, str +] +TimedeltaConvertibleTypes = Union[ + "Timedelta", timedelta, np.timedelta64, np.int64, float, str +] +Timezone = Union[str, tzinfo] + +ToTimestampHow = Literal["s", "e", "start", "end"] + +# NDFrameT is stricter and ensures that the same subclass of NDFrame always is +# used. E.g. `def func(a: NDFrameT) -> NDFrameT: ...` means that if a +# Series is passed into a function, a Series is always returned and if a DataFrame is +# passed in, a DataFrame is always returned. +NDFrameT = TypeVar("NDFrameT", bound="NDFrame") + +NumpyIndexT = TypeVar("NumpyIndexT", np.ndarray, "Index") + +AxisInt = int +Axis = Union[AxisInt, Literal["index", "columns", "rows"]] +IndexLabel = Union[Hashable, Sequence[Hashable]] +Level = Hashable +Shape = Tuple[int, ...] +Suffixes = Tuple[Optional[str], Optional[str]] +Ordered = Optional[bool] +JSONSerializable = Optional[Union[PythonScalar, List, Dict]] +Frequency = Union[str, "BaseOffset"] +Axes = ListLike + +RandomState = Union[ + int, + np.ndarray, + np.random.Generator, + np.random.BitGenerator, + np.random.RandomState, +] + +# dtypes +NpDtype = Union[str, np.dtype, type_t[Union[str, complex, bool, object]]] +Dtype = Union["ExtensionDtype", NpDtype] +AstypeArg = Union["ExtensionDtype", "npt.DTypeLike"] +# DtypeArg specifies all allowable dtypes in a functions its dtype argument +DtypeArg = Union[Dtype, Dict[Hashable, Dtype]] +DtypeObj = Union[np.dtype, "ExtensionDtype"] + +# converters +ConvertersArg = Dict[Hashable, Callable[[Dtype], Dtype]] + +# parse_dates +ParseDatesArg = Union[ + bool, List[Hashable], List[List[Hashable]], Dict[Hashable, List[Hashable]] +] + +# For functions like rename that convert one label to another +Renamer = Union[Mapping[Any, Hashable], Callable[[Any], Hashable]] + +# to maintain type information across generic functions and parametrization +T = TypeVar("T") + +# used in decorators to preserve the signature of the function it decorates +# see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators +FuncType = Callable[..., Any] +F = TypeVar("F", bound=FuncType) + +# types of vectorized key functions for DataFrame::sort_values and +# DataFrame::sort_index, among others +ValueKeyFunc = Optional[Callable[["Series"], Union["Series", AnyArrayLike]]] +IndexKeyFunc = Optional[Callable[["Index"], Union["Index", AnyArrayLike]]] + +# types of `func` kwarg for DataFrame.aggregate and Series.aggregate +AggFuncTypeBase = Union[Callable, str] +AggFuncTypeDict = Dict[Hashable, Union[AggFuncTypeBase, List[AggFuncTypeBase]]] +AggFuncType = Union[ + AggFuncTypeBase, + List[AggFuncTypeBase], + AggFuncTypeDict, +] +AggObjType = Union[ + "Series", + "DataFrame", + "GroupBy", + "SeriesGroupBy", + "DataFrameGroupBy", + "BaseWindow", + "Resampler", +] + +PythonFuncType = Callable[[Any], Any] + +# filenames and file-like-objects +AnyStr_co = TypeVar("AnyStr_co", str, bytes, covariant=True) +AnyStr_contra = TypeVar("AnyStr_contra", str, bytes, contravariant=True) + + +class BaseBuffer(Protocol): + @property + def mode(self) -> str: + # for _get_filepath_or_buffer + ... + + def seek(self, __offset: int, __whence: int = ...) -> int: + # with one argument: gzip.GzipFile, bz2.BZ2File + # with two arguments: zip.ZipFile, read_sas + ... + + def seekable(self) -> bool: + # for bz2.BZ2File + ... + + def tell(self) -> int: + # for zip.ZipFile, read_stata, to_stata + ... + + +class ReadBuffer(BaseBuffer, Protocol[AnyStr_co]): + def read(self, __n: int = ...) -> AnyStr_co: + # for BytesIOWrapper, gzip.GzipFile, bz2.BZ2File + ... + + +class WriteBuffer(BaseBuffer, Protocol[AnyStr_contra]): + def write(self, __b: AnyStr_contra) -> Any: + # for gzip.GzipFile, bz2.BZ2File + ... + + def flush(self) -> Any: + # for gzip.GzipFile, bz2.BZ2File + ... + + +class ReadPickleBuffer(ReadBuffer[bytes], Protocol): + def readline(self) -> bytes: + ... + + +class WriteExcelBuffer(WriteBuffer[bytes], Protocol): + def truncate(self, size: int | None = ...) -> int: + ... + + +class ReadCsvBuffer(ReadBuffer[AnyStr_co], Protocol): + def __iter__(self) -> Iterator[AnyStr_co]: + # for engine=python + ... + + def fileno(self) -> int: + # for _MMapWrapper + ... + + def readline(self) -> AnyStr_co: + # for engine=python + ... + + @property + def closed(self) -> bool: + # for enine=pyarrow + ... + + +FilePath = Union[str, "PathLike[str]"] + +# for arbitrary kwargs passed during reading/writing files +StorageOptions = Optional[Dict[str, Any]] + + +# compression keywords and compression +CompressionDict = Dict[str, Any] +CompressionOptions = Optional[ + Union[Literal["infer", "gzip", "bz2", "zip", "xz", "zstd", "tar"], CompressionDict] +] + +# types in DataFrameFormatter +FormattersType = Union[ + List[Callable], Tuple[Callable, ...], Mapping[Union[str, int], Callable] +] +ColspaceType = Mapping[Hashable, Union[str, int]] +FloatFormatType = Union[str, Callable, "EngFormatter"] +ColspaceArgType = Union[ + str, int, Sequence[Union[str, int]], Mapping[Hashable, Union[str, int]] +] + +# Arguments for fillna() +FillnaOptions = Literal["backfill", "bfill", "ffill", "pad"] + +# internals +Manager = Union[ + "ArrayManager", "SingleArrayManager", "BlockManager", "SingleBlockManager" +] +SingleManager = Union["SingleArrayManager", "SingleBlockManager"] +Manager2D = Union["ArrayManager", "BlockManager"] + +# indexing +# PositionalIndexer -> valid 1D positional indexer, e.g. can pass +# to ndarray.__getitem__ +# ScalarIndexer is for a single value as the index +# SequenceIndexer is for list like or slices (but not tuples) +# PositionalIndexerTuple is extends the PositionalIndexer for 2D arrays +# These are used in various __getitem__ overloads +# TODO(typing#684): add Ellipsis, see +# https://github.com/python/typing/issues/684#issuecomment-548203158 +# https://bugs.python.org/issue41810 +# Using List[int] here rather than Sequence[int] to disallow tuples. +ScalarIndexer = Union[int, np.integer] +SequenceIndexer = Union[slice, List[int], np.ndarray] +PositionalIndexer = Union[ScalarIndexer, SequenceIndexer] +PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer] +PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple] +if TYPE_CHECKING: + TakeIndexer = Union[Sequence[int], Sequence[np.integer], npt.NDArray[np.integer]] +else: + TakeIndexer = Any + +# Shared by functions such as drop and astype +IgnoreRaise = Literal["ignore", "raise"] + +# Windowing rank methods +WindowingRankType = Literal["average", "min", "max"] + +# read_csv engines +CSVEngine = Literal["c", "python", "pyarrow", "python-fwf"] + +# read_json engines +JSONEngine = Literal["ujson", "pyarrow"] + +# read_xml parsers +XMLParsers = Literal["lxml", "etree"] + +# Interval closed type +IntervalLeftRight = Literal["left", "right"] +IntervalClosedType = Union[IntervalLeftRight, Literal["both", "neither"]] + +# datetime and NaTType +DatetimeNaTType = Union[datetime, "NaTType"] +DateTimeErrorChoices = Union[IgnoreRaise, Literal["coerce"]] + +# sort_index +SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] +NaPosition = Literal["first", "last"] + +# Arguments for nsmalles and n_largest +NsmallestNlargestKeep = Literal["first", "last", "all"] + +# quantile interpolation +QuantileInterpolation = Literal["linear", "lower", "higher", "midpoint", "nearest"] + +# plotting +PlottingOrientation = Literal["horizontal", "vertical"] + +# dropna +AnyAll = Literal["any", "all"] + +# merge +MergeHow = Literal["left", "right", "inner", "outer", "cross"] +MergeValidate = Literal[ + "one_to_one", + "1:1", + "one_to_many", + "1:m", + "many_to_one", + "m:1", + "many_to_many", + "m:m", +] + +# join +JoinHow = Literal["left", "right", "inner", "outer"] +JoinValidate = Literal[ + "one_to_one", + "1:1", + "one_to_many", + "1:m", + "many_to_one", + "m:1", + "many_to_many", + "m:m", +] + +# reindex +ReindexMethod = Union[FillnaOptions, Literal["nearest"]] + +MatplotlibColor = Union[str, Sequence[float]] +TimeGrouperOrigin = Union[ + "Timestamp", Literal["epoch", "start", "start_day", "end", "end_day"] +] +TimeAmbiguous = Union[Literal["infer", "NaT", "raise"], "npt.NDArray[np.bool_]"] +TimeNonexistent = Union[ + Literal["shift_forward", "shift_backward", "NaT", "raise"], timedelta +] +DropKeep = Literal["first", "last", False] +CorrelationMethod = Union[ + Literal["pearson", "kendall", "spearman"], Callable[[np.ndarray, np.ndarray], float] +] +AlignJoin = Literal["outer", "inner", "left", "right"] +DtypeBackend = Literal["pyarrow", "numpy_nullable"] + +TimeUnit = Literal["s", "ms", "us", "ns"] +OpenFileErrors = Literal[ + "strict", + "ignore", + "replace", + "surrogateescape", + "xmlcharrefreplace", + "backslashreplace", + "namereplace", +] + +# update +UpdateJoin = Literal["left"] + +# applymap +NaAction = Literal["ignore"] + +# from_dict +FromDictOrient = Literal["columns", "index", "tight"] + +# to_gbc +ToGbqIfexist = Literal["fail", "replace", "append"] + +# to_stata +ToStataByteorder = Literal[">", "<", "little", "big"] diff --git a/third_party/bigframes_vendored/sklearn/COPYING b/third_party/bigframes_vendored/sklearn/COPYING new file mode 100644 index 00000000000..b161c890897 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/COPYING @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2007-2023 The scikit-learn developers. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py new file mode 100644 index 00000000000..925bb76e1a1 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -0,0 +1,128 @@ +"""Base classes for all estimators.""" + +# Author: Gael Varoquaux +# License: BSD 3 clause +# Original location: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/base.py + +from __future__ import annotations + +import inspect +from typing import Any, Dict, List + + +class BaseEstimator: + """Base class for all estimators. + + Notes: + All estimators should specify all the parameters that can be set + at the class level in their ``__init__`` as explicit keyword + arguments (no ``*args`` or ``**kwargs``). + """ + + @classmethod + def _get_param_names(cls: type[BaseEstimator]) -> List[str]: + """Get parameter names for the estimator""" + # fetch the constructor or the original constructor before + # deprecation wrapping if any + init = getattr(cls.__init__, "deprecated_original", cls.__init__) + if init is object.__init__: + # No explicit constructor to introspect + return [] + + # introspect the constructor arguments to find the model parameters + # to represent + init_signature = inspect.signature(init) + # Consider the constructor parameters excluding 'self' + parameters = [ + p + for p in init_signature.parameters.values() + if p.name != "self" and p.kind != p.VAR_KEYWORD + ] + for p in parameters: + if p.kind == p.VAR_POSITIONAL: + raise RuntimeError( + "Estimators should always " + "specify their parameters in the signature" + " of their __init__ (no varargs)." + " %s with constructor %s doesn't " + " follow this convention." % (cls, init_signature) + ) + # Extract and sort argument names excluding 'self' + return sorted([p.name for p in parameters]) + + def get_params(self, deep: bool = True) -> Dict[str, Any]: + """Get parameters for this estimator. + + Args: + deep: + Default ``True``. If True, will return the parameters for this + estimator and contained subobjects that are estimators. + + Returns: + A dictionary of parameter names mapped to their values. + """ + out: Dict = dict() + for key in self._get_param_names(): + value = getattr(self, key) + if deep and hasattr(value, "get_params") and not isinstance(value, type): + deep_items = value.get_params().items() + out.update((key + "__" + k, val) for k, val in deep_items) + out[key] = value + return out + + +class ClassifierMixin: + """Mixin class for all classifiers.""" + + _estimator_type = "classifier" + + def score(self, X, y): + """Return the mean accuracy on the given test data and labels. + + In multi-label classification, this is the subset accuracy + which is a harsh metric since you require for each sample that + each label set be correctly predicted. + + Args: + X: + DataFrame of shape (n_samples, n_features). Test samples. + + y: + DataFrame of shape (n_samples,) or (n_samples, n_outputs). True + labels for `X`. + + Returns: + A DataFrame of the evaluation result. + """ + raise NotImplementedError("abstract method") + + +class RegressorMixin: + """Mixin class for all regression estimators.""" + + _estimator_type = "regressor" + + def score(self, X, y): + """Return the evaluation metrics of the model. + + Args: + X: + DataFrame of shape (n_samples, n_features). Test samples. For + some estimators this may be a precomputed kernel matrix or a + list of generic objects instead with shape + ``(n_samples, n_samples_fitted)``, where ``n_samples_fitted`` + is the number of samples used in the fitting for the estimator. + + y: + DataFrame of shape (n_samples,) or (n_samples, n_outputs). True + values for `X`. + + Returns: + A DataFrame of the evaluation result. + """ + raise NotImplementedError("abstract method") + + +class MetaEstimatorMixin: + _required_parameters = ["estimator"] + """Mixin class for all meta estimators in scikit-learn.""" diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py new file mode 100644 index 00000000000..8f67cab1c31 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -0,0 +1,71 @@ +"""K-means clustering.""" + +# Authors: Gael Varoquaux +# Thomas Rueckstiess +# James Bergstra +# Jan Schlueter +# Nelle Varoquaux +# Peter Prettenhofer +# Olivier Grisel +# Mathieu Blondel +# Robert Layton +# License: BSD 3 clause + +from abc import ABC +from typing import List, Optional + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class _BaseKMeans(BaseEstimator, ABC): + """Base class for KMeans and MiniBatchKMeans""" + + def predict(self, X): + """Predict the closest cluster each sample in X belongs to. + + Args: + X: + DataFrame of shape (n_samples, n_features). The data matrix for + which we want to get the predictions. + + Returns: + DataFrame of shape (n_samples,), containing the class labels for + each sample. + """ + raise NotImplementedError("abstract method") + + +class KMeans(_BaseKMeans): + """K-Means clustering. + + Args: + n_clusters: int, default=8 + The number of clusters to form as well as the number of centroids to generate. + Default to 8. + """ + + def fit( + self, + X, + y=None, + transforms: Optional[List[str]] = None, + ): + """Compute k-means clustering. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training data. + y: Ignored + Not used, present here for API consistency by convention. + + transforms: + An optional list of SQL expressions to apply over top of the + model inputs as preprocessing. This preprocessing will be + automatically reapplied to new input data (e.g. in .predict), + and may contain steps (like ML.STANDARD_SCALER) that fit to the + training data. + + Returns: + Fitted Estimator. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py new file mode 100644 index 00000000000..dadfa5d0131 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py @@ -0,0 +1,61 @@ +"""Utilities for meta-estimators""" +# Author: Joel Nothman +# Andreas Mueller +# License: BSD + + +from abc import ABCMeta + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class _BaseComposition(BaseEstimator, metaclass=ABCMeta): + """Handles parameter management for classifiers composed of named estimators.""" + + +class ColumnTransformer(_BaseComposition): + """Applies transformers to columns of BigQuery DataFrames. + + This estimator allows different columns or column subsets of the input + to be transformed separately and the features generated by each transformer + will be concatenated to form a single feature space. + This is useful for heterogeneous or columnar data, to combine several + feature extraction mechanisms or transformations into a single transformer. + + Args: + transformers: + List of (name, transformer, columns) tuples specifying the transformer + objects to be applied to subsets of the data. + """ + + def fit( + self, + X, + ): + """Fit all transformers using X. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training vector, + where `n_samples` is the number of samples and `n_features` is + the number of features. + + Returns: + Fitted estimator. + """ + raise NotImplementedError("abstract method") + + def transform( + self, + X, + ): + """Transform X separately by each transformer, concatenate results. + + Args: + X: + The DataFrame to be transformed by subset. + + Returns: + Transformed result. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py new file mode 100644 index 00000000000..6f0d17b2e6d --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py @@ -0,0 +1,53 @@ +""" Principal Component Analysis. +""" + +# Author: Alexandre Gramfort +# Olivier Grisel +# Mathieu Blondel +# Denis A. Engemann +# Michael Eickenberg +# Giorgio Patrini +# +# License: BSD 3 clause + +from abc import ABCMeta + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class PCA(BaseEstimator, metaclass=ABCMeta): + """Principal component analysis (PCA). + + Linear dimensionality reduction using Singular Value Decomposition of the + data to project it to a lower dimensional space. The input data is centered + but not scaled for each feature before applying the SVD. + + It uses the LAPACK implementation of the full SVD or a randomized truncated + SVD by the method of Halko et al. 2009, depending on the shape of the input + data and the number of components to extract. + + It can also use the scipy.sparse.linalg ARPACK implementation of the + truncated SVD. + + Args: + n_components: Optional[int] + Number of components to keep. if n_components is not set all components are kept. + + """ + + def fit( + self, + X, + ): + """Fit the model according to the given training data. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training vector, + where `n_samples` is the number of samples and `n_features` is + the number of features. + + Returns: + Fitted estimator. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py new file mode 100644 index 00000000000..c91821c7622 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py @@ -0,0 +1,191 @@ +""" +Forest of trees-based ensemble methods. +Those methods include random forests and extremely randomized trees. +The module structure is the following: +- The ``BaseForest`` base class implements a common ``fit`` method for all + the estimators in the module. The ``fit`` method of the base ``Forest`` + class calls the ``fit`` method of each sub-estimator on random samples + (with replacement, a.k.a. bootstrap) of the training set. + The init of the sub-estimator is further delegated to the + ``BaseEnsemble`` constructor. +- The ``ForestClassifier`` and ``ForestRegressor`` base classes further + implement the prediction logic by computing an average of the predicted + outcomes of the sub-estimators. +- The ``RandomForestClassifier`` and ``RandomForestRegressor`` derived + classes provide the user with concrete implementations of + the forest ensemble method using classical, deterministic + ``DecisionTreeClassifier`` and ``DecisionTreeRegressor`` as + sub-estimator implementations. +- The ``ExtraTreesClassifier`` and ``ExtraTreesRegressor`` derived + classes provide the user with concrete implementations of the + forest ensemble method using the extremely randomized trees + ``ExtraTreeClassifier`` and ``ExtraTreeRegressor`` as + sub-estimator implementations. +Single and multi-output problems are both handled. +""" + +# Authors: Gilles Louppe +# Brian Holt +# Joly Arnaud +# Fares Hedayati +# +# License: BSD 3 clause + +from abc import ABCMeta + +from ..base import BaseEstimator, ClassifierMixin, MetaEstimatorMixin, RegressorMixin + + +class BaseForest(MetaEstimatorMixin, BaseEstimator, metaclass=ABCMeta): + """ + Base class for forests of trees. + """ + + def fit(self, X, y): + """Build a forest of trees from the training set (X, y). + + Args: + X: + DataFrame of shape (n_samples, n_features). Training data. + + y: + DataFrame of shape (n_samples,) or (n_samples, n_targets). + Target values. Will be cast to X's dtype if necessary. + + Returns: + Fitted Estimator. + """ + raise NotImplementedError("abstract method") + + +class ForestRegressor(RegressorMixin, BaseForest, metaclass=ABCMeta): + """ + Base class for forest of trees-based regressors. + """ + + def predict(self, X): + """Predict regression target for X. + + The predicted regression target of an input sample is computed as the + mean predicted regression targets of the trees in the forest. + + Args: + X: + DataFrame of shape (n_samples, n_features). The data matrix for + which we want to get the predictions. + + Returns: + The predicted values. + """ + raise NotImplementedError("abstract method") + + +class RandomForestRegressor(ForestRegressor): + """A random forest regressor. + + A random forest is a meta estimator that fits a number of classifying + decision trees on various sub-samples of the dataset and uses averaging + to improve the predictive accuracy and control over-fitting. + + Args: + num_parallel_tree: Optional[int] + Number of parallel trees constructed during each iteration. Default to 100. Minimum value is 2. + tree_method: Optional[str] + Specify which tree method to use. Default to "auto". If this parameter is set to + default, XGBoost will choose the most conservative option available. Possible values: ""exact", "approx", + "hist". + min_child_weight : Optional[float] + Minimum sum of instance weight(hessian) needed in a child. Default to 1. + colsample_bytree : Optional[float] + Subsample ratio of columns when constructing each tree. Default to 1.0. The value should be between 0 and 1. + colsample_bylevel : Optional[float] + Subsample ratio of columns for each level. Default to 1.0. The value should be between 0 and 1. + colsample_bynode : Optional[float] + Subsample ratio of columns for each split. Default to 0.8. The value should be between 0 and 1. + gamma : Optional[float] + (min_split_loss) Minimum loss reduction required to make a further partition on a + leaf node of the tree. Default to 0.0. + max_depth : Optional[int] + Maximum tree depth for base learners. Default to 15. The value should be greater than 0 and less than 1. + subsample : Optional[float] + Subsample ratio of the training instance. Default to 0.8. The value should be greater than 0 and less than 1. + reg_alpha : Optional[float] + L1 regularization term on weights (xgb's alpha). Default to 0.0. + reg_lambda : Optional[float] + L2 regularization term on weights (xgb's lambda). Default to 1.0. + early_stop: Optional[bool] + Whether training should stop after the first iteration. Default to True. + min_rel_progress: Optional[float] + Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. + enable_global_explain: Optional[bool] + Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. + xgboost_version: Optional[str] + Specifies the Xgboost version for model training. Default to "0.9". Possible values: "0.9", "1.1". + """ + + +class ForestClassifier(ClassifierMixin, BaseForest, metaclass=ABCMeta): + """ + Base class for forest of trees-based classifiers. + """ + + def predict(self, X): + """Predict regression target for X. + + The predicted regression target of an input sample is computed as the + mean predicted regression targets of the trees in the forest. + + Args: + X: + DataFrame of shape (n_samples, n_features). The data matrix for + which we want to get the predictions. + + Returns: + The predicted values. + """ + raise NotImplementedError("abstract method") + + +class RandomForestClassifier(ForestClassifier): + """ + A random forest classifier. + + A random forest is a meta estimator that fits a number of decision tree + classifiers on various sub-samples of the dataset and uses averaging to + improve the predictive accuracy and control over-fitting. + + Args: + num_parallel_tree: Optional[int] + Number of parallel trees constructed during each iteration. Default to 100. Minimum value is 2. + tree_method: Optional[str] + Specify which tree method to use. Default to "auto". If this parameter is set to + default, XGBoost will choose the most conservative option available. Possible values: ""exact", "approx", + "hist". + min_child_weight : Optional[float] + Minimum sum of instance weight(hessian) needed in a child. Default to 1. + colsample_bytree : Optional[float] + Subsample ratio of columns when constructing each tree. Default to 1.0. The value should be between 0 and 1. + colsample_bylevel : Optional[float] + Subsample ratio of columns for each level. Default to 1.0. The value should be between 0 and 1. + colsample_bynode : Optional[float] + Subsample ratio of columns for each split. Default to 0.8. The value should be between 0 and 1. + gamma : Optional[float] + (min_split_loss) Minimum loss reduction required to make a further partition on a + leaf node of the tree. Default to 0.0. + max_depth : Optional[int] + Maximum tree depth for base learners. Default to 15. The value should be greater than 0 and less than 1. + subsample : Optional[float] + Subsample ratio of the training instance. Default to 0.8. The value should be greater than 0 and less than 1. + reg_alpha : Optional[float] + L1 regularization term on weights (xgb's alpha). Default to 0.0. + reg_lambda : Optional[float] + L2 regularization term on weights (xgb's lambda). Default to 1.0. + early_stop: Optional[bool] + Whether training should stop after the first iteration. Default to True. + min_rel_progress: Optional[float] + Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. + enable_global_explain: Optional[bool] + Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. + xgboost_version: Optional[str] + Specifies the Xgboost version for model training. Default to "0.9". Possible values: "0.9", "1.1".ß + """ diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py new file mode 100644 index 00000000000..e67edac1aa4 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -0,0 +1,98 @@ +""" +Generalized Linear Models. +""" + +# Author: Alexandre Gramfort +# Fabian Pedregosa +# Olivier Grisel +# Vincent Michel +# Peter Prettenhofer +# Mathieu Blondel +# Lars Buitinck +# Maryan Morel +# Giorgio Patrini +# Maria Telenczuk +# License: BSD 3 clause +# Original location: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/linear_model/_base.py + +from abc import ABCMeta +from typing import List, Optional + +from third_party.bigframes_vendored.sklearn.base import ( + BaseEstimator, + ClassifierMixin, + RegressorMixin, +) + + +class LinearModel(BaseEstimator, metaclass=ABCMeta): + def predict(self, X): + """Predict using the linear model. + + Args: + X: + DataFrame of shape (n_samples, n_features). Samples. + + Returns: + DataFrame of shape (n_samples,). Returns predicted values. + """ + raise NotImplementedError("abstract method") + + +class LinearClassifierMixin(ClassifierMixin): + def predict(self, X): + """Predict class labels for samples in X. + + Args: + X: + DataFrame of shape (n_samples, n_features). The data matrix for + which we want to get the predictions. + + Returns: + DataFrame of shape (n_samples,), containing the class labels for + each sample. + """ + raise NotImplementedError("abstract method") + + +class LinearRegression(RegressorMixin, LinearModel): + """Ordinary least squares Linear Regression. + + LinearRegression fits a linear model with coefficients w = (w1, ..., wp) + to minimize the residual sum of squares between the observed targets in + the dataset, and the targets predicted by the linear approximation. + + Args: + fit_intercept: + Default ``True``. Whether to calculate the intercept for this + model. If set to False, no intercept will be used in calculations + (i.e. data is expected to be centered). + """ + + def fit( + self, + X, + y, + transforms: Optional[List[str]] = None, + ): + """Fit linear model. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training data. + + y: + DataFrame of shape (n_samples,) or (n_samples, n_targets). + Target values. Will be cast to X's dtype if necessary. + + transforms: + An optional list of SQL expressions to apply over top of the + model inputs as preprocessing. This preprocessing will be + automatically reapplied to new input data (e.g. in .predict), + and may contain steps (like ML.STANDARD_SCALER) that fit to the + training data. + + Returns: + Fitted Estimator. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py new file mode 100644 index 00000000000..6de5bf65e25 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py @@ -0,0 +1,83 @@ +""" +Logistic Regression +""" + +# Author: Gael Varoquaux +# Fabian Pedregosa +# Alexandre Gramfort +# Manoj Kumar +# Lars Buitinck +# Simon Wu +# Arthur Mensch ) + + Hsiang-Fu Yu, Fang-Lan Huang, Chih-Jen Lin (2011). Dual coordinate descent + methods for logistic regression and maximum entropy models. + Machine Learning 85(1-2):41-75. + https://www.csie.ntu.edu.tw/~cjlin/papers/maxent_dual.pdf + """ + + def fit( + self, + X, + y, + transforms: Optional[List[str]] = None, + ): + """Fit the model according to the given training data. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training vector, + where `n_samples` is the number of samples and `n_features` is + the number of features. + + y: + DataFrame of shape (n_samples,). Target vector relative to X. + + transforms: + An optional list of SQL expressions to apply over top of the + model inputs as preprocessing. This preprocessing will be + automatically reapplied to new input data (e.g. in .predict), + and may contain steps (like ML.STANDARD_SCALER) that fit to the + training data. + + Returns: + Fitted estimator. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/metrics/_classification.py b/third_party/bigframes_vendored/sklearn/metrics/_classification.py new file mode 100644 index 00000000000..05fb5f6fb6e --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/metrics/_classification.py @@ -0,0 +1,173 @@ +"""Metrics to assess performance on classification task given class prediction. +Functions named as ``*_score`` return a scalar value to maximize: the higher +the better. +Function named as ``*_error`` or ``*_loss`` return a scalar value to minimize: +the lower the better. +""" + +# Authors: Alexandre Gramfort +# Mathieu Blondel +# Olivier Grisel +# Arnaud Joly +# Jochen Wersdorfer +# Lars Buitinck +# Joel Nothman +# Noel Dawe +# Jatin Shah +# Saurabh Jha +# Bernardo Stein +# Shangwu Yao +# Michal Karbownik +# License: BSD 3 clause + + +def accuracy_score(y_true, y_pred, normalize=True) -> float: + """Accuracy classification score. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) labels. + y_pred: DataFrame of shape (n_samples,) + Predicted labels, as returned by a classifier. + normalize: bool, default=True + Default to True. If ``False``, return the number of correctly + classified samples. Otherwise, return the fraction of correctly + classified samples. + + Returns: + score: float. + If ``normalize == True``, return the fraction of correctly + classified samples (float), else returns the number of correctly + classified samples (int). + """ + raise NotImplementedError("abstract method") + + +def confusion_matrix( + y_true, + y_pred, +): + """Compute confusion matrix to evaluate the accuracy of a classification. + + By definition a confusion matrix :math:`C` is such that :math:`C_{i, j}` + is equal to the number of observations known to be in group :math:`i` and + predicted to be in group :math:`j`. + + Thus in binary classification, the count of true negatives is + :math:`C_{0,0}`, false negatives is :math:`C_{1,0}`, true positives is + :math:`C_{1,1}` and false positives is :math:`C_{0,1}`. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) target values. + y_pred: DataFrame of shape (n_samples,) + Estimated targets as returned by a classifier. + + Returns: + C: DataFrame of shape (n_samples, n_features). + Confusion matrix whose i-th row and j-th + column entry indicates the number of + samples with true label being i-th class + and predicted label being j-th class. + """ + raise NotImplementedError("abstract method") + + +def recall_score( + y_true, + y_pred, + average: str = "binary", +): + """Compute the recall. + + The recall is the ratio ``tp / (tp + fn)`` where ``tp`` is the number of + true positives and ``fn`` the number of false negatives. The recall is + intuitively the ability of the classifier to find all the positive samples. + + The best value is 1 and the worst value is 0. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) target values. + y_pred: DataFrame of shape (n_samples,) + Estimated targets as returned by a classifier. + average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ + default='binary' + This parameter is required for multiclass/multilabel targets. + Possible values are 'None', 'micro', 'macro', 'samples', 'weighted', 'binary'. + + Returns: + recall: float (if average is not None) or Series of float of shape \ + (n_unique_labels,). + Recall of the positive class in binary classification or weighted + average of the recall of each class for the multiclass task. + """ + raise NotImplementedError("abstract method") + + +def precision_score( + y_true, + y_pred, + average: str = "binary", +): + """Compute the precision. + + The precision is the ratio ``tp / (tp + fp)`` where ``tp`` is the number of + true positives and ``fp`` the number of false positives. The precision is + intuitively the ability of the classifier not to label as positive a sample + that is negative. + + The best value is 1 and the worst value is 0. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) target values. + y_pred: DataFrame of shape (n_samples,) + Estimated targets as returned by a classifier. + average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ + default='binary' + This parameter is required for multiclass/multilabel targets. + Possible values are 'None', 'micro', 'macro', 'samples', 'weighted', 'binary'. + + Returns: + precision: float (if average is not None) or Series of float of shape \ + (n_unique_labels,). + Precision of the positive class in binary classification or weighted + average of the precision of each class for the multiclass task. + """ + raise NotImplementedError("abstract method") + + +def f1_score( + y_true, + y_pred, + average: str = "binary", +): + """Compute the F1 score, also known as balanced F-score or F-measure. + + The F1 score can be interpreted as a harmonic mean of the precision and + recall, where an F1 score reaches its best value at 1 and worst score at 0. + The relative contribution of precision and recall to the F1 score are + equal. The formula for the F1 score is: F1 = 2 * (precision * recall) / (precision + recall) + + In the multi-class and multi-label case, this is the average of + the F1 score of each class with weighting depending on the ``average`` + parameter. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) target values. + y_pred: DataFrame of shape (n_samples,) + Estimated targets as returned by a classifier. + average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ + default='binary' + This parameter is required for multiclass/multilabel targets. + Possible values are 'None', 'micro', 'macro', 'samples', 'weighted', 'binary'. + + Returns: + f1_score: float or Series of float, shape = [n_unique_labels] + F1 score of the positive class in binary classification or weighted + average of the F1 scores of each class for the multiclass task. + + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/metrics/_ranking.py b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py new file mode 100644 index 00000000000..b7625171b56 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py @@ -0,0 +1,100 @@ +"""Metrics to assess performance on classification task given scores. +Functions named as ``*_score`` return a scalar value to maximize: the higher +the better. +Function named as ``*_error`` or ``*_loss`` return a scalar value to minimize: +the lower the better. +""" + +# Authors: Alexandre Gramfort +# Mathieu Blondel +# Olivier Grisel +# Arnaud Joly +# Jochen Wersdorfer +# Lars Buitinck +# Joel Nothman +# Noel Dawe +# Michal Karbownik +# License: BSD 3 clause + + +def auc(x, y) -> float: + """Compute Area Under the Curve (AUC) using the trapezoidal rule. + + This is a general function, given points on a curve. For computing the + area under the ROC-curve, see :func:`roc_auc_score`. For an alternative + way to summarize a precision-recall curve, see + :func:`average_precision_score`. + + Args: + x : DataFrame of shape (n_samples,) + X coordinates. These must be either monotonic increasing or monotonic + decreasing. + y : DataFrame of shape (n_samples,) + Y coordinates. + + Returns: + auc : float + Area Under the Curve. + """ + raise NotImplementedError("abstract method") + + +def roc_auc_score(y_true, y_score) -> float: + """Compute Area Under the Receiver Operating Characteristic Curve (ROC AUC) \ + from prediction scores. + + Args: + y_true: DataFrame of shape (n_samples,) + True labels or binary label indicators. The binary and multiclass cases + expect labels with shape (n_samples,) while the multilabel case expects + binary label indicators with shape (n_samples, n_classes). + y_score: DataFrame of shape (n_samples,) + Target scores. + * In the binary case, it corresponds to an array of shape + `(n_samples,)`. Both probability estimates and non-thresholded + decision values can be provided. The probability estimates correspond + to the **probability of the class with the greater label**, + i.e. `estimator.classes_[1]` and thus + `estimator.predict_proba(X, y)[:, 1]`. The decision values + corresponds to the output of `estimator.decision_function(X, y)`. + + Returns: + auc: float. + Area Under the Curve score. + """ + raise NotImplementedError("abstract method") + + +def roc_curve( + y_true, + y_score, + drop_intermediate: bool = True, +): + """Compute Receiver operating characteristic (ROC). + + Args: + y_true: DataFrame of shape (n_samples,) + True binary labels. If labels are not either {-1, 1} or {0, 1}, then + pos_label should be explicitly given. + y_score: DataFrame of shape (n_samples,) + Target scores, can either be probability estimates of the positive + class, confidence values, or non-thresholded measure of decisions + (as returned by "decision_function" on some classifiers). + drop_intermediate: bool, default=True + Default to True. Whether to drop some suboptimal thresholds which would not appear + on a plotted ROC curve. This is useful in order to create lighter + ROC curves. + + Returns: + fpr: + Increasing false positive rates such that element i is the false + positive rate of predictions with score >= `thresholds[i]`. + tpr: + Increasing true positive rates such that element `i` is the true + positive rate of predictions with score >= `thresholds[i]`. + thresholds: + Decreasing thresholds on the decision function used to compute + fpr and tpr. `thresholds[0]` represents no instances being predicted + and is arbitrarily set to `max(y_score) + 1`. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/metrics/_regression.py b/third_party/bigframes_vendored/sklearn/metrics/_regression.py new file mode 100644 index 00000000000..44ceab48a6a --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/metrics/_regression.py @@ -0,0 +1,53 @@ +"""Metrics to assess performance on regression task. +Functions named as ``*_score`` return a scalar value to maximize: the higher +the better. +Function named as ``*_error`` or ``*_loss`` return a scalar value to minimize: +the lower the better. +""" + +# Authors: Alexandre Gramfort +# Mathieu Blondel +# Olivier Grisel +# Arnaud Joly +# Jochen Wersdorfer +# Lars Buitinck +# Joel Nothman +# Karan Desai +# Noel Dawe +# Manoj Kumar +# Michael Eickenberg +# Konstantin Shmelkov +# Christian Lorentzen +# Ashutosh Hathidara +# Uttam kumar +# Sylvain Marie +# Ohad Michel +# License: BSD 3 clause + + +def r2_score(y_true, y_pred, force_finite=True) -> float: + """:math:`R^2` (coefficient of determination) regression score function. + + Best possible score is 1.0 and it can be negative (because the + model can be arbitrarily worse). In the general case when the true y is + non-constant, a constant model that always predicts the average y + disregarding the input features would get a :math:`R^2` score of 0.0. + + In the particular case when ``y_true`` is constant, the :math:`R^2` score + is not finite: it is either ``NaN`` (perfect predictions) or ``-Inf`` + (imperfect predictions). To prevent such non-finite numbers to pollute + higher-level experiments such as a grid search cross-validation, by default + these cases are replaced with 1.0 (perfect predictions) or 0.0 (imperfect + predictions) respectively. + + Args: + y_true: DataFrame of shape (n_samples,) + Ground truth (correct) target values. + y_pred: DataFrame of shape (n_samples,) + Estimated target values. + + Returns: + z: float. + The :math:`R^2` score. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/pipeline.py b/third_party/bigframes_vendored/sklearn/pipeline.py new file mode 100644 index 00000000000..7374edf990a --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/pipeline.py @@ -0,0 +1,83 @@ +""" +The :mod:`sklearn.pipeline` module implements utilities to build a composite +estimator, as a chain of transforms and estimators. +""" +# Author: Edouard Duchesnay +# Gael Varoquaux +# Virgile Fritsch +# Alexandre Gramfort +# Lars Buitinck +# License: BSD + +from abc import ABCMeta + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class Pipeline(BaseEstimator, metaclass=ABCMeta): + """Pipeline of transforms with a final estimator. + + Sequentially apply a list of transforms and a final estimator. + Intermediate steps of the pipeline must be 'transforms', that is, they + must implement `fit` and `transform` methods. + The final estimator only needs to implement `fit`. + + The purpose of the pipeline is to assemble several steps that can be + cross-validated together while setting different parameters. This simplifies code, and allows deploying an estimator + and peprocessing together, e.g. with Pipeline.to_gbq(...) + """ + + def fit( + self, + X, + y, + ): + """Fit the model. + + Fit all the transformers one after the other and transform the + data. Finally, fit the transformed data using the final estimator. + + Args: + X: + A BigQuery DataFrames representing training data. Must match the + input requirements of the first step of the pipeline. + y: + A BigQuery DataFrames representing training targets, if applicable. + + Returns: + Pipeline with fitted steps. + """ + raise NotImplementedError("abstract method") + + +def score(self, X, y): + """Transform the data, and apply `score` with the final estimator. + + Call `transform` of each transformer in the pipeline. The transformed + data are finally passed to the final estimator that calls + `score` method. Only valid if the final estimator implements `score`. + + Args: + X: + A BigQuery DataFrames as evaluation data. + y: + A BigQuery DataFrames as evaluation labels. + + Returns: + A BigQuery DataFrames representing the result of calling + `score` on the final estimator. + """ + raise NotImplementedError("abstract method") + + +def predict(self, X): + """Predict the pipeline result for each sample in X. + + Args: + X: + A BigQuery DataFrames to predict. + + Returns: + A BigQuery DataFrames Dataframe representing predicted result. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py new file mode 100644 index 00000000000..7a44c4e66d3 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -0,0 +1,53 @@ +# Authors: Alexandre Gramfort +# Mathieu Blondel +# Olivier Grisel +# Andreas Mueller +# Eric Martin +# Giorgio Patrini +# Eric Chang +# License: BSD 3 clause + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class StandardScaler(BaseEstimator): + """Standardize features by removing the mean and scaling to unit variance. + + The standard score of a sample `x` is calculated as:z = (x - u) / s + where `u` is the mean of the training samples or zero if `with_mean=False`, + and `s` is the standard deviation of the training samples or one if + `with_std=False`. + + Centering and scaling happen independently on each feature by computing + the relevant statistics on the samples in the training set. Mean and + standard deviation are then stored to be used on later data using + :meth:`transform`. + + Standardization of a dataset is a common requirement for many + machine learning estimators: they might behave badly if the + individual features do not more or less look like standard normally + distributed data (e.g. Gaussian with 0 mean and unit variance). + """ + + def fit(self, X): + """Compute the mean and std to be used for later scaling. + + Args: + X: + A dataframe with training data. + + Returns: + Fitted scaler. + """ + raise NotImplementedError("abstract method") + + def transform(self, X): + """Perform standardization by centering and scaling. + + Args: + X: + The DataFrame to be transformed. + + Returns: + Transformed result.""" + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py new file mode 100644 index 00000000000..6f0565ac136 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -0,0 +1,41 @@ +# Authors: Andreas Mueller +# Joris Van den Bossche +# License: BSD 3 clause + +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class OneHotEncoder(BaseEstimator): + """Encode categorical features as a one-hot format. + + The input to this transformer should be an array-like of integers or + strings, denoting the values taken on by categorical (discrete) features. + The features are encoded using a one-hot (aka 'one-of-K' or 'dummy') + encoding scheme. + + Note that this method deviates from Scikit-Learn; instead of producing sparse + binary columns, the encoding is a single column of STRUCT + """ + + def fit(self, X): + """Fit OneHotEncoder to X. + + Args: + X: + A dataframe with training data. + + Returns: + Fitted encoder. + """ + raise NotImplementedError("abstract method") + + def transform(self, X): + """Transform X using one-hot encoding. + + Args: + X: + The DataFrame to be transformed. + + Returns: + Transformed result.""" + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/xgboost/LICENSE b/third_party/bigframes_vendored/xgboost/LICENSE new file mode 100644 index 00000000000..82832bd6e46 --- /dev/null +++ b/third_party/bigframes_vendored/xgboost/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright (c) 2019 by Contributors + + 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/third_party/bigframes_vendored/xgboost/sklearn.py b/third_party/bigframes_vendored/xgboost/sklearn.py new file mode 100644 index 00000000000..268484b2c23 --- /dev/null +++ b/third_party/bigframes_vendored/xgboost/sklearn.py @@ -0,0 +1,145 @@ +"""Scikit-Learn Wrapper interface for XGBoost.""" + +from typing import Any + +from ..sklearn.base import BaseEstimator as XGBModelBase +from ..sklearn.base import ClassifierMixin as XGBClassifierBase +from ..sklearn.base import RegressorMixin as XGBRegressorBase + + +class XGBModel(XGBModelBase): + def predict(self, X): + """Predict using the XGB model. + + Args: + X: + DataFrame of shape (n_samples, n_features). Samples. + + Returns: + DataFrame of shape (n_samples,). Returns predicted values. + """ + raise NotImplementedError("abstract method") + + def fit(self, X, y): + """Fit gradient boosting model. + + Note that calling ``fit()`` multiple times will cause the model object to be + re-fit from scratch. To resume training from a previous checkpoint, explicitly + pass ``xgb_model`` argument. + + Args: + X: + DataFrame of shape (n_samples, n_features). Training data. + + y: + DataFrame of shape (n_samples,) or (n_samples, n_targets). + Target values. Will be cast to X's dtype if necessary. + + Returns: + Fitted Estimator. + """ + raise NotImplementedError("abstract method") + + +class XGBClassifierMixIn: + """MixIn for classification.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class XGBRegressor(XGBModel, XGBRegressorBase): + """ + XGBoost regression model. + + Args: + num_parallel_tree: Optional[int] + Number of parallel trees constructed during each iteration. Default to 1. + booster: Optional[str] + Specify which booster to use: gbtree or dart. Default to "gbtree". + dart_normalized_type": Optional[str] + Type of normalization algorithm for DART booster. Possible values: "TREE", "FOREST". Default to "TREE". + tree_method: Optional[str] + Specify which tree method to use. Default to "auto". If this parameter is set to + default, XGBoost will choose the most conservative option available. + min_child_weight : Optional[float] + Minimum sum of instance weight(hessian) needed in a child. Default to 1. + colsample_bytree : Optional[float] + Subsample ratio of columns when constructing each tree. Default to 1.0. + colsample_bylevel : Optional[float] + Subsample ratio of columns for each level. Default to 1.0. + colsample_bynode : Optional[float] + Subsample ratio of columns for each split. Default to 1.0. + gamma : Optional[float] + (min_split_loss) Minimum loss reduction required to make a further partition on a + leaf node of the tree. Default to 0.0. + max_depth : Optional[int] + Maximum tree depth for base learners. Default to 6. + subsample : Optional[float] + Subsample ratio of the training instance. Default to 1.0. + reg_alpha : Optional[float] + L1 regularization term on weights (xgb's alpha). Default to 0.0. + reg_lambda : Optional[float] + L2 regularization term on weights (xgb's lambda). Default to 1.0. + early_stop: Optional[bool] + Whether training should stop after the first iteration. Default to True. + learning_rate: Optional[float] + Boosting learning rate (xgb's "eta"). Default to 0.3. + max_iterations: Optional[int] + Maximum number of rounds for boosting. Default to 20. + min_rel_progress: Optional[float] + Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. + enable_global_explain: Optional[bool] + Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. + xgboost_version: Optional[str] + Specifies the Xgboost version for model training. Default to "0.9". + """ + + +class XGBClassifier(XGBModel, XGBClassifierMixIn, XGBClassifierBase): + """ + XGBoost classifier model. + + Args: + num_parallel_tree: Optional[int] + Number of parallel trees constructed during each iteration. Default to 1. + booster: Optional[str] + Specify which booster to use: gbtree or dart. Default to "gbtree". + dart_normalized_type": Optional[str] + Type of normalization algorithm for DART booster. Possible values: "TREE", "FOREST". Default to "TREE". + tree_method: Optional[str] + Specify which tree method to use. Default to "auto". If this parameter is set to + default, XGBoost will choose the most conservative option available. Possible values: ""exact", "approx", + "hist". + min_child_weight : Optional[float] + Minimum sum of instance weight(hessian) needed in a child. Default to 1. + colsample_bytree : Optional[float] + Subsample ratio of columns when constructing each tree. Default to 1.0. + colsample_bylevel : Optional[float] + Subsample ratio of columns for each level. Default to 1.0. + colsample_bynode : Optional[float] + Subsample ratio of columns for each split. Default to 1.0. + gamma : Optional[float] + (min_split_loss) Minimum loss reduction required to make a further partition on a + leaf node of the tree. Default to 0.0. + max_depth : Optional[int] + Maximum tree depth for base learners. Default to 6. + subsample : Optional[float] + Subsample ratio of the training instance. Default to 1.0. + reg_alpha : Optional[float] + L1 regularization term on weights (xgb's alpha). Default to 0.0. + reg_lambda : Optional[float] + L2 regularization term on weights (xgb's lambda). Default to 1.0. + early_stop: Optional[bool] + Whether training should stop after the first iteration. Default to True. + learning_rate: Optional[float] + Boosting learning rate (xgb's "eta"). Default to 0.3. + max_iterations: Optional[int] + Maximum number of rounds for boosting. Default to 20. + min_rel_progress: Optional[float] + Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. + enable_global_explain: Optional[bool] + Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. + xgboost_version: Optional[str] + Specifies the Xgboost version for model training. Default to "0.9". Possible values: "0.9", "1.1". + """ From 264142e54498d50ec7427c8b3e45265f19fcd874 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 11 Aug 2023 14:50:04 -0500 Subject: [PATCH 03/80] chore: release 0.1.0 (#4) Syncs to internal commit 792794841b8ae0ca502723a2abead4f5fd5d11a0 Change-Id: Ib799a9c5e5a18d9b471756410aa5c87cb4932fe8 --- .kokoro/continuous/e2e.cfg | 2 +- .kokoro/continuous/nightly.cfg | 2 +- .kokoro/presubmit/e2e.cfg | 2 +- .kokoro/release-nightly.sh | 5 + .repo-metadata.json | 16 + CHANGELOG.md | 4 +- README.rst | 232 +- bigframes/_config/__init__.py | 11 + bigframes/_config/bigquery_options.py | 21 +- bigframes/_config/display_options.py | 3 +- bigframes/_config/sampling_options.py | 30 + bigframes/constants.py | 23 + bigframes/core/__init__.py | 372 ++- bigframes/core/blocks.py | 428 ++- bigframes/core/groupby/__init__.py | 169 +- bigframes/core/indexers.py | 62 +- bigframes/core/indexes/index.py | 75 +- bigframes/core/io.py | 90 + bigframes/core/joins/row_identity.py | 79 +- bigframes/core/joins/single_column.py | 235 +- bigframes/core/ordering.py | 201 +- bigframes/core/reshape/__init__.py | 105 +- bigframes/core/scalar.py | 29 +- bigframes/core/utils.py | 22 + bigframes/dataframe.py | 649 ++-- bigframes/dtypes.py | 69 +- bigframes/formatting_helpers.py | 207 +- bigframes/ml/base.py | 30 +- bigframes/ml/cluster.py | 69 +- bigframes/ml/compose.py | 72 +- bigframes/ml/core.py | 171 +- bigframes/ml/decomposition.py | 70 +- bigframes/ml/ensemble.py | 249 +- bigframes/ml/forecasting.py | 88 +- bigframes/ml/imported.py | 144 +- bigframes/ml/linear_model.py | 132 +- bigframes/ml/llm.py | 90 +- bigframes/ml/loader.py | 100 +- bigframes/ml/metrics.py | 220 +- bigframes/ml/model_selection.py | 46 +- bigframes/ml/pipeline.py | 144 +- bigframes/ml/preprocessing.py | 143 +- bigframes/ml/sql.py | 36 +- bigframes/ml/utils.py | 58 + bigframes/operations/__init__.py | 145 +- bigframes/operations/aggregations.py | 117 +- bigframes/operations/base.py | 15 +- bigframes/operations/strings.py | 109 +- bigframes/pandas/__init__.py | 99 +- bigframes/remote_function.py | 455 ++- bigframes/series.py | 354 ++- bigframes/session.py | 451 ++- docs/getting_started/index.rst | 27 - docs/index.rst | 4 +- docs/reference/bigframes.ml/cluster.rst | 2 + docs/reference/bigframes.ml/compose.rst | 2 + docs/reference/bigframes.ml/decomposition.rst | 2 + docs/reference/bigframes.ml/ensemble.rst | 7 + docs/reference/bigframes.ml/forecasting.rst | 2 + docs/reference/bigframes.ml/imported.rst | 2 + docs/reference/bigframes.ml/index.rst | 2 + docs/reference/bigframes.ml/linear_model.rst | 2 + docs/reference/bigframes.ml/llm.rst | 2 + docs/reference/bigframes.ml/metrics.rst | 2 + .../bigframes.ml/model_selection.rst | 2 + docs/reference/bigframes.ml/pipeline.rst | 2 + docs/reference/bigframes.ml/preprocessing.rst | 2 + docs/reference/bigframes/options.rst | 2 + .../bigframes.pandas/remote_functions.rst | 105 +- notebooks/01 - Getting Started.ipynb | 231 +- notebooks/02 - DataFrame.ipynb | 2 +- ...Using ML - SKLearn linear regression.ipynb | 1166 ++++--- ... - Using ML - Easy linear regression.ipynb | 1136 ++++--- ...6 - Using ML - Large Language Models.ipynb | 2 +- notebooks/10 - Regionalized.ipynb | 2800 +++++++++++++++++ noxfile.py | 54 +- samples/snippets/quickstart.py | 62 +- scripts/upload_to_google_drive.py | 71 + setup.py | 7 +- testing/constraints-3.9.txt | 4 +- tests/system/conftest.py | 24 + tests/system/large/ml/test_cluster.py | 35 +- tests/system/large/ml/test_compose.py | 8 +- tests/system/large/ml/test_core.py | 16 +- tests/system/large/ml/test_decomposition.py | 34 +- tests/system/large/ml/test_ensemble.py | 25 +- tests/system/large/ml/test_forecasting.py | 2 +- tests/system/large/ml/test_linear_model.py | 8 +- tests/system/large/ml/test_pipeline.py | 301 +- tests/system/large/test_remote_function.py | 20 +- tests/system/small/ml/conftest.py | 70 +- tests/system/small/ml/test_cluster.py | 117 +- tests/system/small/ml/test_core.py | 12 +- tests/system/small/ml/test_decomposition.py | 71 +- tests/system/small/ml/test_ensemble.py | 154 +- tests/system/small/ml/test_forecasting.py | 46 +- tests/system/small/ml/test_imported.py | 22 +- tests/system/small/ml/test_linear_model.py | 90 +- tests/system/small/ml/test_llm.py | 34 +- tests/system/small/ml/test_metrics.py | 278 +- tests/system/small/ml/test_model_selection.py | 43 +- tests/system/small/ml/test_preprocessing.py | 136 +- tests/system/small/ml/test_register.py | 8 +- .../system/small/operations/test_datetimes.py | 20 +- tests/system/small/operations/test_strings.py | 155 +- tests/system/small/test_dataframe.py | 418 ++- tests/system/small/test_dataframe_io.py | 100 +- tests/system/small/test_groupby.py | 54 +- tests/system/small/test_ibis.py | 39 + tests/system/small/test_index.py | 12 +- tests/system/small/test_multiindex.py | 224 +- tests/system/small/test_pandas.py | 57 +- tests/system/small/test_pandas_options.py | 194 ++ tests/system/small/test_progress_bar.py | 97 +- tests/system/small/test_remote_function.py | 313 +- tests/system/small/test_series.py | 622 +++- tests/system/small/test_session.py | 97 +- tests/system/small/test_window.py | 43 +- tests/unit/ml/test_sql.py | 28 +- tests/unit/test_core.py | 78 +- tests/unit/test_dtypes.py | 9 +- tests/unit/test_formatting_helper.py | 17 + tests/unit/test_formatting_helpers.py | 46 + tests/unit/test_pandas.py | 32 +- .../bigframes_vendored/ibis/LICENSE.txt | 202 ++ third_party/bigframes_vendored/ibis/README.md | 196 ++ .../bigframes_vendored/ibis/__init__.py | 0 .../ibis/backends/__init__.py | 0 .../ibis/backends/bigquery/__init__.py | 0 .../ibis/backends/bigquery/registry.py | 19 + .../bigframes_vendored/ibis/expr/__init__.py | 0 .../ibis/expr/operations/__init__.py | 4 + .../ibis/expr/operations/reductions.py | 23 + .../pandas/core/config_init.py | 46 +- .../bigframes_vendored/pandas/core/frame.py | 629 ++-- .../bigframes_vendored/pandas/core/generic.py | 146 +- .../pandas/core/groupby/__init__.py | 235 +- .../pandas/core/indexes/accessor.py | 13 +- .../pandas/core/indexes/base.py | 18 +- .../pandas/core/indexing.py | 30 +- .../pandas/core/reshape/__init__.py | 0 .../pandas/core/reshape/concat.py | 3 + .../pandas/core/reshape/tile.py | 65 + .../bigframes_vendored/pandas/core/series.py | 1185 +++---- .../pandas/core/strings/accessor.py | 183 +- .../bigframes_vendored/pandas/io/common.py | 5 +- .../bigframes_vendored/pandas/io/gbq.py | 12 +- .../bigframes_vendored/pandas/io/parquet.py | 15 +- .../pandas/io/parsers/readers.py | 53 +- .../bigframes_vendored/sklearn/base.py | 22 +- .../sklearn/cluster/_kmeans.py | 60 +- .../sklearn/compose/_column_transformer.py | 12 +- .../sklearn/decomposition/_pca.py | 48 +- .../sklearn/ensemble/_forest.py | 13 +- .../sklearn/linear_model/_base.py | 35 +- .../sklearn/linear_model/_logistic.py | 43 +- .../sklearn/metrics/_classification.py | 45 +- .../sklearn/metrics/_ranking.py | 18 +- .../sklearn/metrics/_regression.py | 7 +- .../bigframes_vendored/sklearn/pipeline.py | 37 +- .../sklearn/preprocessing/_data.py | 33 +- .../sklearn/preprocessing/_encoder.py | 49 +- .../bigframes_vendored/xgboost/sklearn.py | 99 +- 163 files changed, 15353 insertions(+), 4913 deletions(-) create mode 100644 .repo-metadata.json create mode 100644 bigframes/_config/sampling_options.py create mode 100644 bigframes/constants.py create mode 100644 bigframes/core/io.py create mode 100644 bigframes/core/utils.py create mode 100644 bigframes/ml/utils.py delete mode 100644 docs/getting_started/index.rst create mode 100644 docs/reference/bigframes.ml/ensemble.rst create mode 100644 notebooks/10 - Regionalized.ipynb create mode 100644 scripts/upload_to_google_drive.py create mode 100644 tests/system/small/test_ibis.py create mode 100644 tests/unit/test_formatting_helper.py create mode 100644 tests/unit/test_formatting_helpers.py create mode 100644 third_party/bigframes_vendored/ibis/LICENSE.txt create mode 100644 third_party/bigframes_vendored/ibis/README.md create mode 100644 third_party/bigframes_vendored/ibis/__init__.py create mode 100644 third_party/bigframes_vendored/ibis/backends/__init__.py create mode 100644 third_party/bigframes_vendored/ibis/backends/bigquery/__init__.py create mode 100644 third_party/bigframes_vendored/ibis/backends/bigquery/registry.py create mode 100644 third_party/bigframes_vendored/ibis/expr/__init__.py create mode 100644 third_party/bigframes_vendored/ibis/expr/operations/__init__.py create mode 100644 third_party/bigframes_vendored/ibis/expr/operations/reductions.py create mode 100644 third_party/bigframes_vendored/pandas/core/reshape/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/reshape/tile.py diff --git a/.kokoro/continuous/e2e.cfg b/.kokoro/continuous/e2e.cfg index e2ca8bc78db..d875f360603 100644 --- a/.kokoro/continuous/e2e.cfg +++ b/.kokoro/continuous/e2e.cfg @@ -3,5 +3,5 @@ # Only run this nox session. env_vars: { key: "NOX_SESSION" - value: "system_prerelease system_noextras e2e notebook samples" + value: "system_noextras e2e notebook samples" } diff --git a/.kokoro/continuous/nightly.cfg b/.kokoro/continuous/nightly.cfg index 5d1ec3a8f75..ac34c4b0c60 100644 --- a/.kokoro/continuous/nightly.cfg +++ b/.kokoro/continuous/nightly.cfg @@ -2,7 +2,7 @@ env_vars: { key: "NOX_SESSION" - value: "unit unit_prerelease system system_prerelease cover lint lint_setup_py mypy format docs e2e notebook" + value: "unit system cover lint lint_setup_py mypy format docs e2e notebook" } build_file: "bigframes/.kokoro/release-nightly.sh" diff --git a/.kokoro/presubmit/e2e.cfg b/.kokoro/presubmit/e2e.cfg index e2ca8bc78db..d875f360603 100644 --- a/.kokoro/presubmit/e2e.cfg +++ b/.kokoro/presubmit/e2e.cfg @@ -3,5 +3,5 @@ # Only run this nox session. env_vars: { key: "NOX_SESSION" - value: "system_prerelease system_noextras e2e notebook samples" + value: "system_noextras e2e notebook samples" } diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh index 8742a0fc790..e3b6b4d4498 100755 --- a/.kokoro/release-nightly.sh +++ b/.kokoro/release-nightly.sh @@ -211,3 +211,8 @@ gcs_docs () { } gcs_docs + +if ! [ ${DRY_RUN} ]; then + # Copy docs and wheels to Google Drive + python3.10 scripts/upload_to_google_drive.py +fi diff --git a/.repo-metadata.json b/.repo-metadata.json new file mode 100644 index 00000000000..0efaa967d2c --- /dev/null +++ b/.repo-metadata.json @@ -0,0 +1,16 @@ +{ + "name": "bigframes", + "name_pretty": "A unified Python API in BigQuery", + "product_documentation": "https://cloud.google.com/bigquery", + "client_documentation": "https://cloud.google.com/python/docs/reference/bigframes/latest", + "issue_tracker": "https://github.com/googleapis/python-bigquery-dataframes/issues", + "release_level": "preview", + "language": "python", + "library_type": "INTEGRATION", + "repo": "googleapis/python-bigquery-dataframes", + "distribution_name": "bigframes", + "api_id": "bigquery.googleapis.com", + "default_version": "", + "codeowner_team": "@googleapis/api-bigquery-dataframe", + "api_shortname": "bigquery" +} diff --git a/CHANGELOG.md b/CHANGELOG.md index 39353a3f1ba..9d966220bd2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,14 +4,14 @@ [1]: https://pypi.org/project/bigframes/#history -## 0.1.0 (TBD) +## 0.1.0 (2023-08-11) ### Features * Add `bigframes.pandas` package with an API compatible with [pandas](https://pandas.pydata.org/). Supported data sources include: BigQuery SQL queries, BigQuery tables, CSV (local and GCS), Parquet (local - and GCS), and more. + and Cloud Storage), and more. * Add `bigframes.ml` package with an API inspired by [scikit-learn](https://scikit-learn.org/stable/). Train machine learning models and run batch predicition, powered by [BigQuery diff --git a/README.rst b/README.rst index 5f473e11890..6f51dfde23e 100644 --- a/README.rst +++ b/README.rst @@ -4,5 +4,233 @@ BigQuery DataFrames BigQuery DataFrames provides a Pythonic DataFrame and machine learning (ML) API powered by the BigQuery engine. -* ``bigframes.pandas`` provides a pandas-like API for analytics. -* ``bigframes.ml`` provides a Scikit-Learn-like API for ML. +* ``bigframes.pandas`` provides a pandas-compatible API for analytics. +* ``bigframes.ml`` provides a scikit-learn-like API for ML. + +Documentation +------------- + +* `BigQuery DataFrames sample notebooks `_ +* `BigQuery DataFrames API reference `_ +* `BigQuery documentation `_ + + +Quickstart +---------- + +Prerequisites +^^^^^^^^^^^^^ + +* Install the ``bigframes`` package. +* Create a Google Cloud project and billing account. +* When running locally, authenticate with application default credentials. See + the `gcloud auth application-default login + `_ + reference. + +Code sample +^^^^^^^^^^^ + +Import ``bigframes.pandas`` for a pandas-like interface. The ``read_gbq`` +method accepts either a fully-qualified table ID or a SQL query. + +.. code-block:: python + + import bigframes.pandas as bpd + + df1 = bpd.read_gbq("project.dataset.table") + df2 = bpd.read_gbq("SELECT a, b, c, FROM `project.dataset.table`") + +* `More code samples `_ + + +Locations +--------- +BigQuery DataFrames uses a +`BigQuery session `_ +internally to manage metadata on the service side. This session is tied to a +`location `_ . +BigQuery DataFrames uses the US multi-region as the default location, but you +can use ``session_options.location`` to set a different location. Every query +in a session is executed in the location where the session was created. + +If you want to reset the location of the created DataFrame or Series objects, +can reset the session by executing ``bigframes.pandas.reset_session()``. +After that, you can reuse ``bigframes.pandas.options.bigquery.location`` to +specify another location. + + +``read_gbq()`` requires you to specify a location if the dataset you are +querying is not in the US multi-region. If you try to read a table from another +location, you get a NotFound exception. + + +ML locations +------------ + +``bigframes.ml`` supports the same locations as BigQuery ML. BigQuery ML model +prediction and other ML functions are supported in all BigQuery regions. Support +for model training varies by region. For more information, see +`BigQuery ML locations `_. + + +Data types +---------- + +BigQuery DataFrames supports the following numpy and pandas dtypes: + +* ``numpy.dtype("O")`` +* ``pandas.BooleanDtype()`` +* ``pandas.Float64Dtype()`` +* ``pandas.Int64Dtype()`` +* ``pandas.StringDtype(storage="pyarrow")`` +* ``pandas.ArrowDtype(pa.date32())`` +* ``pandas.ArrowDtype(pa.time64("us"))`` +* ``pandas.ArrowDtype(pa.timestamp("us"))`` +* ``pandas.ArrowDtype(pa.timestamp("us", tz="UTC"))`` + +BigQuery DataFrames doesn’t support the following BigQuery data types: + +* ``ARRAY`` +* ``NUMERIC`` +* ``BIGNUMERIC`` +* ``INTERVAL`` +* ``STRUCT`` +* ``JSON`` + +All other BigQuery data types display as the object type. + + +Remote functions +---------------- + +BigQuery DataFrames gives you the ability to turn your custom scalar functions +into `BigQuery remote functions +`_ . Creating a remote +function in BigQuery DataFrames creates a BigQuery remote function, a `BigQuery +connection +`_ , +and a `Cloud Functions (2nd gen) function +`_ . + +BigQuery connections are created in the same location as the BigQuery +DataFrames session, using the name you provide in the custom function +definition. To view and manage connections, do the following: + +1. Go to `BigQuery Studio `__. +2. Select the project in which you created the remote function. +3. In the Explorer pane, expand that project and then expand External connections. + +BigQuery remote functions are created in the dataset you specify, or +in a dataset with the name ``bigframes_temp_location``, where location is +the location used by the BigQuery DataFrames session. For example, +``bigframes_temp_us_central1``. To view and manage remote functions, do +the following: + +1. Go to `BigQuery Studio `__. +2. Select the project in which you created the remote function. +3. In the Explorer pane, expand that project, expand the dataset in which you + created the remote function, and then expand Routines. + +To view and manage Cloud Functions functions, use the +`Functions `_ +page and use the project picker to select the project in which you +created the function. For easy identification, the names of the functions +created by BigQuery DataFrames are prefixed by ``bigframes-``. + +**Requirements** + +BigQuery DataFrames uses the ``gcloud`` command-line interface internally, +so you must run ``gcloud auth login`` before using remote functions. + +To use BigQuery DataFrames remote functions, you must enable the following APIs: + +* The BigQuery API (bigquery.googleapis.com) +* The BigQuery Connection API (bigqueryconnection.googleapis.com) +* The Cloud Functions API (cloudfunctions.googleapis.com) +* The Cloud Run API (run.googleapis.com) +* The Artifact Registry API (artifactregistry.googleapis.com) +* The Cloud Build API (cloudbuild.googleapis.com ) +* The Cloud Resource Manager API (cloudresourcemanager.googleapis.com) + +To use BigQuery DataFrames remote functions, you must be granted the +following IAM roles: + +* BigQuery Data Editor (roles/bigquery.dataEditor) +* BigQuery Connection Admin (roles/bigquery.connectionAdmin) +* Cloud Functions Developer (roles/cloudfunctions.developer) +* Service Account User (roles/iam.serviceAccountUser) +* Storage Object Viewer (roles/storage.objectViewer) +* Project IAM Admin (roles/resourcemanager.projectIamAdmin) + +**Limitations** + +* Remote functions take about 90 seconds to become available when you first create them. +* Trivial changes in the notebook, such as inserting a new cell or renaming a variable, + might cause the remote function to be re-created, even if these changes are unrelated + to the remote function code. +* BigQuery DataFrames does not differentiate any personal data you include in the remote + function code. The remote function code is serialized as an opaque box to deploy it as a + Cloud Functions function. +* The Cloud Functions (2nd gen) functions, BigQuery connections, and BigQuery remote + functions created by BigQuery DataFrames persist in Google Cloud. If you don’t want to + keep these resources, you must delete them separately using an appropriate Cloud Functions + or BigQuery interface. +* A project can have up to 1000 Cloud Functions (2nd gen) functions at a time. See Cloud + Functions quotas for all the limits. + + +Quotas and limits +----------------- + +`BigQuery quotas `_ +including hardware, software, and network components. + + +Session termination +------------------- + +Each BigQuery DataFrames DataFrame or Series object is tied to a BigQuery +DataFrames session, which is in turn based on a BigQuery session. BigQuery +sessions +`auto-terminate `_ +; when this happens, you can’t use previously +created DataFrame or Series objects and must re-create them using a new +BigQuery DataFrames session. You can do this by running +``bigframes.pandas.reset_session()`` and then re-running the BigQuery +DataFrames expressions. + + +Data processing location +------------------------ + +BigQuery DataFrames is designed for scale, which it achieves by keeping data +and processing on the BigQuery service. However, you can bring data into the +memory of your client machine by calling ``.execute()`` on a DataFrame or Series +object. If you choose to do this, the memory limitation of your client machine +applies. + + +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 `_ + +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/bigframes/_config/__init__.py b/bigframes/_config/__init__.py index 54c529fb0ca..e26eaf88009 100644 --- a/bigframes/_config/__init__.py +++ b/bigframes/_config/__init__.py @@ -19,6 +19,7 @@ import bigframes._config.bigquery_options as bigquery_options import bigframes._config.display_options as display_options +import bigframes._config.sampling_options as sampling_options class Options: @@ -27,6 +28,7 @@ class Options: def __init__(self): self._bigquery_options = bigquery_options.BigQueryOptions() self._display_options = display_options.DisplayOptions() + self._sampling_options = sampling_options.SamplingOptions() @property def bigquery(self) -> bigquery_options.BigQueryOptions: @@ -38,6 +40,15 @@ def display(self) -> display_options.DisplayOptions: """Options controlling object representation.""" return self._display_options + @property + def sampling(self) -> sampling_options.SamplingOptions: + """Options controlling downsampling when downloading data + to memory. The data will be downloaded into memory explicitly + (e.g., to_pandas, to_numpy, values) or implicitly (e.g., + matplotlib plotting). This option can be overriden by + parameters in specific functions.""" + return self._sampling_options + options = Options() """Global options for default session.""" diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index 5eb5ba7bc85..a103abe1903 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -21,11 +21,14 @@ import google.api_core.exceptions import google.auth.credentials -SESSION_STARTED_MESSAGE = "Cannot change '{attribute}' once a session has started." +SESSION_STARTED_MESSAGE = ( + "Cannot change '{attribute}' once a session has started. " + "Call bigframes.pandas.reset_session() first, if you are using the bigframes.pandas API." +) class BigQueryOptions: - """Encapsulates configuration for working with an Session.""" + """Encapsulates configuration for working with a session.""" def __init__( self, @@ -55,7 +58,7 @@ def credentials(self, value: Optional[google.auth.credentials.Credentials]): @property def location(self) -> Optional[str]: - """Default location for jobs / datasets / tables. + """Default location for job, datasets, and tables. See: https://cloud.google.com/bigquery/docs/locations """ @@ -69,7 +72,7 @@ def location(self, value: Optional[str]): @property def project(self) -> Optional[str]: - """Google Cloud project ID to use for billing and default data project.""" + """Google Cloud project ID to use for billing and as the default project.""" return self._project @project.setter @@ -80,10 +83,12 @@ def project(self, value: Optional[str]): @property def remote_udf_connection(self) -> Optional[str]: - """Name of the BigQuery connection for the purpose of remote UDFs. + """Name of the BigQuery connection to use for remote functions. - It should be either pre created in `location`, or the user should have - privilege to create one. + You should either have the connection already created in the + location you have chosen, or you should have the Project IAM + Admin role to enable the service to create the connection for you if you + need it. """ return self._remote_udf_connection @@ -97,7 +102,7 @@ def remote_udf_connection(self, value: Optional[str]): @property def use_regional_endpoints(self) -> bool: - """In preview. Flag to connect to regional API endpoints. + """Flag to connect to regional API endpoints. Requires ``location`` to also be set. For example, set ``location='asia-northeast1'`` and ``use_regional_endpoints=True`` to diff --git a/bigframes/_config/display_options.py b/bigframes/_config/display_options.py index 43faad0e785..8bd2743f175 100644 --- a/bigframes/_config/display_options.py +++ b/bigframes/_config/display_options.py @@ -16,7 +16,7 @@ import contextlib import dataclasses -from typing import Optional +from typing import Literal, Optional import pandas as pd @@ -30,6 +30,7 @@ class DisplayOptions: max_columns: int = 20 max_rows: int = 25 progress_bar: Optional[str] = "auto" + repr_mode: Literal["head", "deferred"] = "head" @contextlib.contextmanager diff --git a/bigframes/_config/sampling_options.py b/bigframes/_config/sampling_options.py new file mode 100644 index 00000000000..1742dabe17a --- /dev/null +++ b/bigframes/_config/sampling_options.py @@ -0,0 +1,30 @@ +# 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. + +"""Options for downsampling.""" + +import dataclasses +from typing import Literal, Optional + +import third_party.bigframes_vendored.pandas.core.config_init as vendored_pandas_config + + +@dataclasses.dataclass +class SamplingOptions: + __doc__ = vendored_pandas_config.sampling_options_doc + + max_download_size: Optional[int] = 500 + enable_downsampling: bool = False + sampling_method: Literal["head", "uniform"] = "uniform" + random_state: Optional[int] = None diff --git a/bigframes/constants.py b/bigframes/constants.py new file mode 100644 index 00000000000..3f3f1557336 --- /dev/null +++ b/bigframes/constants.py @@ -0,0 +1,23 @@ +# 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. + +"""Constants used across BigQuery DataFrames. + +This module should not depend on any others in the package. +""" + +FEEDBACK_LINK = ( + "Share your usecase with the BigQuery DataFrames team at the " + "https://bit.ly/bigframes-feedback survey." +) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 70565fddf56..075e27b0c23 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -25,11 +25,15 @@ import ibis.expr.types as ibis_types import pandas +import bigframes.constants as constants import bigframes.core.guid from bigframes.core.ordering import ( + encode_order_string, ExpressionOrdering, + IntegerEncoding, OrderingColumnReference, - stringify_order_id, + reencode_order_string, + StringEncoding, ) import bigframes.dtypes import bigframes.operations as ops @@ -85,27 +89,22 @@ def __init__( self, session: Session, table: ibis_types.Table, - columns: Optional[Sequence[ibis_types.Value]] = None, + columns: Sequence[ibis_types.Value], hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None, - ordering: Optional[ExpressionOrdering] = None, + ordering: ExpressionOrdering = ExpressionOrdering(), predicates: Optional[Collection[ibis_types.BooleanValue]] = None, ): self._session = session self._table = table self._predicates = tuple(predicates) if predicates is not None else () # TODO: Validate ordering - self._ordering = ordering or ExpressionOrdering() + if not ordering.total_ordering_columns: + raise ValueError("Must have total ordering defined by one or more columns") + self._ordering = ordering # Allow creating a DataFrame directly from an Ibis table expression. - if columns is None: - self._columns = tuple( - table[key] - for key in table.columns - if ordering is None or key != ordering.ordering_id - ) - else: - # TODO(swast): Validate that each column references the same table (or - # no table for literal values). - self._columns = tuple(columns) + # TODO(swast): Validate that each column references the same table (or + # no table for literal values). + self._columns = tuple(columns) # Meta columns store ordering, or other data that doesn't correspond to dataframe columns self._hidden_ordering_columns = ( @@ -120,6 +119,20 @@ def __init__( self._hidden_ordering_column_names = { column.get_name(): column for column in self._hidden_ordering_columns } + ### Validation + value_col_ids = self._column_names.keys() + hidden_col_ids = self._hidden_ordering_column_names.keys() + + all_columns = value_col_ids | hidden_col_ids + ordering_valid = all( + col.column_id in all_columns for col in ordering.all_ordering_columns + ) + if value_col_ids & hidden_col_ids: + raise ValueError( + f"Keys in both hidden and exposed list: {value_col_ids & hidden_col_ids}" + ) + if not ordering_valid: + raise ValueError(f"Illegal ordering keys: {ordering.all_ordering_columns}") @classmethod def mem_expr_from_pandas( @@ -146,7 +159,8 @@ def mem_expr_from_pandas( session, # type: ignore # Session cannot normally be none, see "caution" above keys_memtable, ordering=ExpressionOrdering( - ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN) + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), ), hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), ) @@ -155,10 +169,6 @@ def mem_expr_from_pandas( def table(self) -> ibis_types.Table: return self._table - @property - def predicates(self) -> typing.Tuple[ibis_types.BooleanValue, ...]: - return self._predicates - @property def reduced_predicate(self) -> typing.Optional[ibis_types.BooleanValue]: """Returns the frame's predicates as an equivalent boolean value, useful where a single predicate value is preferred.""" @@ -181,18 +191,12 @@ def hidden_ordering_columns(self) -> typing.Tuple[ibis_types.Value, ...]: return self._hidden_ordering_columns @property - def ordering(self) -> Sequence[ibis_types.Value]: + def _ibis_order(self) -> Sequence[ibis_types.Value]: """Returns a sequence of ibis values which can be directly used to order a table expression. Has direction modifiers applied.""" - if not self._ordering: - return [] - else: - # TODO(swast): When we assign literals / scalars, we might not - # have a true Column. Do we need to check this before trying to - # sort by such a column? - return _convert_ordering_to_table_values( - {**self._column_names, **self._hidden_ordering_column_names}, - self._ordering.all_ordering_columns, - ) + return _convert_ordering_to_table_values( + {**self._column_names, **self._hidden_ordering_column_names}, + self._ordering.all_ordering_columns, + ) def builder(self) -> ArrayValueBuilder: """Creates a mutable builder for expressions.""" @@ -202,17 +206,12 @@ def builder(self) -> ArrayValueBuilder: return ArrayValueBuilder( self._session, self._table, - self._columns, - self._hidden_ordering_columns, + columns=self._columns, + hidden_ordering_columns=self._hidden_ordering_columns, ordering=self._ordering, predicates=self._predicates, ) - def insert_column(self, index: int, column: ibis_types.Value) -> ArrayValue: - expr = self.builder() - expr.columns.insert(index, column) - return expr.build() - def drop_columns(self, columns: Iterable[str]) -> ArrayValue: # Must generate offsets if we are dropping a column that ordering depends on expr = self @@ -229,7 +228,9 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue: return expr_builder.build() def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: - ibis_type = typing.cast(bigframes.dtypes.IbisDtype, self.get_column(key).type()) + ibis_type = typing.cast( + bigframes.dtypes.IbisDtype, self.get_any_column(key).type() + ) return typing.cast( bigframes.dtypes.Dtype, bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), @@ -286,8 +287,7 @@ def apply_limit(self, max_results: int) -> ArrayValue: def filter(self, predicate: ibis_types.BooleanValue) -> ArrayValue: """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" expr = self.builder() - if expr.ordering: - expr.ordering = expr.ordering.with_is_sequential(False) + expr.ordering = expr.ordering.with_non_sequential() expr.predicates = [*self._predicates, predicate] return expr.build() @@ -303,13 +303,33 @@ def reversed(self) -> ArrayValue: expr_builder.ordering = self._ordering.with_reverse() return expr_builder.build() + def _uniform_sampling(self, fraction: float) -> ArrayValue: + table = self.to_ibis_expr( + ordering_mode="order_by", expose_hidden_cols=True, fraction=fraction + ) + columns = [table[column_name] for column_name in self._column_names] + hidden_ordering_columns = [ + table[column_name] for column_name in self._hidden_ordering_column_names + ] + return ArrayValue( + self._session, + table, + columns=columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=self._ordering, + ) + @property def offsets(self): if not self._ordering.is_sequential: raise ValueError( "Expression does not have offsets. Generate them first using project_offsets." ) - return self._get_hidden_ordering_column(self._ordering.ordering_id) + if not self._ordering.total_order_col: + raise ValueError( + "Ordering is invalid. Marked as sequential but no total order columns." + ) + return self.get_any_column(self._ordering.total_order_col.column_id) def project_offsets(self) -> ArrayValue: """Create a new expression that contains offsets. Should only be executed when offsets are needed for an operations. Has no effect on expression semantics.""" @@ -321,8 +341,9 @@ def project_offsets(self) -> ArrayValue: ) columns = [table[column_name] for column_name in self._column_names] ordering = ExpressionOrdering( - ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN), - is_sequential=True, + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(True, is_sequential=True), ) return ArrayValue( self._session, @@ -342,13 +363,7 @@ def _hide_column(self, column_id) -> ArrayValue: *self._hidden_ordering_columns, self.get_column(column_id).name(new_name), ] - - ordering_columns = [ - col if col.column_id != column_id else col.with_name(new_name) - for col in self._ordering.ordering_value_columns - ] - - expr_builder.ordering = self._ordering.with_ordering_columns(ordering_columns) + expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) return expr_builder.build() def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: @@ -358,12 +373,12 @@ def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: # Special case: offsets already exist ordering = self._ordering - if (not ordering.is_sequential) or (not ordering.ordering_id): + if (not ordering.is_sequential) or (not ordering.total_order_col): return self.project_offsets().promote_offsets() col_id = bigframes.core.guid.generate_guid() expr_builder = self.builder() expr_builder.columns = [ - self._get_hidden_ordering_column(ordering.ordering_id).name(col_id), + self.get_any_column(ordering.total_order_col.column_id).name(col_id), *self.columns, ] return expr_builder.build(), col_id @@ -409,13 +424,13 @@ def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: prefix_size = math.ceil(math.log(len(other) + 1, prefix_base)) # Must normalize all ids to the same encoding size max_encoding_size = max( - self._ordering.ordering_encoding_size, - *[expression._ordering.ordering_encoding_size for expression in other], + self._ordering.string_encoding.length, + *[expression._ordering.string_encoding.length for expression in other], ) for i, expr in enumerate([self, *other]): ordering_prefix = str(i).zfill(prefix_size) table = expr.to_ibis_expr( - ordering_mode="ordered_col", order_col_name=ORDER_ID_COLUMN + ordering_mode="string_encoded", order_col_name=ORDER_ID_COLUMN ) # Rename the value columns based on horizontal offset before applying union. table = table.select( @@ -424,7 +439,9 @@ def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: if col != ORDER_ID_COLUMN else ( ordering_prefix - + stringify_order_id(table[ORDER_ID_COLUMN], max_encoding_size) + + reencode_order_string( + table[ORDER_ID_COLUMN], max_encoding_size + ) ).name(ORDER_ID_COLUMN) for i, col in enumerate(table.columns) ] @@ -432,8 +449,9 @@ def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: tables.append(table) combined_table = ibis.union(*tables) ordering = ExpressionOrdering( - ordering_id_column=OrderingColumnReference(ORDER_ID_COLUMN), - ordering_encoding_size=prefix_size + max_encoding_size, + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + string_encoding=StringEncoding(True, prefix_size + max_encoding_size), ) return ArrayValue( self._session, @@ -498,7 +516,7 @@ def aggregate( aggregations: input_column_id, operation, output_column_id tuples dropna: whether null keys should be dropped """ - table = self.to_ibis_expr() + table = self.to_ibis_expr(ordering_mode="unordered") stats = { col_out: agg_op._as_ibis(table[col_in]) for col_in, agg_op, col_out in aggregations @@ -510,9 +528,11 @@ def aggregate( [ OrderingColumnReference(column_id=column_id) for column_id in by_column_ids - ] + ], + total_ordering_columns=frozenset(by_column_ids), ) - expr = ArrayValue(self._session, result, ordering=ordering) + columns = tuple(result[key] for key in result.columns) + expr = ArrayValue(self._session, result, columns=columns, ordering=ordering) if dropna: for column_id in by_column_ids: expr = expr.filter( @@ -525,8 +545,9 @@ def aggregate( result = table.aggregate(**aggregates) # Ordering is irrelevant for single-row output, but set ordering id regardless as other ops(join etc.) expect it. ordering = ExpressionOrdering( - ordering_id_column=OrderingColumnReference(column_id=ORDER_ID_COLUMN), - is_sequential=True, + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), ) return ArrayValue( self._session, @@ -589,10 +610,12 @@ def project_window_op( def to_ibis_expr( self, ordering_mode: Literal[ - "order_by", "ordered_col", "offset_col", "unordered" + "order_by", "string_encoded", "offset_col", "unordered" ] = "order_by", order_col_name: Optional[str] = ORDER_ID_COLUMN, expose_hidden_cols: bool = False, + fraction: Optional[float] = None, + col_id_overrides: typing.Mapping[str, str] = {}, ): """ Creates an Ibis table expression representing the DataFrame. @@ -604,8 +627,7 @@ def to_ibis_expr( column, however there will be an order_by clause applied to the ouput. * "offset_col": Zero-based offsets are generated as a column, this will not sort the rows however. - * "ordered_col": An ordered column is provided in output table, without - guarantee that the values are sequential + * "string_encoded": An ordered string column is provided in output table. * "unordered": No ordering information will be provided in output. Only value columns are projected. @@ -624,12 +646,14 @@ def to_ibis_expr( If True, include the hidden ordering columns in the results. Only compatible with `order_by` and `unordered` ``ordering_mode``. + col_id_overrides: + overrides the column ids for the result Returns: An ibis expression representing the data help by the ArrayValue object. """ assert ordering_mode in ( "order_by", - "ordered_col", + "string_encoded", "offset_col", "unordered", ) @@ -638,87 +662,107 @@ def to_ibis_expr( f"Cannot expose hidden ordering columns with ordering_mode {ordering_mode}" ) - table = self._table columns = list(self._columns) - hidden_ordering_columns = [ - col.column_id - for col in self._ordering.all_ordering_columns - if col.column_id not in self._column_names.keys() - ] + columns_to_drop: list[ + str + ] = [] # Ordering/Filtering columns that will be dropped at end if self.reduced_predicate is not None: columns.append(self.reduced_predicate) - if ordering_mode in ("offset_col", "ordered_col"): - # Generate offsets if current ordering id semantics are not sufficiently strict - if (ordering_mode == "offset_col" and not self._ordering.is_sequential) or ( - ordering_mode == "ordered_col" and not self._ordering.order_id_defined - ): - window = ibis.window(order_by=self.ordering) - if self._predicates: - window = window.group_by(self.reduced_predicate) - columns.append(ibis.row_number().name(order_col_name).over(window)) - elif self._ordering.ordering_id: - columns.append( - self._get_hidden_ordering_column(self._ordering.ordering_id).name( - order_col_name - ) - ) - else: - # Should not be possible. - raise ValueError( - "Expression does not have ordering id and none was generated." - ) - elif ordering_mode == "order_by": - columns.extend( - [ - self._get_hidden_ordering_column(name) - for name in hidden_ordering_columns - ] - ) + # Usually drop predicate as it is will be all TRUE after filtering + if not expose_hidden_cols: + columns_to_drop.append(self.reduced_predicate.get_name()) - # We already need to add the hidden ordering columns for "order_by" so - # we can order by them. - if expose_hidden_cols and ordering_mode != "order_by": - columns.extend( - [ - self._get_hidden_ordering_column(name) - for name in hidden_ordering_columns - ] - ) + order_columns = self._create_order_columns( + ordering_mode, order_col_name, expose_hidden_cols + ) + columns.extend(order_columns) + if (ordering_mode == "order_by") and not expose_hidden_cols: + columns_to_drop.extend(col.get_name() for col in order_columns) # Special case for empty tables, since we can't create an empty # projection. if not columns: return ibis.memtable([]) - table = table.select(columns) + # Make sure all dtypes are the "canonical" ones for BigFrames. This is # important for operations like UNION where the schema must match. - table = bigframes.dtypes.ibis_table_to_canonical_types(table) - + table = self._table.select( + bigframes.dtypes.ibis_value_to_canonical_type(column) for column in columns + ) + base_table = table if self.reduced_predicate is not None: - table = table.filter(table[PREDICATE_COLUMN]) - # Drop predicate as it is will be all TRUE after filtering - table = table.drop(PREDICATE_COLUMN) + table = table.filter(base_table[PREDICATE_COLUMN]) if ordering_mode == "order_by": - # Some ordering columns are value columns, while other are used purely for ordering. - # We drop the non-value columns after the ordering table = table.order_by( _convert_ordering_to_table_values( - {col: table[col] for col in table.columns}, + {col: base_table[col] for col in table.columns}, self._ordering.all_ordering_columns, ) # type: ignore ) - # TODO(swast): We should be able to avoid this subquery by ordering - # by columns that don't have to be in the SELECT clause. - if not expose_hidden_cols: - table = table.drop(*hidden_ordering_columns) - + table = table.drop(*columns_to_drop) + if col_id_overrides: + table = table.relabel(col_id_overrides) + if fraction is not None: + table = table.filter(ibis.random() < ibis.literal(fraction)) return table + def _create_order_columns( + self, + ordering_mode: str, + order_col_name: Optional[str], + expose_hidden_cols: bool, + ) -> typing.Sequence[ibis_types.Value]: + # Generate offsets if current ordering id semantics are not sufficiently strict + if ordering_mode == "offset_col": + return (self._create_offset_column().name(order_col_name),) + elif ordering_mode == "string_encoded": + return (self._create_string_ordering_column().name(order_col_name),) + elif ordering_mode == "order_by" or expose_hidden_cols: + return self.hidden_ordering_columns + return () + + def _create_offset_column(self) -> ibis_types.IntegerColumn: + if self._ordering.total_order_col and self._ordering.is_sequential: + offsets = self.get_any_column(self._ordering.total_order_col.column_id) + return typing.cast(ibis_types.IntegerColumn, offsets) + else: + window = ibis.window(order_by=self._ibis_order) + if self._predicates: + window = window.group_by(self.reduced_predicate) + offsets = ibis.row_number().over(window) + return typing.cast(ibis_types.IntegerColumn, offsets) + + def _create_string_ordering_column(self) -> ibis_types.StringColumn: + if self._ordering.total_order_col and self._ordering.is_string_encoded: + string_order_ids = self.get_any_column( + self._ordering.total_order_col.column_id + ) + return typing.cast(ibis_types.StringColumn, string_order_ids) + if ( + self._ordering.total_order_col + and self._ordering.integer_encoding.is_encoded + ): + # Special case: non-negative integer ordering id can be converted directly to string without regenerating row numbers + int_values = self.get_any_column(self._ordering.total_order_col.column_id) + return encode_order_string( + typing.cast(ibis_types.IntegerColumn, int_values), + ) + else: + # Have to build string from scratch + window = ibis.window(order_by=self._ibis_order) + if self._predicates: + window = window.group_by(self.reduced_predicate) + row_nums = typing.cast( + ibis_types.IntegerColumn, ibis.row_number().over(window) + ) + return encode_order_string(row_nums) + def start_query( self, job_config: Optional[bigquery.job.QueryJobConfig] = None, max_results: Optional[int] = None, + expose_extra_columns: bool = False, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """Execute a query and return metadata about the results.""" # TODO(swast): Cache the job ID so we can look it up again if they ask @@ -731,7 +775,7 @@ def start_query( # a LocalSession for unit testing. # TODO(swast): Add a timeout here? If the query is taking a long time, # maybe we just print the job metadata that we have so far? - table = self.to_ibis_expr() + table = self.to_ibis_expr(expose_hidden_cols=expose_extra_columns) sql = self._session.ibis_client.compile(table) # type:ignore return self._session._start_query( sql=sql, @@ -739,6 +783,9 @@ def start_query( max_results=max_results, ) + def _get_table_size(self, destination_table): + return self._session._get_table_size(destination_table) + def _reproject_to_table(self) -> ArrayValue: """ Internal operators that projects the internal representation into a @@ -749,12 +796,16 @@ def _reproject_to_table(self) -> ArrayValue: """ table = self.to_ibis_expr( ordering_mode="unordered", - order_col_name=self._ordering.ordering_id, expose_hidden_cols=True, ) columns = [table[column_name] for column_name in self._column_names] + ordering_col_ids = [ + ref.column_id for ref in self._ordering.all_ordering_columns + ] hidden_ordering_columns = [ - table[column_name] for column_name in self._hidden_ordering_column_names + table[column_name] + for column_name in self._hidden_ordering_column_names + if column_name in ordering_col_ids ] return ArrayValue( self._session, @@ -782,10 +833,10 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal ) if not allow_ties: # Most operator need an unambiguous ordering, so the table's total ordering is appended - order_by = tuple([*order_by, *self.ordering]) + order_by = tuple([*order_by, *self._ibis_order]) elif (window_spec.following is not None) or (window_spec.preceding is not None): # If window spec has following or preceding bounds, we need to apply an unambiguous ordering. - order_by = tuple(self.ordering) + order_by = tuple(self._ibis_order) else: # Unbound grouping window. Suitable for aggregations but not for analytic function application. order_by = None @@ -796,47 +847,56 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal group_by=group_by, ) - def transpose_single_row( + def unpivot_single_row( self, - labels, + row_labels: typing.Sequence[typing.Optional[str]], + unpivot_columns: typing.Sequence[typing.Tuple[str, typing.Sequence[str]]], *, index_col_id: str = "index", - value_col_id: str = "values", dtype=pandas.Float64Dtype(), ) -> ArrayValue: - """Pivot a single row into a 3 column expression with index, values and offsets. Only works if all values can be cast to a common type.""" + """Unpivot a single row.""" + # TODO: Generalize to multiple row input table = self.to_ibis_expr(ordering_mode="unordered") sub_expressions = [] - for i, col_id in enumerate(self._column_names.keys()): + + # TODO: validate all columns are equal length, as well as row labels + row_n = len(row_labels) + if not all( + len(source_columns) == row_n for _, source_columns in unpivot_columns + ): + raise ValueError("Columns and row labels must all be same length.") + + # Select each column + for i in range(row_n): + values = [] + for result_col, source_cols in unpivot_columns: + values.append( + ops.AsTypeOp(dtype)._as_ibis(table[source_cols[i]]).name(result_col) + ) + sub_expr = table.select( - ibis_types.literal(labels[i]).name(index_col_id), - ops.AsTypeOp(dtype)._as_ibis(table[col_id]).name(value_col_id), + ibis_types.literal(row_labels[i]).name(index_col_id), + *values, ibis_types.literal(i).name(ORDER_ID_COLUMN), ) sub_expressions.append(sub_expr) rotated_table = ibis.union(*sub_expressions) + + value_columns = [ + rotated_table[value_col_id] for value_col_id, _ in unpivot_columns + ] return ArrayValue( session=self._session, table=rotated_table, - columns=[rotated_table[index_col_id], rotated_table[value_col_id]], + columns=[rotated_table[index_col_id], *value_columns], hidden_ordering_columns=[rotated_table[ORDER_ID_COLUMN]], ordering=ExpressionOrdering( - ordering_id_column=OrderingColumnReference(column_id=ORDER_ID_COLUMN), + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), ), ) - # TODO(b/282041134) Remove deprecate_rename_column once label/id separation in dataframe - def deprecated_rename_column(self, old_id, new_id) -> ArrayValue: - """ - Don't use this, temporary measure until dataframe supports sqlid!=dataframe col id. - In future, caller shouldn't need to control internal column id strings. - """ - if new_id == old_id: - return self - return self._set_or_replace_by_id(new_id, self.get_column(old_id)).drop_columns( - [old_id] - ) - def assign(self, source_id: str, destination_id: str) -> ArrayValue: return self._set_or_replace_by_id(destination_id, self.get_column(source_id)) @@ -850,12 +910,20 @@ def assign_constant( ibis_value = bigframes.dtypes.literal_to_ibis_scalar(value, dtype) if ibis_value is None: raise NotImplementedError( - f"Type not supported as scalar value {type(value)}" + f"Type not supported as scalar value {type(value)}. {constants.FEEDBACK_LINK}" ) expr = self._set_or_replace_by_id(destination_id, ibis_value) return expr._reproject_to_table() - def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value): + def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value) -> ArrayValue: + """Safely assign by id while maintaining ordering integrity.""" + # TODO: Split into explicit set and replace methods + ordering_col_ids = [ + col_ref.column_id for col_ref in self._ordering.ordering_value_columns + ] + if id in ordering_col_ids: + return self._hide_column(id)._set_or_replace_by_id(id, new_value) + builder = self.builder() if id in self.column_names: builder.columns = [ @@ -937,9 +1005,9 @@ def __init__( self, session: Session, table: ibis_types.Table, + ordering: ExpressionOrdering, columns: Collection[ibis_types.Value] = (), hidden_ordering_columns: Collection[ibis_types.Value] = (), - ordering: Optional[ExpressionOrdering] = None, predicates: Optional[Collection[ibis_types.BooleanValue]] = None, ): self.session = session diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index fe7261522db..5ef92630727 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -23,8 +23,10 @@ import functools import itertools +import random import typing from typing import Iterable, List, Optional, Sequence, Tuple +import warnings import geopandas as gpd # type: ignore import google.cloud.bigquery as bigquery @@ -34,10 +36,12 @@ import pandas as pd import pyarrow as pa # type: ignore +import bigframes.constants as constants import bigframes.core as core import bigframes.core.guid as guid import bigframes.core.indexes as indexes import bigframes.core.ordering as ordering +import bigframes.core.utils import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -45,6 +49,15 @@ # Type constraint for wherever column labels are used Label = typing.Optional[str] +# Bytes to Megabyte Conversion +_BYTES_TO_KILOBYTES = 1024 +_BYTES_TO_MEGABYTES = _BYTES_TO_KILOBYTES * 1024 + +# All sampling method +_HEAD = "head" +_UNIFORM = "uniform" +_SAMPLING_METHODS = (_HEAD, _UNIFORM) + class BlockHolder(typing.Protocol): """Interface for mutable objects with state represented by a block value object.""" @@ -89,6 +102,11 @@ def __init__( raise ValueError( f"'value_columns' (size {len(self.value_columns)}) and 'column_labels' (size {len(self._column_labels)}) must have equal length" ) + # col_id -> [stat_name -> scalar] + # TODO: Preserve cache under safe transforms (eg. drop column, reorder) + self._stats_cache: dict[str, dict[str, typing.Any]] = { + col_id: {} for col_id in self.value_columns + } @property def index(self) -> indexes.IndexValue: @@ -241,7 +259,11 @@ def reset_index(self, drop: bool = True) -> Block: return block def set_index( - self, col_ids: typing.Sequence[str], drop: bool = True, append: bool = False + self, + col_ids: typing.Sequence[str], + drop: bool = True, + append: bool = False, + index_labels: typing.Sequence[Label] = (), ) -> Block: """Set the index of the block to @@ -249,6 +271,7 @@ def set_index( ids: columns to be converted to index columns drop: whether to drop the new index columns as value columns append: whether to discard the existing index or add on to it + index_labels: new index labels Returns: Block with new index @@ -269,6 +292,9 @@ def set_index( else: expr = expr.drop_columns(self.index_columns) + if index_labels: + new_index_labels = list(index_labels) + block = Block( expr, index_columns=new_index_columns, @@ -325,42 +351,222 @@ def _to_dataframe(self, result, schema: ibis_schema.Schema) -> pd.DataFrame: ) return df - def compute( - self, value_keys: Optional[Iterable[str]] = None, max_results=None + def to_pandas( + self, + value_keys: Optional[Iterable[str]] = None, + max_results: Optional[int] = None, + max_download_size: Optional[int] = None, + sampling_method: Optional[str] = None, + random_state: Optional[int] = None, ) -> Tuple[pd.DataFrame, bigquery.QueryJob]: """Run query and download results as a pandas DataFrame.""" + if max_download_size is None: + max_download_size = bigframes.options.sampling.max_download_size + if sampling_method is None: + sampling_method = ( + bigframes.options.sampling.sampling_method + if bigframes.options.sampling.sampling_method is not None + else _UNIFORM + ) + if random_state is None: + random_state = bigframes.options.sampling.random_state + + sampling_method = sampling_method.lower() + if sampling_method not in _SAMPLING_METHODS: + raise NotImplementedError( + f"The downsampling method {sampling_method} is not implemented, " + f"please choose from {','.join(_SAMPLING_METHODS)}." + ) + df, _, query_job = self._compute_and_count( - value_keys=value_keys, max_results=max_results + value_keys=value_keys, + max_results=max_results, + max_download_size=max_download_size, + sampling_method=sampling_method, + random_state=random_state, ) return df, query_job def _compute_and_count( - self, value_keys: Optional[Iterable[str]] = None, max_results=None + self, + value_keys: Optional[Iterable[str]] = None, + max_results: Optional[int] = None, + max_download_size: Optional[int] = None, + sampling_method: Optional[str] = None, + random_state: Optional[int] = None, ) -> Tuple[pd.DataFrame, int, bigquery.QueryJob]: """Run query and download results as a pandas DataFrame. Return the total number of results as well.""" # TODO(swast): Allow for dry run and timeout. - expr = self._expr + expr = self._apply_value_keys_to_expr(value_keys=value_keys) - value_column_names = value_keys or self.value_columns - if value_keys is not None: - index_columns = ( - expr.get_column(column_name) for column_name in self._index_columns + results_iterator, query_job = expr.start_query( + max_results=max_results, expose_extra_columns=True + ) + + table_size = expr._get_table_size(query_job.destination) / _BYTES_TO_MEGABYTES + fraction = ( + max_download_size / table_size + if (max_download_size is not None) and (table_size != 0) + else 2 + ) + + if fraction < 1: + if not bigframes.options.sampling.enable_downsampling: + raise RuntimeError( + f"The data size ({table_size:.2f} MB) exceeds the maximum download limit of " + f"{max_download_size} MB. You can:\n\t* Enable downsampling in global options:\n" + "\t\t`bigframes.options.sampling.enable_downsampling = True`\n" + "\t* Update the global `max_download_size` option. Please make sure " + "there is enough memory available:\n" + "\t\t`bigframes.options.sampling.max_download_size = desired_size`" + " # Setting it to None will download all the data\n" + f"{constants.FEEDBACK_LINK}" + ) + + warnings.warn( + f"The data size ({table_size:.2f} MB) exceeds the maximum download limit of" + f"({max_download_size} MB). It will be downsampled to {max_download_size} MB for download." + "\nPlease refer to the documentation for configuring the downloading limit.", + UserWarning, ) - value_columns = (expr.get_column(column_name) for column_name in value_keys) - expr = expr.projection(itertools.chain(index_columns, value_columns)) + if sampling_method == _HEAD: + total_rows = int(results_iterator.total_rows * fraction) + results_iterator.max_results = total_rows + df = self._to_dataframe(results_iterator, expr.to_ibis_expr().schema()) + + if self.index_columns: + df.set_index(list(self.index_columns), inplace=True) + df.index.names = self.index.names # type: ignore + + df.drop( + [col for col in df.columns if col not in self.value_columns], + axis=1, + inplace=True, + ) + elif (sampling_method == _UNIFORM) and (random_state is None): + filtered_expr = self.expr._uniform_sampling(fraction) + block = Block( + filtered_expr, + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + df, total_rows, _ = block._compute_and_count(max_download_size=None) + elif sampling_method == _UNIFORM: + block = self._split( + fracs=(max_download_size / table_size,), + random_state=random_state, + preserve_order=True, + )[0] + df, total_rows, _ = block._compute_and_count(max_download_size=None) + else: + # This part should never be called, just in case. + raise NotImplementedError( + f"The downsampling method {sampling_method} is not implemented, " + f"please choose from {','.join(_SAMPLING_METHODS)}." + ) + else: + total_rows = results_iterator.total_rows + df = self._to_dataframe(results_iterator, expr.to_ibis_expr().schema()) + + if self.index_columns: + df.set_index(list(self.index_columns), inplace=True) + df.index.names = self.index.names # type: ignore + + df.drop( + [col for col in df.columns if col not in self.value_columns], + axis=1, + inplace=True, + ) + + return df, total_rows, query_job - results_iterator, query_job = expr.start_query(max_results=max_results) - df = self._to_dataframe( - results_iterator, - expr.to_ibis_expr().schema(), + def _split( + self, + ns: Iterable[int] = (), + fracs: Iterable[float] = (), + *, + random_state: Optional[int] = None, + preserve_order: Optional[bool] = False, + ) -> List[Block]: + """Internal function to support splitting Block to multiple parts along index axis. + + At most one of ns and fracs can be passed in. If neither, default to ns = (1,). + Return a list of sampled Blocks. + """ + block = self + if ns and fracs: + raise ValueError("Only one of 'ns' or 'fracs' parameter must be specified.") + + if not ns and not fracs: + ns = (1,) + + if ns: + sample_sizes = ns + else: + total_rows = block.shape[0] + # Round to nearest integer. "round half to even" rule applies. + # At least to be 1. + sample_sizes = [round(frac * total_rows) or 1 for frac in fracs] + + if random_state is None: + random_state = random.randint(-(2**63), 2**63 - 1) + + # Create a new column with random_state value. + block, random_state_col = block.create_constant(str(random_state)) + + # Create an ordering col and convert to string + block, ordering_col = block.promote_offsets() + block, string_ordering_col = block.apply_unary_op( + ordering_col, ops.AsTypeOp("string[pyarrow]") + ) + + # Apply hash method to sum col and order by it. + block, string_sum_col = block.apply_binary_op( + string_ordering_col, random_state_col, ops.concat_op ) + block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op) + block = block.order_by([ordering.OrderingColumnReference(hash_string_sum_col)]) - df = df.loc[:, [*self.index_columns, *value_column_names]] - if self.index_columns: - df = df.set_index(list(self.index_columns)) - df.index.names = self.index.names # type: ignore + intervals = [] + cur = 0 - return df, results_iterator.total_rows, query_job + for sample_size in sample_sizes: + intervals.append((cur, cur + sample_size)) + cur += sample_size + + sliced_blocks = [ + typing.cast(Block, block.slice(start=lower, stop=upper)) + for lower, upper in intervals + ] + if preserve_order: + sliced_blocks = [ + sliced_block.order_by([ordering.OrderingColumnReference(ordering_col)]) + for sliced_block in sliced_blocks + ] + + drop_cols = [ + random_state_col, + ordering_col, + string_ordering_col, + string_sum_col, + hash_string_sum_col, + ] + return [sliced_block.drop_columns(drop_cols) for sliced_block in sliced_blocks] + + def _compute_dry_run( + self, value_keys: Optional[Iterable[str]] = None + ) -> bigquery.QueryJob: + expr = self._apply_value_keys_to_expr(value_keys=value_keys) + job_config = bigquery.QueryJobConfig(dry_run=True) + _, query_job = expr.start_query(job_config=job_config) + return query_job + + def _apply_value_keys_to_expr(self, value_keys: Optional[Iterable[str]] = None): + expr = self._expr + if value_keys is not None: + expr = expr.select_columns(itertools.chain(self._index_columns, value_keys)) + return expr def with_column_labels(self, value: typing.Iterable[Label]) -> Block: label_list = tuple(value) @@ -553,10 +759,17 @@ def assign_label(self, column_id: str, new_label: Label) -> Block: new_labels[col_index] = new_label return self.with_column_labels(new_labels) - def filter(self, column_name: str): + def filter(self, column_name: str, keep_null: bool = False): condition = typing.cast( ibis_types.BooleanValue, self._expr.get_column(column_name) ) + if keep_null: + condition = typing.cast( + ibis_types.BooleanValue, + condition.fillna( + typing.cast(ibis_types.BooleanScalar, ibis_types.literal(True)) + ), + ) filtered_expr = self.expr.filter(condition) return Block( filtered_expr, @@ -576,10 +789,10 @@ def aggregate_all_and_pivot( aggregations = [(col_id, operation, col_id) for col_id in self.value_columns] result_expr = self.expr.aggregate( aggregations, dropna=dropna - ).transpose_single_row( - labels=self.column_labels, + ).unpivot_single_row( + row_labels=self.column_labels, index_col_id="index", - value_col_id=value_col_id, + unpivot_columns=[(value_col_id, self.value_columns)], dtype=dtype, ) return Block(result_expr, index_columns=["index"], column_labels=[None]) @@ -614,8 +827,8 @@ def rename(self, *, columns: typing.Mapping[Label, Label]): def aggregate( self, - by_column_ids: typing.Sequence[str], - aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp]], + by_column_ids: typing.Sequence[str] = (), + aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp]] = (), *, as_index: bool = True, dropna: bool = True, @@ -623,7 +836,7 @@ def aggregate( """ Apply aggregations to the block. Callers responsible for setting index column(s) after. Arguments: - by_column_id: column id of the aggregation key, this is preserved through the transform and used as index + by_column_id: column id of the aggregation key, this is preserved through the transform and used as index. aggregations: input_column_id, operation tuples as_index: if True, grouping keys will be index columns in result, otherwise they will be non-index columns. dropna: whether null keys should be dropped @@ -639,15 +852,12 @@ def aggregate( [agg[0] for agg in aggregations] ) if as_index: - # TODO: Generalize to multi-index names: typing.List[Label] = [] for by_col_id in by_column_ids: - if by_col_id in self.index_columns: - # Groupby level 0 case, keep index name - index_name = self.col_id_to_index_name[by_col_id] + if by_col_id in self.value_columns: + names.append(self.col_id_to_label[by_col_id]) else: - index_name = self.col_id_to_label[by_col_id] - names.append(index_name) + names.append(self.col_id_to_index_name[by_col_id]) return ( Block( result_expr, @@ -657,10 +867,89 @@ def aggregate( ), output_col_ids, ) - else: - by_column_labels = self._get_labels_for_columns(by_column_ids) + else: # as_index = False + # If as_index=False, drop grouping levels, but keep grouping value columns + by_value_columns = [ + col for col in by_column_ids if col in self.value_columns + ] + by_column_labels = self._get_labels_for_columns(by_value_columns) labels = (*by_column_labels, *aggregate_labels) - return Block(result_expr, column_labels=labels), output_col_ids + result_expr_pruned = result_expr.select_columns( + [*by_value_columns, *output_col_ids] + ) + return Block(result_expr_pruned, column_labels=labels), output_col_ids + + def get_stat(self, column_id: str, stat: agg_ops.AggregateOp): + """Gets aggregates immediately, and caches it""" + if stat.name in self._stats_cache[column_id]: + return self._stats_cache[column_id][stat.name] + + # TODO: Convert nonstandard stats into standard stats where possible (popvar, etc.) + # if getting a standard stat, just go get the rest of them + standard_stats = self._standard_stats(column_id) + stats_to_fetch = standard_stats if stat in standard_stats else [stat] + + aggregations = [(column_id, stat, stat.name) for stat in stats_to_fetch] + expr = self.expr.aggregate(aggregations) + block = Block(expr, column_labels=[s.name for s in stats_to_fetch]) + df, _ = block.to_pandas() + + # Carefully extract stats such that they aren't coerced to a common type + stats_map = {stat_name: df.loc[0, stat_name] for stat_name in df.columns} + self._stats_cache[column_id].update(stats_map) + return stats_map[stat.name] + + def summarize( + self, + column_ids: typing.Sequence[str], + stats: typing.Sequence[agg_ops.AggregateOp], + ): + """Get a list of stats as a deferred block object.""" + label_col_id = guid.generate_guid() + labels = [stat.name for stat in stats] + aggregations = [ + (col_id, stat, f"{col_id}-{stat.name}") + for stat in stats + for col_id in column_ids + ] + columns = [ + (col_id, [f"{col_id}-{stat.name}" for stat in stats]) + for col_id in column_ids + ] + expr = self.expr.aggregate(aggregations).unpivot_single_row( + labels, + unpivot_columns=columns, + index_col_id=label_col_id, + ) + labels = self._get_labels_for_columns(column_ids) + return Block(expr, column_labels=labels, index_columns=[label_col_id]) + + def _standard_stats(self, column_id) -> typing.Sequence[agg_ops.AggregateOp]: + """ + Gets a standard set of stats to preemptively fetch for a column if + any other stat is fetched. + Helps prevent repeat scanning of the same column to fetch statistics. + Standard stats should be: + - commonly used + - efficiently computable. + """ + # TODO: annotate aggregations themself with this information + dtype = self.expr.get_column_type(column_id) + stats: list[agg_ops.AggregateOp] = [agg_ops.count_op] + if dtype not in bigframes.dtypes.UNORDERED_DTYPES: + stats += [agg_ops.min_op, agg_ops.max_op] + if dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES: + # Notable exclusions: + # prod op tends to cause overflows + # Also, var_op is redundant as can be derived from std + stats += [ + agg_ops.std_op, + agg_ops.mean_op, + agg_ops.var_op, + agg_ops.sum_op, + ] + + return stats def _get_labels_for_columns(self, column_ids: typing.Sequence[str]): """Get column label for value columns, or index name for index columns""" @@ -699,6 +988,29 @@ def slice( ) return block + def retrieve_repr_request_results( + self, max_results: int + ) -> Tuple[pd.DataFrame, int, bigquery.QueryJob]: + """ + Retrieves a pandas dataframe containing only max_results many rows for use + with printing methods. + + Returns a tuple of the dataframe and the overall number of rows of the query. + """ + # TODO(swast): Select a subset of columns if max_columns is less than the + # number of columns in the schema. + count = self.shape[0] + if count > max_results: + head_block = self.slice(0, max_results) + computed_df, query_job = head_block.to_pandas(max_results=max_results) + else: + head_block = self + computed_df, query_job = head_block.to_pandas() + formatted_df = computed_df.set_axis(self.column_labels, axis=1) + # we reset the axis and substitute the bf index name for the default + formatted_df.index.name = self.index.name + return formatted_df, count, query_job + def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]: expr, result_id = self._expr.promote_offsets() return ( @@ -712,7 +1024,7 @@ def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]: ) def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: - axis_number = _get_axis_number(axis) + axis_number = bigframes.core.utils.get_axis_number(axis) if axis_number == 0: expr = self._expr for index_col in self._index_columns: @@ -735,7 +1047,7 @@ def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: ) def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block: - axis_number = _get_axis_number(axis) + axis_number = bigframes.core.utils.get_axis_number(axis) if axis_number == 0: expr = self._expr for index_col in self._index_columns: @@ -788,6 +1100,15 @@ def concat( result_block = result_block.reset_index() return result_block + def _force_reproject(self) -> Block: + """Forces a reprojection of the underlying tables expression. Used to force predicate/order application before subsequent operations.""" + return Block( + self._expr._reproject_to_table(), + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index.names, + ) + def block_from_local(data, session=None, use_index=True) -> Block: # TODO(tbergeron): Handle duplicate column labels @@ -795,14 +1116,20 @@ def block_from_local(data, session=None, use_index=True) -> Block: column_labels = list(pd_data.columns) if not all((label is None) or isinstance(label, str) for label in column_labels): - raise NotImplementedError("Only string column labels supported") + raise NotImplementedError( + f"Only string column labels supported. {constants.FEEDBACK_LINK}" + ) if use_index: if pd_data.index.nlevels > 1: - raise NotImplementedError("multi-indices not supported.") + raise NotImplementedError( + f"multi-indices not supported. {constants.FEEDBACK_LINK}" + ) index_label = pd_data.index.name if (index_label is not None) and (not isinstance(index_label, str)): - raise NotImplementedError("Only string index names supported") + raise NotImplementedError( + f"Only string index names supported. {constants.FEEDBACK_LINK}" + ) index_id = guid.generate_guid() pd_data = pd_data.reset_index(names=index_id) @@ -853,11 +1180,11 @@ def _align_indices(blocks: typing.Sequence[Block]) -> typing.Sequence[Label]: for block in blocks[1:]: if len(names) != block.index.nlevels: raise NotImplementedError( - "Cannot combine indices with different number of levels. Use 'ignore_index'=True." + f"Cannot combine indices with different number of levels. Use 'ignore_index'=True. {constants.FEEDBACK_LINK}" ) if block.index.dtypes != types: raise NotImplementedError( - "Cannot combine different index dtypes. Use 'ignore_index'=True." + f"Cannot combine different index dtypes. Use 'ignore_index'=True. {constants.FEEDBACK_LINK}" ) names = [ lname if lname == rname else None @@ -875,7 +1202,7 @@ def _combine_schema_inner( if label in right: if type != right[label]: raise ValueError( - f"Cannot concat rows with label {label} due to mismatched types" + f"Cannot concat rows with label {label} due to mismatched types. {constants.FEEDBACK_LINK}" ) result[label] = type return result @@ -889,7 +1216,7 @@ def _combine_schema_outer( for label, type in left.items(): if (label in right) and (type != right[label]): raise ValueError( - f"Cannot concat rows with label {label} due to mismatched types" + f"Cannot concat rows with label {label} due to mismatched types. {constants.FEEDBACK_LINK}" ) result[label] = type for label, type in right.items(): @@ -906,12 +1233,3 @@ def _get_block_schema( for label, dtype in zip(block.column_labels, block.dtypes): result[label] = typing.cast(bigframes.dtypes.Dtype, dtype) return result - - -def _get_axis_number(axis: str | int | None) -> typing.Literal[0, 1]: - if axis in {0, "index", "rows", None}: - return 0 - elif axis in {1, "columns"}: - return 1 - else: - raise ValueError(f"Not a valid axis: {axis}") diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index a8655116e93..288dcdd5b0b 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -16,12 +16,16 @@ import typing +import pandas as pd +import typing_extensions + +import bigframes.constants as constants import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.ordering as order import bigframes.core.window as windows import bigframes.dataframe as df -import bigframes.dtypes +import bigframes.dtypes as dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops import bigframes.series as series @@ -36,6 +40,7 @@ def __init__( block: blocks.Block, by_col_ids: typing.Sequence[str], *, + selected_cols: typing.Optional[typing.Sequence[str]] = None, dropna: bool = True, as_index: bool = True, ): @@ -48,8 +53,52 @@ def __init__( ) } self._by_col_ids = by_col_ids - self._dropna = dropna # Applies to aggregations but not windowing + + self._dropna = dropna self._as_index = as_index + if selected_cols: + for col in selected_cols: + if col not in self._block.value_columns: + raise ValueError(f"Invalid column selection: {col}") + self._selected_cols = selected_cols + else: + self._selected_cols = [ + col_id + for col_id in self._block.value_columns + if col_id not in self._by_col_ids + ] + + def __getitem__( + self, + key: typing.Union[ + blocks.Label, + typing.Sequence[blocks.Label], + ], + ): + if _is_list_like(key): + keys = list(key) + else: + keys = [key] + columns = [ + col_id for col_id, label in self._col_id_labels.items() if label in keys + ] + + if len(columns) > 1 or (not self._as_index): + return DataFrameGroupBy( + self._block, + self._by_col_ids, + selected_cols=columns, + dropna=self._dropna, + as_index=self._as_index, + ) + else: + return SeriesGroupBy( + self._block, + columns[0], + self._by_col_ids, + value_name=self._col_id_labels[columns[0]], + dropna=self._dropna, + ) def sum(self, numeric_only: bool = False, *args) -> df.DataFrame: if not numeric_only: @@ -61,15 +110,22 @@ def mean(self, numeric_only: bool = False, *args) -> df.DataFrame: self._raise_on_non_numeric("mean") return self._aggregate(agg_ops.mean_op, numeric_only=True) - def min(self, numeric_only: bool = False, *args) -> df.DataFrame: + def median( + self, numeric_only: bool = False, *, exact: bool = False + ) -> df.DataFrame: + if exact: + raise NotImplementedError( + f"Only approximate median is supported. {constants.FEEDBACK_LINK}" + ) if not numeric_only: - self._raise_on_non_numeric("min") - return self._aggregate(agg_ops.min_op, numeric_only=True) + self._raise_on_non_numeric("median") + return self._aggregate(agg_ops.median_op, numeric_only=True) + + def min(self, numeric_only: bool = False, *args) -> df.DataFrame: + return self._aggregate(agg_ops.min_op, numeric_only=numeric_only) def max(self, numeric_only: bool = False, *args) -> df.DataFrame: - if not numeric_only: - self._raise_on_non_numeric("max") - return self._aggregate(agg_ops.max_op, numeric_only=True) + return self._aggregate(agg_ops.max_op, numeric_only=numeric_only) def std( self, @@ -101,45 +157,40 @@ def count(self) -> df.DataFrame: def cumsum(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("cumsum") - window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) - return self._apply_window_op(agg_ops.sum_op, window, numeric_only=True) + return self._apply_window_op(agg_ops.sum_op, numeric_only=True) def cummin(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: - if not numeric_only: - self._raise_on_non_numeric("cummin") - window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) - return self._apply_window_op(agg_ops.min_op, window, numeric_only=True) + return self._apply_window_op(agg_ops.min_op, numeric_only=numeric_only) def cummax(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: - if not numeric_only: - self._raise_on_non_numeric("cummax") - window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) - return self._apply_window_op(agg_ops.max_op, window, numeric_only=True) + return self._apply_window_op(agg_ops.max_op, numeric_only=numeric_only) def cumprod(self, *args, **kwargs) -> df.DataFrame: - window = bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0) - return self._apply_window_op(agg_ops.product_op, window, numeric_only=True) + return self._apply_window_op(agg_ops.product_op, numeric_only=True) def _raise_on_non_numeric(self, op: str): if not all( - dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES - for dtype in self._block.dtypes + dtype in dtypes.NUMERIC_BIGFRAMES_TYPES for dtype in self._block.dtypes ): raise NotImplementedError( - f"'{op}' does not support non-numeric columns. Set 'numeric_only'=True to ignore non-numeric columns" + f"'{op}' does not support non-numeric columns. " + "Set 'numeric_only'=True to ignore non-numeric columns. " + f"{constants.FEEDBACK_LINK}" ) return self - def _aggregated_columns(self, numeric_only: bool = False): - return [ - col_id - for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) - if col_id not in self._by_col_ids - and ( - (not numeric_only) - or (dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) - ) - ] + def _aggregated_columns(self, numeric_only: bool = False) -> typing.Sequence[str]: + valid_agg_cols: list[str] = [] + for col_id in self._selected_cols: + is_numeric = self._column_type(col_id) in dtypes.NUMERIC_BIGFRAMES_TYPES + if is_numeric or not numeric_only: + valid_agg_cols.append(col_id) + return valid_agg_cols + + def _column_type(self, col_id: str) -> dtypes.Dtype: + col_offset = self._block.value_columns.index(col_id) + dtype = self._block.dtypes[col_offset] + return dtype def _aggregate( self, aggregate_op: agg_ops.AggregateOp, numeric_only: bool = False @@ -147,8 +198,8 @@ def _aggregate( aggregated_col_ids = self._aggregated_columns(numeric_only=numeric_only) aggregations = [(col_id, aggregate_op) for col_id in aggregated_col_ids] result_block, _ = self._block.aggregate( - self._by_col_ids, - aggregations, + by_column_ids=self._by_col_ids, + aggregations=aggregations, as_index=self._as_index, dropna=self._dropna, ) @@ -157,15 +208,16 @@ def _aggregate( def _apply_window_op( self, op: agg_ops.WindowOp, - window_spec: bigframes.core.WindowSpec, + window: typing.Optional[core.WindowSpec] = None, numeric_only: bool = False, ): + """Apply window op to groupby. Defaults to grouped cumulative window.""" + window_spec = window or core.WindowSpec( + grouping_keys=self._by_col_ids, following=0 + ) columns = self._aggregated_columns(numeric_only=numeric_only) - block = self._block.select_columns([*columns, *window_spec.grouping_keys]) block = self._block.multi_apply_window_op( - columns, - op, - window_spec=window_spec, + columns, op, window_spec=window_spec, skip_null_groups=self._dropna ) block = block.select_columns(columns) return df.DataFrame(block) @@ -190,7 +242,7 @@ def __init__( self._dropna = dropna # Applies to aggregations but not windowing @property - def value(self): + def _value(self): return self._block.expr.get_column(self._value_column) def all(self) -> series.Series: @@ -199,16 +251,24 @@ def all(self) -> series.Series: def any(self) -> series.Series: return self._aggregate(agg_ops.any_op) + def min(self, *args) -> series.Series: + return self._aggregate(agg_ops.min_op) + + def max(self, *args) -> series.Series: + return self._aggregate(agg_ops.max_op) + def count(self) -> series.Series: return self._aggregate(agg_ops.count_op) def sum(self, *args) -> series.Series: - """Sums the numeric values for each group in the series. Ignores null/nan.""" return self._aggregate(agg_ops.sum_op) def mean(self, *args) -> series.Series: return self._aggregate(agg_ops.mean_op) + def median(self, *args, **kwargs) -> series.Series: + return self._aggregate(agg_ops.mean_op) + def std(self, *args, **kwargs) -> series.Series: return self._aggregate(agg_ops.std_op) @@ -221,45 +281,39 @@ def prod(self, *args) -> series.Series: def cumsum(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.sum_op, - bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), ) def cumprod(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.product_op, - bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), ) def cummax(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.max_op, - bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), ) def cummin(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.min_op, - bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), ) def cumcount(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.rank_op, - bigframes.core.WindowSpec(grouping_keys=self._by_col_ids, following=0), discard_name=True, )._apply_unary_op(ops.partial_right(ops.sub_op, 1)) def shift(self, periods=1) -> series.Series: """Shift index by desired number of periods.""" - window = bigframes.core.WindowSpec( + window = core.WindowSpec( grouping_keys=self._by_col_ids, preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) - return self._apply_window_op(agg_ops.ShiftOp(periods), window) + return self._apply_window_op(agg_ops.ShiftOp(periods), window=window) def diff(self) -> series.Series: - """Difference between each element and previous element.""" return self._ungroup() - self.shift(1) def rolling(self, window: int, min_periods=None) -> windows.Window: @@ -278,7 +332,9 @@ def rolling(self, window: int, min_periods=None) -> windows.Window: def expanding(self, min_periods: int = 1) -> windows.Window: window_spec = core.WindowSpec( - grouping_keys=self._by_col_ids, following=0, min_periods=min_periods + grouping_keys=self._by_col_ids, + following=0, + min_periods=min_periods, ) block = self._block.order_by( [order.OrderingColumnReference(col) for col in self._by_col_ids], @@ -301,9 +357,14 @@ def _aggregate(self, aggregate_op: agg_ops.AggregateOp) -> series.Series: def _apply_window_op( self, op: agg_ops.WindowOp, - window_spec: bigframes.core.WindowSpec, discard_name=False, + window: typing.Optional[core.WindowSpec] = None, ): + """Apply window op to groupby. Defaults to grouped cumulative window.""" + window_spec = window or core.WindowSpec( + grouping_keys=self._by_col_ids, following=0 + ) + label = self._value_name if not discard_name else None block, result_id = self._block.apply_window_op( self._value_column, @@ -313,3 +374,7 @@ def _apply_window_op( skip_null_groups=self._dropna, ) return series.Series(block.select_column(result_id)) + + +def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: + return pd.api.types.is_list_like(obj) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index f65cb02941e..0aaf169bea4 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -20,6 +20,7 @@ import ibis import pandas as pd +import bigframes.constants as constants import bigframes.core as core import bigframes.core.guid as guid import bigframes.core.indexes as indexes @@ -47,10 +48,14 @@ def __setitem__(self, key, value) -> None: # TODO(swast): support MultiIndex if isinstance(key, slice): # TODO(swast): Implement loc with slices. - raise NotImplementedError("loc does not yet support slices") + raise NotImplementedError( + f"loc does not yet support slices. {constants.FEEDBACK_LINK}" + ) elif isinstance(key, list): # TODO(tbergeron): Implement loc for index label list. - raise NotImplementedError("loc does not yet support index label lists") + raise NotImplementedError( + f"loc does not yet support index label lists. {constants.FEEDBACK_LINK}" + ) # Assume the key is for the index label. block = self._series._block @@ -149,6 +154,7 @@ def __setitem__( ): raise NotImplementedError( "Only setting a column by DataFrame.loc[:, 'column'] is supported." + f"{constants.FEEDBACK_LINK}" ) # TODO(swast): Support setting multiple columns with key[1] as a list @@ -213,19 +219,45 @@ def _loc_getitem_series_or_dataframe( typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], series_or_dataframe.iloc[0:0], ) + + # We can't upload a DataFrame with None as the column name, so set it + # an arbitrary string. index_name = series_or_dataframe.index.name + index_name_is_none = index_name is None + if index_name_is_none: + index_name = "unnamed_col" + keys_df = bigframes.dataframe.DataFrame( {index_name: key}, session=series_or_dataframe._get_block().expr._session ) keys_df = keys_df.set_index(index_name, drop=True) + + if index_name_is_none: + keys_df.index.name = None return _perform_loc_list_join(series_or_dataframe, keys_df) elif isinstance(key, slice): - return series_or_dataframe._slice(key.start, key.stop, key.step) + if (key.start is None) and (key.stop is None) and (key.step is None): + return series_or_dataframe.copy() + raise NotImplementedError( + f"loc does not yet support indexing with a slice. {constants.FEEDBACK_LINK}" + ) elif callable(key): - raise NotImplementedError("loc does not yet support indexing with a callable") + raise NotImplementedError( + f"loc does not yet support indexing with a callable. {constants.FEEDBACK_LINK}" + ) + elif pd.api.types.is_scalar(key): + index_name = "unnamed_col" + keys_df = bigframes.dataframe.DataFrame( + {index_name: [key]}, session=series_or_dataframe._get_block().expr._session + ) + keys_df = keys_df.set_index(index_name, drop=True) + keys_df.index.name = None + return _perform_loc_list_join(series_or_dataframe, keys_df) else: raise TypeError( - "Invalid argument type. loc currently only supports indexing with a boolean bigframes Series or a list of index entries." + "Invalid argument type. loc currently only supports indexing with a " + "boolean bigframes Series, a list of index entries or a single index entry. " + f"{constants.FEEDBACK_LINK}" ) @@ -257,11 +289,13 @@ def _perform_loc_list_join( name = series_or_dataframe.name if series_or_dataframe.name is not None else "0" result = typing.cast( bigframes.series.Series, - series_or_dataframe.to_frame().join(keys_df, how="right")[name], + series_or_dataframe.to_frame()._perform_join_by_index(keys_df, how="right")[ + name + ], ) result = result.rename(original_name) else: - result = series_or_dataframe.join(keys_df, how="right") # type: ignore + result = series_or_dataframe._perform_join_by_index(keys_df, how="right") # type: ignore result = result.rename_axis(original_index_names) return result @@ -284,12 +318,8 @@ def _iloc_getitem_series_or_dataframe( series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, key ) -> bigframes.dataframe.DataFrame | bigframes.series.Series | bigframes.core.scalar.Scalar | pd.Series: if isinstance(key, int): - if key < 0: - raise NotImplementedError( - "iloc does not yet support negative single positional index" - ) internal_slice_result = series_or_dataframe._slice(key, key + 1, 1) - result_pd_df = internal_slice_result.compute() + result_pd_df = internal_slice_result.to_pandas() if result_pd_df.empty: raise IndexError("single positional indexer is out-of-bounds") return result_pd_df.iloc[0] @@ -329,9 +359,11 @@ def _iloc_getitem_series_or_dataframe( elif isinstance(key, tuple): raise NotImplementedError( - "iloc does not yet support indexing with a (row, column) tuple" + f"iloc does not yet support indexing with a (row, column) tuple. {constants.FEEDBACK_LINK}" ) elif callable(key): - raise NotImplementedError("iloc does not yet support indexing with a callable") + raise NotImplementedError( + f"iloc does not yet support indexing with a callable. {constants.FEEDBACK_LINK}" + ) else: - raise TypeError("Invalid argument type.") + raise TypeError(f"Invalid argument type. {constants.FEEDBACK_LINK}") diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index f7fa7ac1195..bb1fcc4b175 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -20,8 +20,9 @@ from typing import Callable, Tuple import numpy as np -import pandas as pd +import pandas +import bigframes.constants as constants import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.joins as joins @@ -45,6 +46,7 @@ def name(self, value: blocks.Label): @property def names(self) -> typing.Sequence[blocks.Label]: + """Returns the names of the Index.""" return self._data._get_block()._index_labels @names.setter @@ -57,21 +59,40 @@ def shape(self) -> typing.Tuple[int]: @property def size(self) -> int: + """Returns the size of the Index.""" return self.shape[0] @property def empty(self) -> bool: + """Returns True if the Index is empty, otherwise returns False.""" return self.shape[0] == 0 - def to_pandas(self) -> pd.Index: - """Get the Index as a pandas Index.""" - return IndexValue(self._data._get_block()).compute() + def __getitem__(self, key: int) -> typing.Any: + if isinstance(key, int): + result_pd_df, _ = self._data._get_block().slice(key, key + 1, 1).to_pandas() + if result_pd_df.empty: + raise IndexError("single positional indexer is out-of-bounds") + return result_pd_df.index[0] + else: + raise NotImplementedError(f"Index key not supported {key}") + + def to_pandas(self) -> pandas.Index: + """Gets the Index as a pandas Index. + + Returns: + pandas.Index: + A pandas Index with all of the labels from this Index. + """ + return IndexValue(self._data._get_block()).to_pandas() + + def to_numpy(self, dtype=None, **kwargs) -> np.ndarray: + return self.to_pandas().to_numpy(dtype, **kwargs) + + __array__ = to_numpy def __len__(self): return self.shape[0] - compute = to_pandas - class IndexValue: """An immutable index.""" @@ -107,10 +128,10 @@ def __repr__(self) -> str: # maybe we just print the job metadata that we have so far? # TODO(swast): Avoid downloading the whole index by using job # metadata, like we do with DataFrame. - preview = self.compute() + preview = self.to_pandas() return repr(preview) - def compute(self) -> pd.Index: + def to_pandas(self) -> pandas.Index: """Executes deferred operations and downloads the results.""" # Project down to only the index column. So the query can be cached to visualize other data. index_column = self._block.index_columns[0] @@ -123,26 +144,35 @@ def compute(self) -> pd.Index: return index def join( - self, other: IndexValue, *, how="left", sort=False + self, + other: IndexValue, + *, + how="left", + sort=False, + block_identity_join: bool = False, ) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: if not isinstance(other, IndexValue): # TODO(swast): We need to improve this error message to be more # actionable for the user. For example, it's possible they # could call set_index and try again to resolve this error. raise ValueError( - "Can't mixed objects with explicit Index and ImpliedJoiner" + f"Tried to join with an unexpected type: {type(other)}. {constants.FEEDBACK_LINK}" ) # TODO(swast): Support cross-joins (requires reindexing). if how not in {"outer", "left", "right", "inner"}: raise NotImplementedError( - "Only how='outer','left','right','inner' currently supported" + f"Only how='outer','left','right','inner' currently supported. {constants.FEEDBACK_LINK}" ) if self.nlevels == other.nlevels == 1: - return join_mono_indexed(self, other, how=how, sort=sort) + return join_mono_indexed( + self, other, how=how, sort=sort, block_identity_join=block_identity_join + ) else: # Always sort mult-index join - return join_multi_indexed(self, other, how=how, sort=sort) + return join_multi_indexed( + self, other, how=how, sort=sort, block_identity_join=block_identity_join + ) def resolve_level_name(self: IndexValue, label: blocks.Label) -> str: matches = self._block.index_name_to_col_id.get(label, []) @@ -157,7 +187,12 @@ def is_uniquely_named(self: IndexValue): def join_mono_indexed( - left: IndexValue, right: IndexValue, *, how="left", sort=False + left: IndexValue, + right: IndexValue, + *, + how="left", + sort=False, + block_identity_join: bool = False, ) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: ( combined_expr, @@ -170,6 +205,7 @@ def join_mono_indexed( right._block.index_columns, how=how, sort=sort, + allow_row_identity_join=(not block_identity_join), ) # Drop original indices from each side. and used the coalesced combination generated by the join. left_indices = [get_column_left(col_id) for col_id in left._block.index_columns] @@ -188,7 +224,12 @@ def join_mono_indexed( def join_multi_indexed( - left: IndexValue, right: IndexValue, *, how="left", sort=False + left: IndexValue, + right: IndexValue, + *, + how="left", + sort=False, + block_identity_join: bool = False, ) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: if not (left.is_uniquely_named() and right.is_uniquely_named()): raise ValueError("Joins not supported on indices with non-unique level names") @@ -203,6 +244,7 @@ def join_multi_indexed( left_join_ids = [left.resolve_level_name(name) for name in common_names] right_join_ids = [right.resolve_level_name(name) for name in common_names] + names_fully_match = len(left_only_names) == 0 and len(right_only_names) == 0 ( combined_expr, joined_index_col_names, @@ -214,6 +256,9 @@ def join_multi_indexed( right_join_ids, how=how, sort=sort, + # If we're only joining on a subset of the index columns, we need to + # perform a true join. + allow_row_identity_join=names_fully_match and not block_identity_join, ) # Drop original indices from each side. and used the coalesced combination generated by the join. combined_expr = combined_expr.drop_columns( diff --git a/bigframes/core/io.py b/bigframes/core/io.py new file mode 100644 index 00000000000..6b7ed52488d --- /dev/null +++ b/bigframes/core/io.py @@ -0,0 +1,90 @@ +# 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. + +"""Private module: Helpers for I/O operations.""" + +import datetime +import textwrap +from typing import Dict, Union + +import google.cloud.bigquery as bigquery + +IO_ORDERING_ID = "bqdf_row_nums" + + +def create_export_csv_statement( + table_id: str, uri: str, field_delimiter: str, header: bool +) -> str: + return create_export_data_statement( + table_id, + uri, + "CSV", + { + "field_delimiter": field_delimiter, + "header": header, + }, + ) + + +def create_export_data_statement( + table_id: str, uri: str, format: str, export_options: Dict[str, Union[bool, str]] +) -> str: + all_options: Dict[str, Union[bool, str]] = { + "uri": uri, + "format": format, + # TODO(swast): Does pandas have an option not to overwrite files? + "overwrite": True, + } + all_options.update(export_options) + export_options_str = ", ".join( + format_option(key, value) for key, value in all_options.items() + ) + # Manually generate ORDER BY statement since ibis will not always generate + # it in the top level statement. This causes BigQuery to then run + # non-distributed sort and run out of memory. + return textwrap.dedent( + f""" + EXPORT DATA + OPTIONS ( + {export_options_str} + ) AS + SELECT * EXCEPT ({IO_ORDERING_ID}) + FROM `{table_id}` + ORDER BY {IO_ORDERING_ID} + """ + ) + + +def create_snapshot_sql( + table_ref: bigquery.TableReference, current_timestamp: datetime.datetime +) -> str: + """Query a table via 'time travel' for consistent reads.""" + + # If we have a _SESSION table, assume that it's already a copy. Nothing to do here. + if table_ref.dataset_id.upper() == "_SESSION": + return f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" + + return textwrap.dedent( + f""" + SELECT * + FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}` + FOR SYSTEM_TIME AS OF TIMESTAMP({repr(current_timestamp.isoformat())}) + """ + ) + + +def format_option(key: str, value: Union[bool, str]) -> str: + if isinstance(value, bool): + return f"{key}=true" if value else f"{key}=false" + return f"{key}={repr(value)}" diff --git a/bigframes/core/joins/row_identity.py b/bigframes/core/joins/row_identity.py index b77b5d80adc..66eb2239904 100644 --- a/bigframes/core/joins/row_identity.py +++ b/bigframes/core/joins/row_identity.py @@ -23,6 +23,7 @@ import ibis import ibis.expr.types as ibis_types +import bigframes.constants as constants import bigframes.core as core SUPPORTED_ROW_IDENTITY_HOW = {"outer", "left", "inner"} @@ -33,7 +34,9 @@ def join_by_row_identity( ) -> Tuple[core.ArrayValue, Tuple[Callable[[str], str], Callable[[str], str]],]: """Compute join when we are joining by row identity not a specific column.""" if how not in SUPPORTED_ROW_IDENTITY_HOW: - raise NotImplementedError("Only how='outer','left','inner' currently supported") + raise NotImplementedError( + f"Only how='outer','left','inner' currently supported. {constants.FEEDBACK_LINK}" + ) if not left.table.equals(right.table): raise ValueError( @@ -67,47 +70,41 @@ def join_by_row_identity( for key in right.column_names.keys() ] - hidden_ordering_columns = [] - new_ordering = core.ExpressionOrdering() - if left._ordering and right._ordering: - # These ordering columns will be present in the ArrayValue, as we - # haven't hidden any value / index column(s). Code that is aware of - # which columns are index columns / value columns columns will need to - # add the previous columns to hidden columns. - new_ordering = left._ordering.with_ordering_columns( - [ - col_ref.with_name(map_left_id(col_ref.column_id)) - for col_ref in left._ordering.ordering_value_columns - ] - + [ - col_ref.with_name(map_right_id(col_ref.column_id)) - for col_ref in right._ordering.ordering_value_columns - ] - ) - - hidden_ordering_columns = [ - left._get_hidden_ordering_column(key.column_id).name( - map_left_id(key.column_id) - ) - for key in left._ordering.ordering_value_columns - if key.column_id in left._hidden_ordering_column_names.keys() + # If left isn't being masked, can just use left ordering + if not left_mask: + col_mapping = { + order_ref.column_id: map_left_id(order_ref.column_id) + for order_ref in left._ordering.ordering_value_columns + } + new_ordering = left._ordering.with_column_remap(col_mapping) + else: + ordering_columns = [ + col_ref.with_name(map_left_id(col_ref.column_id)) + for col_ref in left._ordering.ordering_value_columns ] + [ - right._get_hidden_ordering_column(key.column_id).name( - map_right_id(key.column_id) - ) - for key in right._ordering.ordering_value_columns - if key.column_id in right._hidden_ordering_column_names.keys() + col_ref.with_name(map_right_id(col_ref.column_id)) + for col_ref in right._ordering.ordering_value_columns ] + left_total_order_cols = frozenset( + map_left_id(col) for col in left._ordering.total_ordering_columns + ) + # Assume that left ordering is sufficient since 1:1 join over same base table + join_total_order_cols = left_total_order_cols + new_ordering = core.ExpressionOrdering( + ordering_columns, total_ordering_columns=join_total_order_cols + ) - left_ordering_id = left._ordering.ordering_id - if left_ordering_id: - new_ordering = new_ordering.with_ordering_id(map_left_id(left_ordering_id)) - if left_ordering_id in left._hidden_ordering_column_names.keys(): - hidden_ordering_columns.append( - left._get_hidden_ordering_column(left_ordering_id).name( - map_left_id(left_ordering_id) - ) - ) + hidden_ordering_columns = [ + left._get_hidden_ordering_column(key.column_id).name(map_left_id(key.column_id)) + for key in left._ordering.ordering_value_columns + if key.column_id in left._hidden_ordering_column_names.keys() + ] + [ + right._get_hidden_ordering_column(key.column_id).name( + map_right_id(key.column_id) + ) + for key in right._ordering.ordering_value_columns + if key.column_id in right._hidden_ordering_column_names.keys() + ] joined_expr = core.ArrayValue( left._session, @@ -169,7 +166,9 @@ def _join_predicates( ) return (*left_predicates, *right_relative_predicates) else: - raise ValueError("Unsupported join_type: " + join_type) + raise ValueError( + f"Unsupported join_type: {join_type}. {constants.FEEDBACK_LINK}" + ) def _get_relative_predicates( diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 34873e6f7ee..4c865fffdf7 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -23,6 +23,7 @@ import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types +import bigframes.constants as constants import bigframes.core as core import bigframes.core.guid import bigframes.core.joins.row_identity @@ -42,7 +43,8 @@ def join_by_column( "right", ], sort: bool = False, - get_both_join_key_cols: bool = False, + coalesce_join_keys: bool = True, + allow_row_identity_join: bool = True, ) -> Tuple[ core.ArrayValue, typing.Sequence[str], @@ -56,8 +58,11 @@ def join_by_column( right: Expression for right table to join. right_column_ids: Column IDs (not label) to join by. how: The type of join to perform. - get_both_join_key_cols: if set to True, returned column ids will contain + coalesce_join_keys: if set to False, returned column ids will contain both left and right join key columns. + allow_row_identity_join (bool): + If True, allow matching by row identity. Set to False to always + perform a true JOIN in generated SQL. Returns: The joined expression and the objects needed to interpret it. @@ -66,19 +71,22 @@ def join_by_column( * Sequence[str]: Column IDs of the coalesced join columns. Sometimes either the left/right table will have missing rows. This column pulls the non-NULL value from either left/right. - If get_both_join_key_cols is True, will return uncombined left and + If coalesce_join_keys is False, will return uncombined left and right key columns. * Tuple[Callable, Callable]: For a given column ID from left or right, respectively, return the new column id from the combined expression. """ - if ( - how in bigframes.core.joins.row_identity.SUPPORTED_ROW_IDENTITY_HOW + allow_row_identity_join + and how in bigframes.core.joins.row_identity.SUPPORTED_ROW_IDENTITY_HOW and left.table.equals(right.table) - # Compare ibis expressions for left/right columns because its possible that - # they both have the same names but were modified in different ways. + # Make sure we're joining on exactly the same column(s), at least with + # regards to value its possible that they both have the same names but + # were modified in different ways. Ignore differences in the names. and all( - left.get_any_column(lcol).equals(right.get_any_column(rcol)) + left.get_any_column(lcol) + .name("index") + .equals(right.get_any_column(rcol).name("index")) for lcol, rcol in zip(left_column_ids, right_column_ids) ) ): @@ -86,15 +94,42 @@ def join_by_column( get_column_left, get_column_right, ) = bigframes.core.joins.row_identity.join_by_row_identity(left, right, how=how) - original_ordering = combined_expr._ordering + left_join_keys = [ + combined_expr.get_column(get_column_left(col)) for col in left_column_ids + ] + right_join_keys = [ + combined_expr.get_column(get_column_right(col)) for col in right_column_ids + ] + join_key_cols = get_join_cols( + left_join_keys, right_join_keys, how, coalesce_join_keys + ) + join_key_ids = [col.get_name() for col in join_key_cols] + combined_expr = combined_expr.projection( + [*join_key_cols, *combined_expr.columns] + ) + if sort: + combined_expr = combined_expr.order_by( + [ + core.OrderingColumnReference(join_col_id) + for join_col_id in join_key_ids + ] + ) + return ( + combined_expr, + join_key_ids, + ( + get_column_left, + get_column_right, + ), + ) else: # Generate offsets if non-default ordering is applied # Assumption, both sides are totally ordered, otherwise offsets will be nondeterministic left_table = left.to_ibis_expr( - ordering_mode="ordered_col", order_col_name=core.ORDER_ID_COLUMN + ordering_mode="string_encoded", order_col_name=core.ORDER_ID_COLUMN ) right_table = right.to_ibis_expr( - ordering_mode="ordered_col", order_col_name=core.ORDER_ID_COLUMN + ordering_mode="string_encoded", order_col_name=core.ORDER_ID_COLUMN ) join_conditions = [ value_to_join_key(left_table[left_index]) @@ -144,66 +179,105 @@ def get_column_right(key: str) -> str: return key left_ordering_encoding_size = ( - left._ordering.ordering_encoding_size - or bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH + left._ordering.string_encoding.length + if left._ordering.is_string_encoded + else bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH ) right_ordering_encoding_size = ( - right._ordering.ordering_encoding_size - or bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH + right._ordering.string_encoding.length + if right._ordering.is_string_encoded + else bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH ) # Preserve original ordering accross joins. left_order_id = get_column_left(core.ORDER_ID_COLUMN) right_order_id = get_column_right(core.ORDER_ID_COLUMN) new_order_id_col = _merge_order_ids( - combined_table[left_order_id], + typing.cast(ibis_types.StringColumn, combined_table[left_order_id]), left_ordering_encoding_size, - combined_table[right_order_id], + typing.cast(ibis_types.StringColumn, combined_table[right_order_id]), right_ordering_encoding_size, how, ) new_order_id = new_order_id_col.get_name() if new_order_id is None: raise ValueError("new_order_id unexpectedly has no name") + hidden_columns = (new_order_id_col,) - original_ordering = core.ExpressionOrdering( - ordering_id_column=core.OrderingColumnReference(new_order_id) - if (new_order_id_col is not None) - else None, - ordering_encoding_size=left_ordering_encoding_size - + right_ordering_encoding_size, + ordering = core.ExpressionOrdering( + # Order id is non-nullable but na_last=False generates simpler sql with current impl + ordering_value_columns=[ + core.OrderingColumnReference(new_order_id, na_last=False) + ], + total_ordering_columns=frozenset([new_order_id]), + string_encoding=core.StringEncoding( + True, left_ordering_encoding_size + right_ordering_encoding_size + ), + ) + + left_join_keys = [ + combined_table[get_column_left(col)] for col in left_column_ids + ] + right_join_keys = [ + combined_table[get_column_right(col)] for col in right_column_ids + ] + join_key_cols = get_join_cols( + left_join_keys, right_join_keys, how, coalesce_join_keys + ) + # We could filter out the original join columns, but predicates/ordering + # might still reference them in implicit joins. + columns = ( + join_key_cols + + [combined_table[get_column_left(col.get_name())] for col in left.columns] + + [ + combined_table[get_column_right(col.get_name())] + for col in right.columns + ] ) combined_expr = core.ArrayValue( left._session, combined_table, + columns=columns, hidden_ordering_columns=hidden_columns, + ordering=ordering, + ) + if sort: + combined_expr = combined_expr.order_by( + [ + core.OrderingColumnReference(join_key_col.get_name()) + for join_key_col in join_key_cols + ] + ) + return ( + combined_expr, + [key.get_name() for key in join_key_cols], + (get_column_left, get_column_right), ) + +def get_join_cols( + left_join_cols: typing.Iterable[ibis_types.Value], + right_join_cols: typing.Iterable[ibis_types.Value], + how: str, + coalesce_join_keys: bool = True, +) -> typing.List[ibis_types.Value]: join_key_cols: list[ibis_types.Value] = [] - for lcol, rcol in zip(left_column_ids, right_column_ids): - if get_both_join_key_cols: + for left_col, right_col in zip(left_join_cols, right_join_cols): + if not coalesce_join_keys: join_key_cols.append( - combined_expr.get_column(get_column_left(lcol)).name( - bigframes.core.guid.generate_guid(prefix="index_") - ) + left_col.name(bigframes.core.guid.generate_guid(prefix="index_")) ) join_key_cols.append( - combined_expr.get_column(get_column_right(rcol)).name( - bigframes.core.guid.generate_guid(prefix="index_") - ) + right_col.name(bigframes.core.guid.generate_guid(prefix="index_")) ) else: if how == "left" or how == "inner": join_key_cols.append( - combined_expr.get_column(get_column_left(lcol)).name( - bigframes.core.guid.generate_guid(prefix="index_") - ) + left_col.name(bigframes.core.guid.generate_guid(prefix="index_")) ) elif how == "right": join_key_cols.append( - combined_expr.get_column(get_column_right(rcol)).name( - bigframes.core.guid.generate_guid(prefix="index_") - ) + right_col.name(bigframes.core.guid.generate_guid(prefix="index_")) ) elif how == "outer": # The left index and the right index might contain null values, for @@ -211,48 +285,25 @@ def get_column_right(key: str) -> str: # these to take the index value from either column. # Use a random name in case the left index and the right index have the # same name. In such a case, _x and _y suffixes will already be used. - join_key_cols.append( - ibis.coalesce( - combined_expr.get_column(get_column_left(lcol)), - combined_expr.get_column(get_column_right(rcol)), - ).name(bigframes.core.guid.generate_guid(prefix="index_")) - ) + # Don't need to coalesce if they are exactly the same column. + if left_col.name("index").equals(right_col.name("index")): + join_key_cols.append( + left_col.name( + bigframes.core.guid.generate_guid(prefix="index_") + ) + ) + else: + join_key_cols.append( + ibis.coalesce( + left_col, + right_col, + ).name(bigframes.core.guid.generate_guid(prefix="index_")) + ) else: - raise ValueError(f"Unexpected join type: {how}") - - # We could filter out the original join columns, but predicates/ordering - # might still reference them in implicit joins. - columns = ( - join_key_cols - + [ - combined_expr.get_column(get_column_left(key)) - for key in left.column_names.keys() - ] - + [ - combined_expr.get_column(get_column_right(key)) - for key in right.column_names.keys() - ] - ) - - if sort: - ordering = original_ordering.with_ordering_columns( - [ - core.OrderingColumnReference(join_key_col.get_name()) - for join_key_col in join_key_cols - ] - ) - else: - ordering = original_ordering - - combined_expr_builder = combined_expr.builder() - combined_expr_builder.columns = columns - combined_expr_builder.ordering = ordering - combined_expr = combined_expr_builder.build() - return ( - combined_expr, - [key.get_name() for key in join_key_cols], - (get_column_left, get_column_right), - ) + raise ValueError( + f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}" + ) + return join_key_cols def value_to_join_key(value: ibis_types.Value): @@ -263,19 +314,31 @@ def value_to_join_key(value: ibis_types.Value): def _merge_order_ids( - left_id: ibis_types.Value, + left_id: ibis_types.StringColumn, left_encoding_size: int, - right_id: ibis_types.Value, + right_id: ibis_types.StringColumn, right_encoding_size: int, how: str, -) -> ibis_types.StringValue: +) -> ibis_types.StringColumn: if how == "right": return _merge_order_ids( right_id, right_encoding_size, left_id, left_encoding_size, "left" ) - return ( - ( - bigframes.core.ordering.stringify_order_id(left_id, left_encoding_size) - + bigframes.core.ordering.stringify_order_id(right_id, right_encoding_size) + + if how == "left": + right_id = typing.cast( + ibis_types.StringColumn, + right_id.fillna(ibis_types.literal(":" * right_encoding_size)), + ) + elif how != "inner": # outer join + left_id = typing.cast( + ibis_types.StringColumn, + left_id.fillna(ibis_types.literal(":" * left_encoding_size)), ) - ).name(bigframes.core.guid.generate_guid(prefix="bigframes_ordering_id_")) + right_id = typing.cast( + ibis_types.StringColumn, + right_id.fillna(ibis_types.literal(":" * right_encoding_size)), + ) + return (left_id + right_id).name( + bigframes.core.guid.generate_guid(prefix="bigframes_ordering_id_") + ) diff --git a/bigframes/core/ordering.py b/bigframes/core/ordering.py index 6f8d35e52ae..d5f07ecf91a 100644 --- a/bigframes/core/ordering.py +++ b/bigframes/core/ordering.py @@ -14,7 +14,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum import math import typing @@ -28,6 +28,8 @@ # Sufficient to store any value up to 2^63 DEFAULT_ORDERING_ID_LENGTH: int = math.ceil(63 * math.log(2, ORDERING_ID_STRING_BASE)) +STABLE_SORTS = ["mergesort", "stable"] + class OrderingDirection(Enum): ASC = 1 @@ -61,121 +63,162 @@ def with_reverse(self): ) +# Encoding classes specify additional properties for some ordering representations @dataclass(frozen=True) -class ExpressionOrdering: - """Immutable object that holds information about the ordering of rows in a ArrayValue object.""" +class StringEncoding: + """String encoded order ids are fixed length and can be concat together in joins.""" - ordering_value_columns: Sequence[OrderingColumnReference] = () - ordering_id_column: Optional[OrderingColumnReference] = None - is_sequential: bool = False + is_encoded: bool = False # Encoding size must be tracked in order to know what how to combine ordering ids across tables (eg how much to pad when combining different length). # Also will be needed to determine when length is too large and need to compact ordering id with a ROW_NUMBER operation. - ordering_encoding_size: int = DEFAULT_ORDERING_ID_LENGTH + length: int = DEFAULT_ORDERING_ID_LENGTH + + +@dataclass(frozen=True) +class IntegerEncoding: + """Integer encoded order ids are guaranteed non-negative.""" + + is_encoded: bool = False + is_sequential: bool = False + - def with_is_sequential(self, is_sequential: bool): +@dataclass(frozen=True) +class ExpressionOrdering: + """Immutable object that holds information about the ordering of rows in a ArrayValue object.""" + + ordering_value_columns: Sequence[OrderingColumnReference] = () + integer_encoding: IntegerEncoding = IntegerEncoding(False) + string_encoding: StringEncoding = StringEncoding(False) + # A table has a total ordering defined by the identities of a set of 1 or more columns. + # These columns must always be part of the ordering, in order to guarantee that the ordering is total. + # Therefore, any modifications(or drops) done to these columns must result in hidden copies being made. + total_ordering_columns: frozenset[str] = field(default_factory=frozenset) + + def with_non_sequential(self): """Create a copy that is marked as non-sequential. This is useful when filtering, but not sorting, an expression. """ - return ExpressionOrdering( - self.ordering_value_columns, - self.ordering_id_column, - is_sequential, - ordering_encoding_size=self.ordering_encoding_size, - ) + if self.integer_encoding.is_sequential: + return ExpressionOrdering( + self.ordering_value_columns, + integer_encoding=IntegerEncoding( + self.integer_encoding.is_encoded, is_sequential=False + ), + total_ordering_columns=self.total_ordering_columns, + ) + + return self def with_ordering_columns( self, ordering_value_columns: Sequence[OrderingColumnReference] = (), stable: bool = False, - ): - """Creates a new ordering that preserves ordering id, but replaces ordering value column list.""" + ) -> ExpressionOrdering: + """Creates a new ordering that reorders by the given columns. + + Args: + ordering_value_columns: + In decreasing precedence order, the values used to sort the ordering + stable: + If True, will use apply a stable sorting, using the old ordering where + the new ordering produces ties. Otherwise, ties will be resolved in + a performance maximizing way, + + Returns: + Modified ExpressionOrdering + """ + col_ids_new = [ + ordering_ref.column_id for ordering_ref in ordering_value_columns + ] if stable: - col_ids_new = [ - ordering_ref.column_id for ordering_ref in ordering_value_columns - ] # Only reference each column once, so discard old referenc if there is a new reference old_ordering_keep = [ ordering_ref for ordering_ref in self.ordering_value_columns if ordering_ref.column_id not in col_ids_new ] - new_ordering = (*ordering_value_columns, *old_ordering_keep) - else: # Not stable, so discard old ordering completely - new_ordering = tuple(ordering_value_columns) + else: + # New ordering needs to keep all total ordering columns no matter what. + # All other old ordering references can be discarded as does not need + # to be a stable sort. + old_ordering_keep = [ + ordering_ref + for ordering_ref in self.ordering_value_columns + if (ordering_ref.column_id not in col_ids_new) + and (ordering_ref.column_id in self.total_ordering_columns) + ] + new_ordering = (*ordering_value_columns, *old_ordering_keep) return ExpressionOrdering( new_ordering, - self.ordering_id_column, - is_sequential=False, - ordering_encoding_size=self.ordering_encoding_size, - ) - - def with_ordering_id(self, ordering_id: str): - """Creates a new ordering that preserves other properties, but with a different ordering id. - - Useful when reprojecting ordering for implicit joins. - """ - return ExpressionOrdering( - self.ordering_value_columns, - OrderingColumnReference(ordering_id), - is_sequential=self.is_sequential, - ordering_encoding_size=self.ordering_encoding_size, + total_ordering_columns=self.total_ordering_columns, ) def with_reverse(self): """Reverses the ordering.""" return ExpressionOrdering( tuple([col.with_reverse() for col in self.ordering_value_columns]), - self.ordering_id_column.with_reverse() - if self.ordering_id_column is not None - else None, - is_sequential=False, - ordering_encoding_size=self.ordering_encoding_size, + total_ordering_columns=self.total_ordering_columns, + ) + + def with_column_remap(self, mapping: typing.Mapping[str, str]): + new_value_columns = [ + col.with_name(mapping.get(col.column_id, col.column_id)) + for col in self.ordering_value_columns + ] + new_total_order = frozenset( + mapping.get(col_id, col_id) for col_id in self.total_ordering_columns + ) + return ExpressionOrdering( + new_value_columns, + integer_encoding=self.integer_encoding, + string_encoding=self.string_encoding, + total_ordering_columns=new_total_order, ) @property - def ordering_id(self) -> Optional[str]: - return self.ordering_id_column.column_id if self.ordering_id_column else None + def total_order_col(self) -> Optional[OrderingColumnReference]: + """Returns column id of columns that defines total ordering, if such as column exists""" + if len(self.ordering_value_columns) != 1: + return None + order_ref = self.ordering_value_columns[0] + if order_ref.direction != OrderingDirection.ASC: + return None + return order_ref @property - def order_id_defined(self) -> bool: - """True if ordering is fully defined in ascending order by its ordering id.""" - return bool( - self.ordering_id_column - and (not self.ordering_value_columns) - and self.ordering_id_column.direction == OrderingDirection.ASC - ) + def is_string_encoded(self) -> bool: + """True if ordering is fully defined by a fixed length string column.""" + return self.string_encoding.is_encoded + + @property + def is_sequential(self) -> bool: + return self.integer_encoding.is_encoded and self.integer_encoding.is_sequential @property def all_ordering_columns(self) -> Sequence[OrderingColumnReference]: - return ( - list(self.ordering_value_columns) - if self.ordering_id_column is None - else [*self.ordering_value_columns, self.ordering_id_column] - ) + return list(self.ordering_value_columns) -def stringify_order_id( - order_id: ibis_types.Value, length: int = DEFAULT_ORDERING_ID_LENGTH -) -> ibis_types.StringValue: +def encode_order_string( + order_id: ibis_types.IntegerColumn, length: int = DEFAULT_ORDERING_ID_LENGTH +) -> ibis_types.StringColumn: """Converts an order id value to string if it is not already a string. MUST produced fixed-length strings.""" - if order_id.type().is_int64(): - # This is very inefficient encoding base-10 string uses only 10 characters per byte(out of 256 bit combinations) - # Furthermore, if know tighter bounds on order id are known, can produce smaller strings. - # 19 characters chosen as it can represent any positive Int64 in base-10 - # For missing values, ":" * 19 is used as it is larger than any other value this function produces, so null values will be last. - string_order_id = ( - typing.cast( - ibis_types.StringValue, - typing.cast(ibis_types.IntegerValue, order_id).cast(ibis_dtypes.string), - ) - .lpad(length, "0") - .fillna(ibis_types.literal(":" * length)) - ) - else: - string_order_id = ( - typing.cast(ibis_types.StringValue, order_id) - .lpad(length, "0") - .fillna(ibis_types.literal(":" * length)) - ) - return typing.cast(ibis_types.StringValue, string_order_id) + # This is very inefficient encoding base-10 string uses only 10 characters per byte(out of 256 bit combinations) + # Furthermore, if know tighter bounds on order id are known, can produce smaller strings. + # 19 characters chosen as it can represent any positive Int64 in base-10 + # For missing values, ":" * 19 is used as it is larger than any other value this function produces, so null values will be last. + string_order_id = typing.cast( + ibis_types.StringValue, + order_id.cast(ibis_dtypes.string), + ).lpad(length, "0") + return typing.cast(ibis_types.StringColumn, string_order_id) + + +def reencode_order_string( + order_id: ibis_types.StringColumn, length: int +) -> ibis_types.StringColumn: + return typing.cast( + ibis_types.StringColumn, + (typing.cast(ibis_types.StringValue, order_id).lpad(length, "0")), + ) diff --git a/bigframes/core/reshape/__init__.py b/bigframes/core/reshape/__init__.py index 1443a8e08fc..339ce7466a1 100644 --- a/bigframes/core/reshape/__init__.py +++ b/bigframes/core/reshape/__init__.py @@ -14,48 +14,107 @@ from __future__ import annotations import typing -from typing import Iterable, Literal, Union +from typing import Iterable, Literal, Optional, Union +import bigframes.constants as constants +import bigframes.core as core +import bigframes.core.utils as utils import bigframes.dataframe +import bigframes.operations.aggregations as agg_ops import bigframes.series @typing.overload def concat( - objs: Iterable[bigframes.dataframe.DataFrame], *, join, ignore_index + objs: Iterable[bigframes.series.Series], + *, + axis: typing.Literal["index", 0] = ..., + join=..., + ignore_index=..., +) -> bigframes.series.Series: + ... + + +@typing.overload +def concat( + objs: Iterable[bigframes.dataframe.DataFrame], + *, + axis: typing.Literal["index", 0] = ..., + join=..., + ignore_index=..., ) -> bigframes.dataframe.DataFrame: ... @typing.overload def concat( - objs: Iterable[bigframes.series.Series], *, join, ignore_index -) -> bigframes.series.Series: + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], + *, + axis: typing.Literal["columns", 1], + join=..., + ignore_index=..., +) -> bigframes.dataframe.DataFrame: ... +@typing.overload def concat( - objs: Union[ - Iterable[bigframes.dataframe.DataFrame], Iterable[bigframes.series.Series] - ], + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], *, + axis=..., + join=..., + ignore_index=..., +) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: + ... + + +def concat( + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], + *, + axis: typing.Union[str, int] = 0, join: Literal["inner", "outer"] = "outer", ignore_index: bool = False, ) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: - contains_dataframes = any( - isinstance(x, bigframes.dataframe.DataFrame) for x in objs - ) - if not contains_dataframes: - # Special case, all series, so align everything into single column even if labels don't match - series = typing.cast(typing.Iterable[bigframes.series.Series], objs) - names = {s.name for s in series} - # For series case, labels are stripped if they don't all match - if len(names) > 1: - blocks = [s._block.with_column_labels([None]) for s in series] - else: - blocks = [s._block for s in series] + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + contains_dataframes = any( + isinstance(x, bigframes.dataframe.DataFrame) for x in objs + ) + if not contains_dataframes: + # Special case, all series, so align everything into single column even if labels don't match + series = typing.cast(typing.Iterable[bigframes.series.Series], objs) + names = {s.name for s in series} + # For series case, labels are stripped if they don't all match + if len(names) > 1: + blocks = [s._block.with_column_labels([None]) for s in series] + else: + blocks = [s._block for s in series] + block = blocks[0].concat(blocks[1:], how=join, ignore_index=ignore_index) + return bigframes.series.Series(block) + blocks = [obj._block for obj in objs] block = blocks[0].concat(blocks[1:], how=join, ignore_index=ignore_index) - return bigframes.series.Series(block) - blocks = [obj._block for obj in objs] - block = blocks[0].concat(blocks[1:], how=join, ignore_index=ignore_index) - return bigframes.dataframe.DataFrame(block) + return bigframes.dataframe.DataFrame(block) + else: + # Note: does not validate inputs + block_list = [obj._block for obj in objs] + block = block_list[0] + for rblock in block_list[1:]: + combined_index, _ = block.index.join(rblock.index, how=join) + block = combined_index._block + return bigframes.dataframe.DataFrame(block) + + +def cut( + x: bigframes.series.Series, + bins: int, + *, + labels: Optional[bool] = None, +) -> bigframes.series.Series: + if bins <= 0: + raise ValueError("`bins` should be a positive integer.") + + if labels is not False: + raise NotImplementedError( + f"Only labels=False is supported in BigQuery DataFrames so far. {constants.FEEDBACK_LINK}" + ) + return x._apply_window_op(agg_ops.CutOp(bins), window_spec=core.WindowSpec()) diff --git a/bigframes/core/scalar.py b/bigframes/core/scalar.py index 67ced525b99..6dfbd31b77b 100644 --- a/bigframes/core/scalar.py +++ b/bigframes/core/scalar.py @@ -15,10 +15,14 @@ from __future__ import annotations import typing -from typing import Any +from typing import Any, Optional +import google.cloud.bigquery as bigquery import ibis.expr.types as ibis_types +import bigframes +import bigframes.formatting_helpers as formatter + if typing.TYPE_CHECKING: import bigframes.session @@ -29,19 +33,36 @@ class DeferredScalar: def __init__(self, value: ibis_types.Scalar, session: bigframes.session.Session): self._value = value self._session = session + self._query_job: Optional[bigquery.QueryJob] = None + + @property + def query_job(self) -> Optional[bigquery.QueryJob]: + """BigQuery job metadata for the most recent query.""" + if self._query_job is None: + self._query_job = self._compute_dry_run() + return self._query_job def __repr__(self) -> str: """Converts a Series to a string.""" # TODO(swast): Add a timeout here? If the query is taking a long time, # maybe we just print the job metadata that we have so far? - return repr(self.compute()) + opts = bigframes.options.display + if opts.repr_mode == "deferred": + return formatter.repr_query_job(self.query_job) + else: + return repr(self.to_pandas()) - def compute(self) -> Any: + def to_pandas(self) -> Any: """Executes deferred operations and downloads the resulting scalar.""" - result, _ = self._session._start_query(self._value.compile()) + result, query_job = self._session._start_query(self._value.compile()) + self._query_job = query_job df = self._session._rows_to_dataframe(result) return df.iloc[0, 0] + def _compute_dry_run(self): + job_config = bigquery.QueryJobConfig(dry_run=True) + return self._session._start_query(self._value.compile(), job_config=job_config) + # All public APIs return Any at present # Later implementation may sometimes return a lazy scalar diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py new file mode 100644 index 00000000000..8ad58fb166e --- /dev/null +++ b/bigframes/core/utils.py @@ -0,0 +1,22 @@ +# 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 typing + + +def get_axis_number(axis: typing.Union[str, int, None]) -> typing.Literal[0, 1]: + if axis in {0, "index", "rows", None}: + return 0 + elif axis in {1, "columns"}: + return 1 + raise ValueError(f"Not a valid axis: {axis}") diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 0155aafa0b0..cb062f09249 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -16,7 +16,6 @@ from __future__ import annotations -import random import re import textwrap import typing @@ -33,13 +32,13 @@ ) import google.cloud.bigquery as bigquery -import ibis.expr.datatypes as ibis_dtypes import numpy -import pandas as pd +import pandas import typing_extensions import bigframes import bigframes._config.display_options as display_options +import bigframes.constants as constants import bigframes.core import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks @@ -47,9 +46,12 @@ import bigframes.core.guid import bigframes.core.indexers as indexers import bigframes.core.indexes as indexes +import bigframes.core.io import bigframes.core.joins as joins import bigframes.core.ordering as order +import bigframes.core.utils as utils import bigframes.dtypes +import bigframes.formatting_helpers as formatter import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops import bigframes.series @@ -64,11 +66,21 @@ # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. # TODO(tbergeron): Convert to bytes-based limit -MAX_INLINE_DF_SIZE = 5000 +# TODO(swast): Address issues with string escaping and empty tables before +# re-enabling inline data (ibis.memtable) feature. +MAX_INLINE_DF_SIZE = -1 -LevelsType = typing.Union[str, int, typing.Sequence[typing.Union[str, int]]] +LevelType = typing.Union[str, int] +LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]] SingleItemValue = Union[bigframes.series.Series, int, float, Callable] +ERROR_IO_ONLY_GS_PATHS = f"Only Google Cloud Storage (gs://...) paths are supported. {constants.FEEDBACK_LINK}" +ERROR_IO_REQUIRES_WILDCARD = ( + "Google Cloud Storage path must contain a wildcard '*' character. See: " + "https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#export_data_statement" + f"{constants.FEEDBACK_LINK}" +) + # Inherits from pandas DataFrame so that we can use the same docstrings. class DataFrame(vendored_pandas_frame.DataFrame): @@ -87,7 +99,9 @@ def __init__( session: typing.Optional[bigframes.session.Session] = None, ): if copy is not None and not copy: - raise ValueError("DataFrame constructor only supports copy=True") + raise ValueError( + f"DataFrame constructor only supports copy=True. {constants.FEEDBACK_LINK}" + ) # Check to see if constructing from BigQuery-backed objects before # falling back to pandas constructor @@ -106,7 +120,9 @@ def __init__( ): if not all(isinstance(data[key], bf_series.Series) for key in data.keys()): # TODO(tbergeron): Support local list/series data by converting to memtable. - raise NotImplementedError("Cannot mix Series with other types.") + raise NotImplementedError( + f"Cannot mix Series with other types. {constants.FEEDBACK_LINK}" + ) keys = list(data.keys()) first_label, first_series = keys[0], data[keys[0]] block = ( @@ -128,7 +144,8 @@ def __init__( if block: if index: raise NotImplementedError( - "DataFrame 'index' constructor parameter not supported when passing BigQuery-backed objects" + "DataFrame 'index' constructor parameter not supported " + f"when passing BigQuery-backed objects. {constants.FEEDBACK_LINK}" ) if columns: block = block.select_columns(list(columns)) # type:ignore @@ -141,7 +158,7 @@ def __init__( else: import bigframes.pandas - pd_dataframe = pd.DataFrame( + pd_dataframe = pandas.DataFrame( data=data, index=index, # type:ignore columns=columns, # type:ignore @@ -182,12 +199,14 @@ def _find_indices( def _resolve_label_exact(self, label) -> str: matches = self._block.label_to_col_id.get(label, []) if len(matches) != 1: - raise ValueError("Index data must be 1-dimensional") + raise ValueError( + f"Index data must be 1-dimensional. {constants.FEEDBACK_LINK}" + ) return matches[0] def _sql_names( self, - columns: Union[blocks.Label, Sequence[blocks.Label], pd.Index], + columns: Union[blocks.Label, Sequence[blocks.Label], pandas.Index], tolerance: bool = False, ) -> Sequence[str]: """Retrieve sql name (column name in BQ schema) of column(s).""" @@ -215,11 +234,11 @@ def iloc(self) -> indexers.ILocDataFrameIndexer: return indexers.ILocDataFrameIndexer(self) @property - def dtypes(self) -> pd.Series: - return pd.Series(data=self._block.dtypes, index=self._block.column_labels) + def dtypes(self) -> pandas.Series: + return pandas.Series(data=self._block.dtypes, index=self._block.column_labels) @property - def columns(self) -> pd.Index: + def columns(self) -> pandas.Index: return self.dtypes.index @property @@ -253,19 +272,21 @@ def astype( ) -> DataFrame: return self._apply_to_rows(ops.AsTypeOp(dtype)) - def to_sql_query( + def _to_sql_query( self, always_include_index: bool ) -> Tuple[str, List[Tuple[str, bool]]]: - """Compiles this dataframe's expression tree to SQL, optionally - including unnamed index columns + """Compiles this DataFrame's expression tree to SQL, optionally + including unnamed index columns. Args: - always_include_index: whether to include unnamed index columns. - If False, only named indexes are included. + always_include_index (bool): + whether to include unnamed index columns. If False, only named + indexes are included. - Returns: a tuple of (sql_string, index_column_list). Each entry in the - index column list is a tuple of (column_name, named). If named is - is false, then the column name exists only in SQL""" + Returns: a tuple of (sql_string, index_column_list) + Each entry in the index column list is a tuple of (column_name, named). + If named is false, then the column name exists only in SQL + """ # Has to be unordered as it is impossible to order the sql without # including metadata columns in selection with ibis. ibis_expr = self._block.expr.to_ibis_expr(ordering_mode="unordered") @@ -325,15 +346,25 @@ def to_sql_query( @property def sql(self) -> str: - """Compiles this dataframe's expression tree to SQL""" - sql, _ = self.to_sql_query(always_include_index=False) + """Compiles this DataFrame's expression tree to SQL.""" + sql, _ = self._to_sql_query(always_include_index=False) return sql @property def query_job(self) -> Optional[bigquery.QueryJob]: - """BigQuery job metadata for the most recent query.""" + """BigQuery job metadata for the most recent query. + + Returns: + The most recent `QueryJob + `_. + """ + if self._query_job is None: + self._set_internal_query_job(self._compute_dry_run()) return self._query_job + def _set_internal_query_job(self, query_job: bigquery.QueryJob): + self._query_job = query_job + @typing.overload def __getitem__(self, key: bigframes.series.Series) -> DataFrame: ... @@ -343,7 +374,7 @@ def __getitem__(self, key: Sequence[blocks.Label]) -> DataFrame: # type:ignore ... @typing.overload - def __getitem__(self, key: pd.Index) -> DataFrame: # type:ignore + def __getitem__(self, key: pandas.Index) -> DataFrame: # type:ignore ... @typing.overload @@ -356,7 +387,7 @@ def __getitem__( blocks.Label, Sequence[blocks.Label], # Index of column labels can be treated the same as a sequence of column labels. - pd.Index, + pandas.Index, bigframes.series.Series, ], ) -> Union[bigframes.series.Series, "DataFrame"]: @@ -397,8 +428,10 @@ def __getitem__( # Bool Series selects rows def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame: - if not key._to_ibis_expr().type() == ibis_dtypes.bool: - raise ValueError("Only boolean series currently supported for indexing.") + if not key.dtype == pandas.BooleanDtype(): + raise NotImplementedError( + f"Only boolean series currently supported for indexing. {constants.FEEDBACK_LINK}" + ) # TODO: enforce stricter alignment combined_index, ( get_column_left, @@ -413,16 +446,12 @@ def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame: def __getattr__(self, key: str): if key in self._block.column_labels: return self.__getitem__(key) - elif hasattr(pd.DataFrame, key): + elif hasattr(pandas.DataFrame, key): raise NotImplementedError( textwrap.dedent( f""" BigQuery DataFrames has not yet implemented an equivalent to - 'pandas.DataFrame.{key}'. Please check - https://github.com/googleapis/python-bigquery-dataframes/issues for - existing feature requests, or file your own. - Please include information about your use case, as well as - relevant code snippets. + 'pandas.DataFrame.{key}'. {constants.FEEDBACK_LINK} """ ) ) @@ -432,14 +461,21 @@ def __getattr__(self, key: str): def __repr__(self) -> str: """Converts a DataFrame to a string. Calls compute. - Only represents the first ``bigframes.options.display.max_rows``. + Only represents the first `bigframes.options.display.max_rows`. """ opts = bigframes.options.display max_results = opts.max_rows + if opts.repr_mode == "deferred": + return formatter.repr_query_job(self.query_job) # TODO(swast): pass max_columns and get the true column count back. Maybe # get 1 more column than we have requested so that pandas can add the # ... for us? - pandas_df, row_count = self._retrieve_repr_request_results(max_results) + pandas_df, row_count, query_job = self._block.retrieve_repr_request_results( + max_results + ) + + self._set_internal_query_job(query_job) + column_count = len(pandas_df.columns) with display_options.pandas_repr(opts): @@ -466,10 +502,17 @@ def _repr_html_(self) -> str: """ opts = bigframes.options.display max_results = bigframes.options.display.max_rows + if opts.repr_mode == "deferred": + return formatter.repr_query_job_html(self.query_job) # TODO(swast): pass max_columns and get the true column count back. Maybe # get 1 more column than we have requested so that pandas can add the # ... for us? - pandas_df, row_count = self._retrieve_repr_request_results(max_results) + pandas_df, row_count, query_job = self._block.retrieve_repr_request_results( + max_results + ) + + self._set_internal_query_job(query_job) + column_count = len(pandas_df.columns) with display_options.pandas_repr(opts): @@ -479,33 +522,6 @@ def _repr_html_(self) -> str: html_string += f"[{row_count} rows x {column_count} columns in total]" return html_string - def _retrieve_repr_request_results( - self, max_results: int - ) -> Tuple[pd.DataFrame, int]: - """ - Retrieves a pandas dataframe containing only max_results many rows for use - with printing methods. - - Returns a tuple of the dataframe and the overall number of rows of the query. - """ - # TODO(swast): Select a subset of columns if max_columns is less than the - # number of columns in the schema. - count = self.shape[0] - if count > max_results: - head_df = self.head(n=max_results) - computed_df, query_job = head_df._block.compute(max_results=max_results) - else: - head_df = self - computed_df, query_job = head_df._block.compute() - - formatted_df = computed_df.set_axis(self._block.column_labels, axis=1) - # don't update details when the cache is hit - if self.query_job is None or not query_job.cache_hit: - self._query_job = query_job - # we reset the axis and substitute the bf index name for the default - formatted_df.index.name = self.index.name - return formatted_df, count - def __setitem__(self, key: str, value: SingleItemValue): """Modify or insert a column into the DataFrame. @@ -527,6 +543,7 @@ def _apply_binop( return self._apply_series_binop(other, op, axis=axis) raise NotImplementedError( f"binary operation is not implemented on the second operand of type {type(other).__name__}." + f"{constants.FEEDBACK_LINK}" ) def _apply_scalar_binop(self, other: float | int, op: ops.BinaryOp) -> DataFrame: @@ -549,7 +566,9 @@ def _apply_series_binop( raise ValueError(f"Invalid input: axis {axis}.") if axis in ("columns", 1): - raise NotImplementedError("Row Series operations haven't been supported.") + raise NotImplementedError( + f"Row Series operations haven't been supported. {constants.FEEDBACK_LINK}" + ) joined_index, (get_column_left, get_column_right) = self._block.index.join( other._block.index, how="outer" @@ -573,6 +592,16 @@ def _apply_series_binop( block = block.with_index_labels(self.index.names) return DataFrame(block) + def eq(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.eq_op, axis=axis) + + def ne(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: + return self._apply_binop(other, ops.ne_op, axis=axis) + + __eq__ = eq # type: ignore + + __ne__ = ne # type: ignore + def le(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: return self._apply_binop(other, ops.le_op, axis=axis) @@ -661,12 +690,47 @@ def rmod(self, other: int | bigframes.series.Series, axis: str | int = "columns" __rmod__ = rmod - def compute(self) -> pd.DataFrame: - """Executes deferred operations and downloads the results.""" + def to_pandas( + self, + max_download_size: Optional[int] = None, + sampling_method: Optional[str] = None, + random_state: Optional[int] = None, + ) -> pandas.DataFrame: + """Write DataFrame to pandas DataFrame. + + Args: + max_download_size (int, default None): + Download size threshold in MB. If max_download_size is exceeded when downloading data + (e.g., to_pandas()), the data will be downsampled if + bigframes.options.sampling.enable_downsampling is True, otherwise, an error will be + raised. If set to a value other than None, this will supersede the global config. + sampling_method (str, default None): + Downsampling algorithms to be chosen from, the choices are: "head": This algorithm + returns a portion of the data from the beginning. It is fast and requires minimal + computations to perform the downsampling; "uniform": This algorithm returns uniform + random samples of the data. If set to a value other than None, this will supersede + the global config. + random_state (int, default None): + The seed for the uniform downsampling algorithm. If provided, the uniform method may + take longer to execute and require more computation. If set to a value other than + None, this will supersede the global config. + + Returns: + pandas.DataFrame: A pandas DataFrame with all rows and columns of this DataFrame if the + data_sampling_threshold_mb is not exceeded; otherwise, a pandas DataFrame with + downsampled rows and all columns of this DataFrame. + """ # TODO(orrbradford): Optimize this in future. Potentially some cases where we can return the stored query job - df, query_job = self._block.compute() - self._query_job = query_job - return df.set_axis(self._block.column_labels, axis=1) + df, query_job = self._block.to_pandas( + max_download_size=max_download_size, + sampling_method=sampling_method, + random_state=random_state, + ) + self._set_internal_query_job(query_job) + return df.set_axis(self._block.column_labels, axis=1, copy=False) + + def _compute_dry_run(self) -> bigquery.QueryJob: + return self._block._compute_dry_run() def copy(self) -> DataFrame: return DataFrame(self._block) @@ -677,12 +741,50 @@ def head(self, n: int = 5) -> DataFrame: def tail(self, n: int = 5) -> DataFrame: return typing.cast(DataFrame, self.iloc[-n:]) - def drop(self, *, columns: Union[str, Iterable[str]]) -> DataFrame: - if not _is_list_like(columns): - columns = [columns] # type:ignore - columns = list(columns) + def drop( + self, + labels: typing.Any = None, + *, + axis: typing.Union[int, str] = 0, + index: typing.Any = None, + columns: Union[blocks.Label, Iterable[blocks.Label]] = None, + level: typing.Optional[LevelType] = None, + ) -> DataFrame: + if labels: + if index or columns: + raise ValueError("Cannot specify both 'labels' and 'index'/'columns") + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + index = labels + else: + columns = labels + + block = self._block + if index: + level_id = self._resolve_levels(level or 0)[0] - block = self._block.drop_columns(self._sql_names(columns)) + if _is_list_like(index): + block, inverse_condition_id = block.apply_unary_op( + level_id, ops.IsInOp(index, match_nulls=True) + ) + block, condition_id = block.apply_unary_op( + inverse_condition_id, ops.invert_op + ) + else: + block, condition_id = block.apply_unary_op( + level_id, ops.partial_right(ops.ne_op, index) + ) + block = block.filter(condition_id, keep_null=True).select_columns( + self._block.value_columns + ) + if columns: + if not _is_list_like(columns): + columns = [columns] # type:ignore + columns = list(columns) + + block = block.drop_columns(self._sql_names(columns)) + if not index and not columns: + raise ValueError("Must specify 'labels' or 'index'/'columns") return DataFrame(block) def droplevel(self, level: LevelsType): @@ -722,7 +824,7 @@ def rename_axis( ) -> DataFrame: if len(kwargs) != 0: raise NotImplementedError( - "rename_axis does not currently support any keyword arguments." + f"rename_axis does not currently support any keyword arguments. {constants.FEEDBACK_LINK}" ) # limited implementation: the new index name is simply the 'mapper' parameter if _is_list_like(mapper): @@ -795,7 +897,7 @@ def _assign_series_join_on_index( # Update case, remove after copying into columns block = block.drop_columns([source_column]) - return DataFrame(block.with_index_labels([self.index.name])) + return DataFrame(block.with_index_labels(self.index.names)) def reset_index(self, *, drop: bool = False) -> DataFrame: block = self._block.reset_index(drop) @@ -814,9 +916,20 @@ def set_index( col_ids = [self._resolve_label_exact(key) for key in keys] return DataFrame(self._block.set_index(col_ids, append=append, drop=drop)) - def sort_index(self) -> DataFrame: + def sort_index( + self, ascending: bool = True, na_position: Literal["first", "last"] = "last" + ) -> DataFrame: + if na_position not in ["first", "last"]: + raise ValueError("Param na_position must be one of 'first' or 'last'") + direction = ( + order.OrderingDirection.ASC if ascending else order.OrderingDirection.DESC + ) + na_last = na_position == "last" index_columns = self._block.index_columns - ordering = [order.OrderingColumnReference(column) for column in index_columns] + ordering = [ + order.OrderingColumnReference(column, direction=direction, na_last=na_last) + for column in index_columns + ] return DataFrame(self._block.order_by(ordering)) def sort_values( @@ -824,6 +937,7 @@ def sort_values( by: str | typing.Sequence[str], *, ascending: bool | typing.Sequence[bool] = True, + kind: str = "quicksort", na_position: typing.Literal["first", "last"] = "last", ) -> DataFrame: if na_position not in {"first", "last"}: @@ -854,8 +968,9 @@ def sort_values( column_id, direction=direction, na_last=na_last ) ) - - return DataFrame(self._block.order_by(ordering)) + return DataFrame( + self._block.order_by(ordering, stable=kind in order.STABLE_SORTS) + ) def value_counts( self, @@ -902,7 +1017,7 @@ def any( else: frame = self._drop_non_bool() block = frame._block.aggregate_all_and_pivot( - agg_ops.any_op, dtype=pd.BooleanDtype() + agg_ops.any_op, dtype=pandas.BooleanDtype() ) return bigframes.series.Series(block.select_column("values")) @@ -912,7 +1027,7 @@ def all(self, *, bool_only: bool = False) -> bigframes.series.Series: else: frame = self._drop_non_bool() block = frame._block.aggregate_all_and_pivot( - agg_ops.all_op, dtype=pd.BooleanDtype() + agg_ops.all_op, dtype=pandas.BooleanDtype() ) return bigframes.series.Series(block.select_column("values")) @@ -932,6 +1047,20 @@ def mean(self, *, numeric_only: bool = False) -> bigframes.series.Series: block = frame._block.aggregate_all_and_pivot(agg_ops.mean_op) return bigframes.series.Series(block.select_column("values")) + def median( + self, *, numeric_only: bool = False, exact: bool = False + ) -> bigframes.series.Series: + if exact: + raise NotImplementedError( + f"Only approximate median is supported. {constants.FEEDBACK_LINK}" + ) + if not numeric_only: + frame = self._raise_on_non_numeric("median") + else: + frame = self._drop_non_numeric() + block = frame._block.aggregate_all_and_pivot(agg_ops.median_op) + return bigframes.series.Series(block.select_column("values")) + def std(self, *, numeric_only: bool = False) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("std") @@ -986,11 +1115,52 @@ def nunique(self) -> bigframes.series.Series: block = self._block.aggregate_all_and_pivot(agg_ops.nunique_op) return bigframes.series.Series(block.select_column("values")) - def _drop_non_numeric(self) -> DataFrame: + def agg( + self, func: str | typing.Sequence[str] + ) -> DataFrame | bigframes.series.Series: + if _is_list_like(func): + if any( + dtype not in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES + for dtype in self.dtypes + ): + raise NotImplementedError( + f"Multiple aggregations only supported on numeric columns. {constants.FEEDBACK_LINK}" + ) + aggregations = [agg_ops.AGGREGATIONS_LOOKUP[f] for f in func] + return DataFrame( + self._block.summarize( + self._block.value_columns, + aggregations, + ) + ) + else: + return bigframes.series.Series( + self._block.aggregate_all_and_pivot( + agg_ops.AGGREGATIONS_LOOKUP[typing.cast(str, func)] + ) + ) + + aggregate = agg + + def describe(self) -> DataFrame: + df_numeric = self._drop_non_numeric(keep_bool=False) + if len(df_numeric.columns) == 0: + raise NotImplementedError( + f"df.describe() currently only supports numeric values. {constants.FEEDBACK_LINK}" + ) + result = df_numeric.agg( + ["count", "mean", "std", "min", "25%", "50%", "75%", "max"] + ) + return typing.cast(DataFrame, result) + + def _drop_non_numeric(self, keep_bool=True) -> DataFrame: + types_to_keep = set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) + if not keep_bool: + types_to_keep -= set(bigframes.dtypes.BOOL_BIGFRAMES_TYPES) non_numeric_cols = [ col_id for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) - if dtype not in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES + if dtype not in types_to_keep ] return DataFrame(self._block.drop_columns(non_numeric_cols)) @@ -1008,7 +1178,8 @@ def _raise_on_non_numeric(self, op: str): for dtype in self._block.dtypes ): raise NotImplementedError( - f"'{op}' does not support non-numeric columns. Set 'numeric_only'=True to ignore non-numeric columns" + f"'{op}' does not support non-numeric columns. " + f"Set 'numeric_only'=True to ignore non-numeric columns. {constants.FEEDBACK_LINK}" ) return self @@ -1018,7 +1189,8 @@ def _raise_on_non_boolean(self, op: str): for dtype in self._block.dtypes ): raise NotImplementedError( - f"'{op}' does not support non-bool columns. Set 'bool_only'=True to ignore non-bool columns" + f"'{op}' does not support non-bool columns. " + f"Set 'bool_only'=True to ignore non-bool columns. {constants.FEEDBACK_LINK}" ) return self @@ -1044,10 +1216,12 @@ def merge( ) -> DataFrame: if on is None: if left_on is None or right_on is None: - raise ValueError("Must specify either on or left_on + right_on.") + raise ValueError("Must specify `on` or `left_on` + `right_on`.") else: if left_on is not None or right_on is not None: - raise ValueError("Can not pass both on and left_on + right_on params.") + raise ValueError( + "Can not pass both `on` and `left_on` + `right_on` params." + ) left_on, right_on = on, on left = self @@ -1075,7 +1249,7 @@ def merge( sort=sort, # In merging on the same column, it only returns 1 key column from coalesced both. # While if 2 different columns, both will be presented in the result. - get_both_join_key_cols=(left_on != right_on), + coalesce_join_keys=(left_on == right_on), ) # TODO(swast): Add suffixes to the column labels instead of reusing the # column IDs as the new labels. @@ -1138,14 +1312,57 @@ def _get_merged_col_labels( return left_col_labels + right_col_labels - def join(self, other: DataFrame, *, how: str = "left") -> DataFrame: - if not self.columns.intersection(other.columns).empty: - raise NotImplementedError("Deduping column names is not implemented") + def join( + self, other: DataFrame, *, on: Optional[str] = None, how: str = "left" + ) -> DataFrame: + left, right = self, other + if not left.columns.intersection(right.columns).empty: + raise NotImplementedError( + f"Deduping column names is not implemented. {constants.FEEDBACK_LINK}" + ) - left = self - right = other - combined_index, (get_column_left, get_column_right) = left._block.index.join( - right._block.index, how=how + # Join left columns with right index + if on is not None: + if other._block.index.nlevels != 1: + raise ValueError( + "Join on columns must match the index level of the other DataFrame. Join on column with multi-index haven't been supported." + ) + # Switch left index with on column + left_columns = left.columns + left_idx_original_names = left.index.names + left_idx_names_in_cols = [ + f"bigframes_left_idx_name_{i}" for i in range(len(left.index.names)) + ] + left.index.names = left_idx_names_in_cols + left = left.reset_index(drop=False) + left = left.set_index(on) + + # Join on index and switch back + combined_df = left._perform_join_by_index(right, how=how) + combined_df.index.name = on + combined_df = combined_df.reset_index(drop=False) + combined_df = combined_df.set_index(left_idx_names_in_cols) + + # To be consistent with Pandas + combined_df.index.names = ( + left_idx_original_names + if how in ("inner", "left") + else ([None] * len(combined_df.index.names)) + ) + + # Reorder columns + combined_df = combined_df[list(left_columns) + list(right.columns)] + return combined_df + + # Join left index with right index + if left._block.index.nlevels != right._block.index.nlevels: + raise ValueError("Index to join on must have the same number of levels.") + + return left._perform_join_by_index(right, how=how) + + def _perform_join_by_index(self, other: DataFrame, *, how: str = "left"): + combined_index, _ = self._block.index.join( + other._block.index, how=how, block_identity_join=True ) return DataFrame(combined_index._block) @@ -1178,7 +1395,7 @@ def _groupby_level( ): return groupby.DataFrameGroupBy( self._block, - self._resolve_levels(level), + by_col_ids=self._resolve_levels(level), as_index=as_index, dropna=dropna, ) @@ -1220,13 +1437,13 @@ def _groupby_series( matches = [*col_matches, *level_matches] if len(matches) != 1: raise ValueError( - f"GroupBy key {key} does not map to unambiguous column or index level" + f"GroupBy key {key} does not match a unique column or index level. BigQuery DataFrames only interprets lists of strings as column or index names, not directly as per-row group assignments." ) col_ids = [*col_ids, matches[0]] return groupby.DataFrameGroupBy( block, - col_ids, + by_col_ids=col_ids, as_index=as_index, dropna=dropna, ) @@ -1311,8 +1528,9 @@ def sample( ns = (n,) if n is not None else () fracs = (frac,) if frac is not None else () - - return self._split(ns=ns, fracs=fracs, random_state=random_state)[0] + return DataFrame( + self._block._split(ns=ns, fracs=fracs, random_state=random_state)[0] + ) def _split( self, @@ -1326,88 +1544,33 @@ def _split( At most one of ns and fracs can be passed in. If neither, default to ns = (1,). Return a list of sampled DataFrames. """ - if ns and fracs: - raise ValueError("Only one of 'ns' or 'fracs' parameter must be specified.") - - block = self._block - if not ns and not fracs: - ns = (1,) - - if ns: - sample_sizes = ns - else: - total_rows = block.shape[0] - # Round to nearest integer. "round half to even" rule applies. - # At least to be 1. - sample_sizes = [round(frac * total_rows) or 1 for frac in fracs] - - # Set random_state if it is not provided - if random_state is None: - random_state = random.randint(-(2**30), 2**30) - - # Create a new column with random_state value. - block, random_state_col = block.create_constant(random_state) - - # Create an ordering col and a new sum col which is ordering+random_state. - block, ordering_col = block.promote_offsets() - block, sum_col = block.apply_binary_op( - ordering_col, random_state_col, ops.add_op - ) - - # Apply hash method to sum col and order by it. - block, string_sum_col = block.apply_unary_op( - sum_col, ops.AsTypeOp("string[pyarrow]") - ) - block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op) - block = block.order_by([order.OrderingColumnReference(hash_string_sum_col)]) - - drop_cols = [ - random_state_col, - ordering_col, - sum_col, - string_sum_col, - hash_string_sum_col, - ] - block = block.drop_columns(drop_cols) - df = DataFrame(block) - - intervals = [] - cur = 0 - for sample_size in sample_sizes: - intervals.append((cur, cur + sample_size)) - cur += sample_size - - # DF.iloc[slice] always returns DF. - return [ - typing.cast(DataFrame, df.iloc[lower:upper]) for lower, upper in intervals - ] + blocks = self._block._split(ns=ns, fracs=fracs, random_state=random_state) + return [DataFrame(block) for block in blocks] - def to_pandas(self) -> pd.DataFrame: - """Writes DataFrame to Pandas DataFrame.""" - # TODO(chelsealin): Support block parameters. - # TODO(chelsealin): Add to_pandas_batches() API. - return self.compute() - - def to_csv(self, path_or_buf: str, *, index: bool = True) -> None: + def to_csv( + self, path_or_buf: str, sep=",", *, header: bool = True, index: bool = True + ) -> None: # TODO(swast): Can we support partition columns argument? # TODO(chelsealin): Support local file paths. # TODO(swast): Some warning that wildcard is recommended for large # query results? See: # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size if not path_or_buf.startswith("gs://"): - raise NotImplementedError( - "Only Google Cloud Storage (gs://...) paths are supported." - ) + raise NotImplementedError(ERROR_IO_ONLY_GS_PATHS) + if "*" not in path_or_buf: + raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) - source_table = self._execute_query(index=index) - job_config = bigquery.ExtractJobConfig( - destination_format=bigquery.DestinationFormat.CSV + result_table = self._run_io_query( + index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID ) - self._block.expr._session._extract_table( - source_table, - destination_uris=[path_or_buf], - job_config=job_config, + export_data_statement = bigframes.core.io.create_export_csv_statement( + f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", + uri=path_or_buf, + field_delimiter=sep, + header=header, ) + _, query_job = self._block.expr._session._start_query(export_data_statement) + self._set_internal_query_job(query_job) def to_json( self, @@ -1421,13 +1584,11 @@ def to_json( ) -> None: # TODO(swast): Can we support partition columns argument? # TODO(chelsealin): Support local file paths. - # TODO(swast): Some warning that wildcard is recommended for large - # query results? See: - # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size if not path_or_buf.startswith("gs://"): - raise NotImplementedError( - "Only Google Cloud Storage (gs://...) paths are supported." - ) + raise NotImplementedError(ERROR_IO_ONLY_GS_PATHS) + + if "*" not in path_or_buf: + raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) if lines is True and orient != "records": raise ValueError( @@ -1438,18 +1599,20 @@ def to_json( # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#to_json_string if lines is False: raise NotImplementedError( - "Only newline delimited JSON format is supported." + f"Only newline delimited JSON format is supported. {constants.FEEDBACK_LINK}" ) - source_table = self._execute_query(index=index) - job_config = bigquery.ExtractJobConfig( - destination_format=bigquery.DestinationFormat.NEWLINE_DELIMITED_JSON + result_table = self._run_io_query( + index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID ) - self._block.expr._session._extract_table( - source_table, - destination_uris=[path_or_buf], - job_config=job_config, + export_data_statement = bigframes.core.io.create_export_data_statement( + f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", + uri=path_or_buf, + format="JSON", + export_options={}, ) + _, query_job = self._block.expr._session._start_query(export_data_statement) + self._set_internal_query_job(query_job) def to_gbq( self, @@ -1457,6 +1620,7 @@ def to_gbq( *, if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", index: bool = True, + ordering_id: Optional[str] = None, ) -> None: if "." not in destination_table: raise ValueError( @@ -1480,12 +1644,12 @@ def to_gbq( ), ) - self._execute_query(index=index, job_config=job_config) + self._run_io_query(index=index, ordering_id=ordering_id, job_config=job_config) def to_numpy( self, dtype=None, copy=False, na_value=None, **kwargs ) -> numpy.ndarray: - return self.compute().to_numpy(dtype, copy, na_value, **kwargs) + return self.to_pandas().to_numpy(dtype, copy, na_value, **kwargs) __array__ = to_numpy @@ -1496,30 +1660,29 @@ def to_parquet(self, path: str, *, index: bool = True) -> None: # query results? See: # https://cloud.google.com/bigquery/docs/exporting-data#limit_the_exported_file_size if not path.startswith("gs://"): - raise NotImplementedError( - "Only Google Cloud Storage (gs://...) paths are supported." - ) + raise NotImplementedError(ERROR_IO_ONLY_GS_PATHS) - source_table = self._execute_query(index=index) - job_config = bigquery.ExtractJobConfig( - destination_format=bigquery.DestinationFormat.PARQUET - ) + if "*" not in path: + raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) - self._block.expr._session._extract_table( - source_table, - destination_uris=[path], - job_config=job_config, + result_table = self._run_io_query( + index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID ) + export_data_statement = bigframes.core.io.create_export_data_statement( + f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", + uri=path, + format="PARQUET", + export_options={}, + ) + _, query_job = self._block.expr._session._start_query(export_data_statement) + self._set_internal_query_job(query_job) def _apply_to_rows(self, operation: ops.UnaryOp): block = self._block.multi_apply_unary_op(self._block.value_columns, operation) return DataFrame(block) - def _execute_query( - self, index: bool, job_config: Optional[bigquery.job.QueryJobConfig] = None - ): - """Executes a query job presenting this dataframe and returns the destination - table.""" + def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: + """Create query text representing this dataframe for I/O.""" expr = self._block.expr session = expr._session columns = list(self._block.value_columns) @@ -1531,26 +1694,53 @@ def _execute_query( if index and self.index.name is not None: columns.extend(self._block.index_columns) column_labels.extend(self.index.names) - # TODO(chelsealin): normalize the file formats if we needs, such as arbitrary - # unicode for column labels. - value_columns = (expr.get_column(column_name) for column_name in columns) - expr = expr.projection(value_columns) + else: + expr = expr.drop_columns(self._block.index_columns) # Make columns in SQL reflect _labels_ not _ids_. Note: This may use # the arbitrary unicode column labels feature in BigQuery, which is # currently (June 2023) in preview. # TODO(swast): Handle duplicate and NULL labels. - ibis_expr = expr.to_ibis_expr() - renamed_columns = [ - ibis_expr[col_id].name(col_label) + id_overrides = { + col_id: col_label for col_id, col_label in zip(columns, column_labels) - ] - ibis_expr = ibis_expr.select(*renamed_columns) - sql = session.ibis_client.compile(ibis_expr) # type: ignore + if col_label + } + + if ordering_id is not None: + ibis_expr = expr.to_ibis_expr( + ordering_mode="offset_col", + col_id_overrides=id_overrides, + order_col_name=ordering_id, + ) + else: + ibis_expr = expr.to_ibis_expr( + ordering_mode="unordered", + col_id_overrides=id_overrides, + ) + + return session.ibis_client.compile(ibis_expr) # type: ignore + + def _run_io_query( + self, + index: bool, + ordering_id: Optional[str] = None, + job_config: Optional[bigquery.job.QueryJobConfig] = None, + ) -> bigquery.TableReference: + """Executes a query job presenting this dataframe and returns the destination + table.""" + expr = self._block.expr + session = expr._session + sql = self._create_io_query(index=index, ordering_id=ordering_id) _, query_job = session._start_query( sql=sql, job_config=job_config # type: ignore ) - return query_job.destination + self._set_internal_query_job(query_job) + + # The query job should have finished, so there should be always be a result table. + result_table = query_job.destination + assert result_table is not None + return result_table def map(self, func, na_action: Optional[str] = None) -> DataFrame: if not callable(func): @@ -1560,18 +1750,29 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: raise ValueError(f"na_action={na_action} not supported") # TODO(shobs): Support **kwargs - - return self._apply_to_rows( + # Reproject as workaround to applying filter too late. This forces the filter + # to be applied before passing data to remote function, protecting from bad + # inputs causing errors. + reprojected_df = DataFrame(self._block._force_reproject()) + return reprojected_df._apply_to_rows( ops.RemoteFunctionOp(func, apply_on_null=(na_action is None)) ) - def drop_duplicates(self, subset=None, *, keep: str = "first") -> DataFrame: + def drop_duplicates( + self, + subset: typing.Union[blocks.Label, typing.Sequence[blocks.Label]] = None, + *, + keep: str = "first", + ) -> DataFrame: if subset is None: column_ids = self._block.value_columns - else: + elif _is_list_like(subset): column_ids = [ id for label in subset for id in self._block.label_to_col_id[label] ] + else: + # interpret as single label + column_ids = self._block.label_to_col_id[typing.cast(blocks.Label, subset)] block = block_ops.drop_duplicates(self._block, column_ids, keep) return DataFrame(block) @@ -1619,8 +1820,8 @@ def _get_block(self) -> blocks.Block: def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: - return pd.api.types.is_list_like(obj) + return pandas.api.types.is_list_like(obj) def _is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]: - return pd.api.types.is_dict_like(obj) + return pandas.api.types.is_dict_like(obj) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index c257d6943eb..95cf737b2e7 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -14,6 +14,7 @@ """Mappings for Pandas dtypes supported by BigQuery DataFrames package""" +import textwrap import typing from typing import Any, Dict, Iterable, Literal, Tuple, Union @@ -25,6 +26,8 @@ import pandas as pd import pyarrow as pa +import bigframes.constants as constants + # Type hints for Pandas dtypes supported by BigQuery DataFrame Dtype = Union[ pd.BooleanDtype, @@ -37,6 +40,9 @@ # Corresponds to the pandas concept of numeric type (such as when 'numeric_only' is specified in an operation) NUMERIC_BIGFRAMES_TYPES = [pd.BooleanDtype(), pd.Float64Dtype(), pd.Int64Dtype()] +# On BQ side, ARRAY, STRUCT, GEOGRAPHY, JSON are not orderable +UNORDERED_DTYPES = [gpd.array.GeometryDtype()] + # Type hints for dtype strings supported by BigQuery DataFrame DtypeString = Literal[ "boolean", @@ -150,7 +156,9 @@ def ibis_dtype_to_bigframes_dtype( if ibis_dtype in IBIS_TO_BIGFRAMES: return IBIS_TO_BIGFRAMES[ibis_dtype] else: - raise ValueError(f"Unexpected Ibis data type {type(ibis_dtype)}") + raise ValueError( + f"Unexpected Ibis data type {type(ibis_dtype)}. {constants.FEEDBACK_LINK}" + ) def ibis_value_to_canonical_type(value: ibis_types.Value) -> ibis_types.Value: @@ -182,14 +190,14 @@ def bigframes_dtype_to_ibis_dtype( """Converts a BigQuery DataFrames supported dtype to an Ibis dtype. Args: - bigframes_dtype: A dtype supported by BigQuery DataFrame + bigframes_dtype: + A dtype supported by BigQuery DataFrame Returns: - The corresponding Ibis type + IbisDtype: The corresponding Ibis type Raises: - ValueError: - If passed a dtype not supported by BigQuery DataFrames. + ValueError: If passed a dtype not supported by BigQuery DataFrames. """ type_string = str(bigframes_dtype) if type_string in BIGFRAMES_STRING_TO_BIGFRAMES: @@ -197,7 +205,23 @@ def bigframes_dtype_to_ibis_dtype( typing.cast(DtypeString, type_string) ] else: - raise ValueError(f"Unexpected data type {bigframes_dtype}") + raise ValueError( + textwrap.dedent( + f""" + Unexpected data type {bigframes_dtype}. The following + str dtypes are supppted: 'boolean','Float64','Int64', 'string', + 'tring[pyarrow]','timestamp[us, tz=UTC][pyarrow]', + 'timestamp[us][pyarrow]','date32[day][pyarrow]', + 'time64[us][pyarrow]'. The following pandas.ExtensionDtype are + supported: pandas.BooleanDtype(), pandas.Float64Dtype(), + pandas.Int64Dtype(), pandas.StringDtype(storage="pyarrow"), + pd.ArrowDtype(pa.date32()), pd.ArrowDtype(pa.time64("us")), + pd.ArrowDtype(pa.timestamp("us")), + pd.ArrowDtype(pa.timestamp("us", tz="UTC")). + {constants.FEEDBACK_LINK} + """ + ) + ) return BIGFRAMES_TO_IBIS[bigframes_dtype] @@ -209,8 +233,10 @@ def literal_to_ibis_scalar( expression with a BigQuery DataFrames compatible data type Args: - literal: any value accepted by Ibis - force_dtype: force the value to a specific dtype + literal: + any value accepted by Ibis + force_dtype: + force the value to a specific dtype validate: If true, will raise ValueError if type cannot be stored in a BigQuery DataFrames object. If used as a subexpression, this should @@ -227,7 +253,9 @@ def literal_to_ibis_scalar( if pd.api.types.is_list_like(literal): if validate: - raise ValueError("List types can't be stored in BigQuery DataFrames") + raise ValueError( + f"List types can't be stored in BigQuery DataFrames. {constants.FEEDBACK_LINK}" + ) # "correct" way would be to use ibis.array, but this produces invalid BQ SQL syntax return tuple(literal) if not pd.api.types.is_list_like(literal) and pd.isna(literal): @@ -246,7 +274,9 @@ def literal_to_ibis_scalar( # TODO(bmil): support other literals that can be coerced to compatible types if validate and (scalar_expr.type() not in BIGFRAMES_TO_IBIS.values()): - raise ValueError(f"Literal did not coerce to a supported data type: {literal}") + raise ValueError( + f"Literal did not coerce to a supported data type: {literal}. {constants.FEEDBACK_LINK}" + ) return scalar_expr @@ -255,9 +285,11 @@ def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.V """Perform compatible type casts of ibis values Args: - value: Ibis value, which could be a literal, scalar, or column + value: + Ibis value, which could be a literal, scalar, or column - to_type: The Ibis type to cast to + to_type: + The Ibis type to cast to Returns: A new Ibis value of type to_type @@ -275,7 +307,7 @@ def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.V ibis_dtypes.float64, ibis_dtypes.string, ), - ibis_dtypes.float64: (ibis_dtypes.string,), + ibis_dtypes.float64: (ibis_dtypes.string, ibis_dtypes.int64), ibis_dtypes.string: (), ibis_dtypes.date: (), ibis_dtypes.time: (), @@ -289,7 +321,9 @@ def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.V return value.cast(to_type) else: # this should never happen - raise TypeError(f"Unexpected value type {value.type()}") + raise TypeError( + f"Unexpected value type {value.type()}. {constants.FEEDBACK_LINK}" + ) # casts that need some encouragement @@ -301,4 +335,9 @@ def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.V if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.float64: return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64) - raise TypeError(f"Unsupported cast {value.type()} to {to_type}") + if value.type() == ibis_dtypes.float64 and to_type == ibis_dtypes.bool: + return value != ibis_types.literal(0) + + raise TypeError( + f"Unsupported cast {value.type()} to {to_type}. {constants.FEEDBACK_LINK}" + ) diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py index fa0ed0c64d1..2c3c50ba754 100644 --- a/bigframes/formatting_helpers.py +++ b/bigframes/formatting_helpers.py @@ -13,28 +13,53 @@ # limitations under the License. """Shared helper functions for formatting jobs related info.""" +# TODO(orrbradford): cleanup up typings and documenttion in this file import datetime -from typing import Optional +from typing import Any, Optional, Union +import google.api_core.exceptions as api_core_exceptions import google.cloud.bigquery as bigquery import humanize import IPython import IPython.display as display import ipywidgets as widgets +import bigframes.constants as constants + +GenericJob = Union[bigquery.LoadJob, bigquery.ExtractJob, bigquery.QueryJob] + query_job_prop_pairs = { "Job Id": "job_id", "Destination Table": "destination", "Slot Time": "slot_millis", - "Bytes Processed": "estimated_bytes_processed", + "Bytes Processed": "total_bytes_processed", "Cache hit": "cache_hit", } -def repr_query_job(query_job: Optional[bigquery.QueryJob]): +def add_feedback_link( + exception: Union[ + api_core_exceptions.RetryError, api_core_exceptions.GoogleAPICallError + ] +): + exception.message = exception.message + f" {constants.FEEDBACK_LINK}" + + +def repr_query_job_html(query_job: Optional[bigquery.QueryJob]): + """Return query job in html format. + Args: + query_job: + The job representing the execution of the query on the server. + Returns: + Pywidget html table. + """ if query_job is None: return widgets.HTML("No job information available") + if query_job.dry_run: + return widgets.HTML( + f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}" + ) table_html = "" table_html += "" for key, value in query_job_prop_pairs.items(): @@ -54,7 +79,35 @@ def repr_query_job(query_job: Optional[bigquery.QueryJob]): return widgets.HTML(table_html) -def wait_for_job( +def repr_query_job(query_job: Optional[bigquery.QueryJob]): + """Return query job as a formatted string. + Args: + query_job: + The job representing the execution of the query on the server. + Returns: + Pywidget html table. + """ + if query_job is None: + return "No job information available" + if query_job.dry_run: + return f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}" + res = "Query Job Info" + for key, value in query_job_prop_pairs.items(): + job_val = getattr(query_job, value) + if job_val is not None: + res += "\n" + if key == "Job Id": # add link to job + res += f"""Job url: {get_job_url(query_job)}""" + elif key == "Slot Time": + res += f"""{key}: {get_formatted_time(job_val)}""" + elif key == "Bytes Processed": + res += f"""{key}: {get_formatted_bytes(job_val)}""" + else: + res += f"""{key}: {job_val}""" + return res + + +def wait_for_query_job( query_job: bigquery.QueryJob, max_results: Optional[int] = None, progress_bar: Optional[str] = None, @@ -72,19 +125,75 @@ def wait_for_job( if progress_bar == "auto": progress_bar = "notebook" if in_ipython() else "terminal" - if progress_bar == "notebook": - display.display(loading_bar) - query_result = query_job.result(max_results=max_results) - query_job.reload() - loading_bar.close() - elif progress_bar == "terminal": - print(get_query_job_loading_string(query_job)) - query_result = query_job.result(max_results=max_results) - query_job.reload() - return query_result + try: + if progress_bar == "notebook": + display.display(loading_bar) + query_result = query_job.result(max_results=max_results) + query_job.reload() + loading_bar.value = get_query_job_loading_html(query_job) + elif progress_bar == "terminal": + initial_loading_bar = get_query_job_loading_string(query_job) + print(initial_loading_bar) + query_result = query_job.result(max_results=max_results) + query_job.reload() + if initial_loading_bar != get_query_job_loading_string(query_job): + print(get_query_job_loading_string(query_job)) + else: + # No progress bar. + query_result = query_job.result(max_results=max_results) + query_job.reload() + return query_result + except api_core_exceptions.RetryError as exc: + add_feedback_link(exc) + raise + except api_core_exceptions.GoogleAPICallError as exc: + add_feedback_link(exc) + raise + + +def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): + """Waits for job results. Displays a progress bar while the job is running + Args: + job: + The bigquery job to be executed + """ + loading_bar = widgets.HTML(get_base_job_loading_html(job)) + if progress_bar == "auto": + progress_bar = "notebook" if in_ipython() else "terminal" + + try: + if progress_bar == "notebook": + display.display(loading_bar) + job.result() + job.reload() + loading_bar.value = get_base_job_loading_html(job) + elif progress_bar == "terminal": + inital_loading_bar = get_base_job_loading_string(job) + print(inital_loading_bar) + job.result() + job.reload() + if get_base_job_loading_string != inital_loading_bar: + print(get_base_job_loading_string(job)) + else: + # No progress bar. + job.result() + job.reload() + except api_core_exceptions.RetryError as exc: + add_feedback_link(exc) + raise + except api_core_exceptions.GoogleAPICallError as exc: + add_feedback_link(exc) + raise -def get_job_url(query_job: bigquery.QueryJob): +def get_job_url(query_job: GenericJob): + """Return url to the query job in cloud console. + Args: + query_job: + The job representing the execution of the query on the server. + Returns: + String url. + """ if ( query_job.project is None or query_job.location is None @@ -95,16 +204,57 @@ def get_job_url(query_job: bigquery.QueryJob): def get_query_job_loading_html(query_job: bigquery.QueryJob): - return f"""Job {query_job.job_id} is {query_job.state}. Open Job""" + """Return progress bar html string + Args: + query_job: + The job representing the execution of the query on the server. + Returns: + Html string. + """ + return f"""Query job {query_job.job_id} is {query_job.state}. {get_bytes_processed_string(query_job.total_bytes_processed)}Open Job""" def get_query_job_loading_string(query_job: bigquery.QueryJob): - return ( - f"""Job {query_job.job_id} is {query_job.state}. \n{get_job_url(query_job)}""" - ) + """Return progress bar string + Args: + query_job: + The job representing the execution of the query on the server. + Returns: + String + """ + return f"""Query job {query_job.job_id} is {query_job.state}.{get_bytes_processed_string(query_job.total_bytes_processed)} \n{get_job_url(query_job)}""" + + +def get_base_job_loading_html(job: GenericJob): + """Return progress bar html string + Args: + job: + The job representing the execution of the query on the server. + Returns: + Html string. + """ + return f"""{job.job_type.capitalize()} job {job.job_id} is {job.state}. Open Job""" + + +def get_base_job_loading_string(job: GenericJob): + """Return progress bar string + Args: + job: + The job representing the execution of the query on the server. + Returns: + String + """ + return f"""{job.job_type.capitalize()} job {job.job_id} is {job.state}. \n{get_job_url(job)}""" def get_formatted_time(val): + """Try to format time + Args: + val: + Time in ms + Returns: + Duration string + """ try: return humanize.naturaldelta(datetime.timedelta(milliseconds=float(val))) except Exception: @@ -112,7 +262,24 @@ def get_formatted_time(val): def get_formatted_bytes(val): - return humanize.naturalsize(val) + """Try to format bytes + Args: + val (Any): + Bytes to format + Returns: + Duration string + """ + if isinstance(val, int): + return humanize.naturalsize(val) + return "N/A" + + +def get_bytes_processed_string(val: Any): + """Try to get bytes processed string. Return empty if passed non int value""" + bytes_processed_string = "" + if isinstance(val, int): + bytes_processed_string = f"""{get_formatted_bytes(val)} processed. """ + return bytes_processed_string def in_ipython(): diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py index e0a1027dfea..e4c68eb17c2 100644 --- a/bigframes/ml/base.py +++ b/bigframes/ml/base.py @@ -22,9 +22,9 @@ """ import abc -from typing import Optional, TypeVar +from typing import cast, Optional, TypeVar -from bigframes.ml.core import BqmlModel +from bigframes.ml import core import third_party.bigframes_vendored.sklearn.base @@ -92,7 +92,7 @@ class Predictor(BaseEstimator): """A BigQuery DataFrames ML Model base class that can be used to predict outputs.""" def __init__(self): - self._bqml_model: Optional[BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None @abc.abstractmethod def predict(self, X): @@ -102,21 +102,32 @@ def predict(self, X): def register(self: _T, vertex_ai_model_id: Optional[str] = None) -> _T: """Register the model to Vertex AI. + + After register, go to https://pantheon.corp.google.com/vertex-ai/models to manage the model registries. + Refer to https://cloud.google.com/vertex-ai/docs/model-registry/introduction for more options. + Args: - vertex_ai_model_id: optional string id as model id in Vertex. If not set, will by default to 'bigframes_{bq_model_id}'. + vertex_ai_model_id (Optional[str], default None): + optional string id as model id in Vertex. If not set, will by default to 'bigframes_{bq_model_id}'. + Vertex Ai model id will be truncated to 63 characters due to its limitation. Returns: BigQuery DataFrames Model after register. """ if not self._bqml_model: - raise RuntimeError("A model must be trained before register.") + # TODO(garrettwu): find a more elegant way to do this. + try: + self._bqml_model = self._create_bqml_model() # type: ignore + except AttributeError: + raise RuntimeError("A model must be trained before register.") + self._bqml_model = cast(core.BqmlModel, self._bqml_model) self._bqml_model.register(vertex_ai_model_id) return self class TrainablePredictor(Predictor): - """A BigQuery DataFrame ML Model base class that can be used to fit and predict outputs. + """A BigQuery DataFrames ML Model base class that can be used to fit and predict outputs. Also the predictor can be attached to a pipeline with transformers.""" @@ -124,10 +135,9 @@ class TrainablePredictor(Predictor): def fit(self, X, y, transforms): pass - # TODO(b/289280565): enable signatures after updating KMeans and PCA - # @abc.abstractmethod - # def score(self, X, y): - # pass + @abc.abstractmethod + def score(self, X, y): + pass # TODO(b/291812029): move to Predictor after implement in LLM and imported models @abc.abstractmethod diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 2dd95bfc1c9..6f6efb9f376 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -13,35 +13,33 @@ # limitations under the License. """Clustering models. This module is styled after Scikit-Learn's cluster module: -https://scikit-learn.org/stable/modules/clustering.html""" +https://scikit-learn.org/stable/modules/clustering.html.""" from __future__ import annotations -from typing import cast, Dict, List, Optional, TYPE_CHECKING +from typing import cast, Dict, List, Optional, Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.core +import bigframes +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.cluster._kmeans class KMeans( third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans.__doc__ def __init__(self, n_clusters=8): self.n_clusters = n_clusters - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> KMeans: + @classmethod + def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> KMeans: assert model.model_type == "KMEANS" kwargs = {} @@ -51,8 +49,8 @@ def _from_bq(session: bigframes.Session, model: bigquery.Model) -> KMeans: if "numClusters" in last_fitting: kwargs["n_clusters"] = int(last_fitting["numClusters"]) - new_kmeans = KMeans(**kwargs) - new_kmeans._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_kmeans = cls(**kwargs) + new_kmeans._bqml_model = core.BqmlModel(session, model) return new_kmeans @property @@ -62,36 +60,55 @@ def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y=None, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> KMeans: + (X,) = utils.convert_to_dataframe(X) + + self._bqml_model = core.create_bqml_model( train_X=X, transforms=transforms, options=self._bqml_options, ) + return self def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") - return cast( - bigframes.dataframe.DataFrame, self._bqml_model.predict(X)[["CENTROID_ID"]] - ) + (X,) = utils.convert_to_dataframe(X) + + return cast(bpd.DataFrame, self._bqml_model.predict(X)[["CENTROID_ID"]]) def to_gbq(self, model_name: str, replace: bool = False) -> KMeans: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + KMeans: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") new_model = self._bqml_model.copy(model_name, replace) return new_model.session.read_gbq_model(model_name) + + def score( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + (X,) = utils.convert_to_dataframe(X) + + return self._bqml_model.evaluate(X) diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index 09645d4cf8e..49b4899bebd 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -12,33 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Build composite transformers on heterogenous data. This module is styled +"""Build composite transformers on heterogeneous data. This module is styled after Scikit-Learn's compose module: -https://scikit-learn.org/stable/modules/classes.html#module-sklearn.compose""" +https://scikit-learn.org/stable/modules/classes.html#module-sklearn.compose.""" from __future__ import annotations import typing -from typing import List, Optional, Tuple, TYPE_CHECKING, Union +from typing import List, Optional, Tuple, Union -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.compose -import bigframes.ml.core -import bigframes.ml.preprocessing +from bigframes import constants +from bigframes.ml import base, core, preprocessing, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.compose._column_transformer CompilablePreprocessorType = Union[ - bigframes.ml.preprocessing.OneHotEncoder, - bigframes.ml.preprocessing.StandardScaler, + preprocessing.OneHotEncoder, + preprocessing.StandardScaler, ] class ColumnTransformer( third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, - bigframes.ml.base.BaseEstimator, + base.BaseEstimator, ): __doc__ = ( third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer.__doc__ @@ -54,14 +50,17 @@ def __init__( ] ], ): + # TODO: if any(transformers) has fitted raise warning self.transformers = transformers - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None + # call self.transformers_ to check chained transformers + self.transformers_ @property def transformers_( self, ) -> List[Tuple[str, CompilablePreprocessorType, str,]]: - """The collection of transformers as tuples of (name, transformer, column)""" + """The collection of transformers as tuples of (name, transformer, column).""" result: List[ Tuple[ str, @@ -69,13 +68,22 @@ def transformers_( str, ] ] = [] + + column_set: set[str] = set() for entry in self.transformers: name, transformer, column_or_columns = entry - if isinstance(column_or_columns, str): - result.append((name, transformer, column_or_columns)) - else: - for column in column_or_columns: - result.append((name, transformer, column)) + columns = ( + column_or_columns + if isinstance(column_or_columns, List) + else [column_or_columns] + ) + + for column in columns: + if column in column_set: + raise NotImplementedError( + f"Chained transformers on the same column isn't supported. {constants.FEEDBACK_LINK}" + ) + result.append((name, transformer, column)) return result @@ -84,9 +92,11 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: a BQML TRANSFORM clause Args: - columns: a list of column names to transform + columns (List[str]): + a list of column names to transform - Returns: a list of tuples of (sql_expression, output_name)""" + Returns: + a list of tuples of (sql_expression, output_name)""" return [ transformer._compile_to_sql([column])[0] for column in columns @@ -96,12 +106,15 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - ): + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> ColumnTransformer: + (X,) = utils.convert_to_dataframe(X) + compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = bigframes.ml.core.create_bqml_model( + self._bqml_model = core.create_bqml_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, @@ -109,15 +122,16 @@ def fit( # The schema of TRANSFORM output is not available in the model API, so save it during fitting self._output_names = [name for _, name in compiled_transforms] + return self - def transform( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("Must be fitted before transform") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.transform(X) return typing.cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[self._output_names], ) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 2663a8e17e5..9629ca0f4d5 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -21,9 +21,10 @@ from google.cloud import bigquery -import bigframes.dataframe -import bigframes.ml.sql -import bigframes.session +import bigframes +import bigframes.constants as constants +from bigframes.ml import sql as ml_sql +import bigframes.pandas as bpd class BqmlModel: @@ -33,7 +34,7 @@ class BqmlModel: BigQuery DataFrames ML. """ - def __init__(self, session: bigframes.session.Session, model: bigquery.Model): + def __init__(self, session: bigframes.Session, model: bigquery.Model): self._session = session self._model = model @@ -52,29 +53,35 @@ def model(self) -> bigquery.Model: """Get the BQML model associated with this wrapper""" return self._model - @staticmethod + @classmethod def _apply_sql( + cls, session: bigframes.Session, - input_data: bigframes.dataframe.DataFrame, + input_data: bpd.DataFrame, func: Callable[[str], str], - ) -> bigframes.dataframe.DataFrame: + ) -> bpd.DataFrame: """Helper to wrap a dataframe in a SQL query, keeping the index intact. Args: - session: the active bigframes.Session + session (bigframes.Session): + the active bigframes.Session - input_data: the dataframe to be wrapped + input_data (bigframes.dataframe.DataFrame): + the dataframe to be wrapped - func: a function that will accept a SQL string and produce a new SQL + func (function): + a function that will accept a SQL string and produce a new SQL string from which to construct the output dataframe. It must include the index columns of the input SQL. """ - source_sql, tagged_index_cols = input_data.to_sql_query( + source_sql, tagged_index_cols = input_data._to_sql_query( always_include_index=True ) if len(tagged_index_cols) != 1: - raise NotImplementedError("Only exactly one index column is supported") + raise NotImplementedError( + f"Only exactly one index column is supported. {constants.FEEDBACK_LINK}" + ) index_col_name, is_named_index = tagged_index_cols[0] sql = func(source_sql) @@ -84,79 +91,75 @@ def _apply_sql( return df - def predict( - self, input_data: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( self._session, input_data, - lambda source_sql: bigframes.ml.sql.ml_predict( + lambda source_sql: ml_sql.ml_predict( model_name=self.model_name, source_sql=source_sql ), ) - def transform( - self, input_data: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def transform(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( self._session, input_data, - lambda source_sql: bigframes.ml.sql.ml_transform( + lambda source_sql: ml_sql.ml_transform( model_name=self.model_name, source_sql=source_sql ), ) def generate_text( self, - input_data: bigframes.dataframe.DataFrame, + input_data: bpd.DataFrame, options: Mapping[str, int | float], - ) -> bigframes.dataframe.DataFrame: + ) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( self._session, input_data, - lambda source_sql: bigframes.ml.sql.ml_generate_text( + lambda source_sql: ml_sql.ml_generate_text( model_name=self.model_name, source_sql=source_sql, - struct_options=bigframes.ml.sql.struct_options(**options), + struct_options=ml_sql.struct_options(**options), ), ) - def embed_text( + def generate_text_embedding( self, - input_data: bigframes.dataframe.DataFrame, + input_data: bpd.DataFrame, options: Mapping[str, int | float], - ) -> bigframes.dataframe.DataFrame: + ) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( self._session, input_data, - lambda source_sql: bigframes.ml.sql.ml_embed_text( + lambda source_sql: ml_sql.ml_generate_text_embedding( model_name=self.model_name, source_sql=source_sql, - struct_options=bigframes.ml.sql.struct_options(**options), + struct_options=ml_sql.struct_options(**options), ), ) - def forecast(self) -> bigframes.dataframe.DataFrame: - sql = bigframes.ml.sql.ml_forecast(self.model_name) + def forecast(self) -> bpd.DataFrame: + sql = ml_sql.ml_forecast(self.model_name) return self._session.read_gbq(sql) - def evaluate(self, input_data: Union[bigframes.dataframe.DataFrame, None] = None): + def evaluate(self, input_data: Optional[bpd.DataFrame] = None): # TODO: validate input data schema # Note: don't need index as evaluate returns a new table source_sql, _ = ( - input_data.to_sql_query(always_include_index=False) + input_data._to_sql_query(always_include_index=False) if (input_data is not None) else (None, None) ) - sql = bigframes.ml.sql.ml_evaluate(self.model_name, source_sql) + sql = ml_sql.ml_evaluate(self.model_name, source_sql) return self._session.read_gbq(sql) - def copy(self, new_model_name, replace=False) -> BqmlModel: + def copy(self, new_model_name: str, replace: bool = False) -> BqmlModel: job_config = bigquery.job.CopyJobConfig() if replace: job_config.write_disposition = "WRITE_TRUNCATE" @@ -173,10 +176,11 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: # vertex id needs to start with letters. https://cloud.google.com/vertex-ai/docs/general/resource-naming vertex_ai_model_id = "bigframes_" + cast(str, self._model.model_id) - options_sql = bigframes.ml.sql.options( - **{"vertex_ai_model_id": vertex_ai_model_id} - ) - sql = bigframes.ml.sql.alter_model(self.model_name, options_sql=options_sql) + # truncate as Vertex ID only accepts 63 characters, easily exceeding the limit for temp models. + # The possibility of conflicts should be low. + vertex_ai_model_id = vertex_ai_model_id[:63] + options_sql = ml_sql.options(**{"vertex_ai_model_id": vertex_ai_model_id}) + sql = ml_sql.alter_model(self.model_name, options_sql=options_sql) # Register the model and wait it to finish self._session.bqclient.query(sql).result() @@ -185,8 +189,8 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: def create_bqml_model( - train_X: bigframes.dataframe.DataFrame, - train_y: Optional[bigframes.dataframe.DataFrame] = None, + train_X: bpd.DataFrame, + train_y: Optional[bpd.DataFrame] = None, transforms: Optional[Iterable[str]] = None, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> BqmlModel: @@ -217,29 +221,22 @@ def create_bqml_model( # for now, drop index to avoid including the index in feature columns input_data = input_data.reset_index(drop=True) - model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" source_sql = input_data.sql - options_sql = bigframes.ml.sql.options(**options) - transform_sql = ( - bigframes.ml.sql.transform(*transforms) if transforms is not None else None - ) - sql = bigframes.ml.sql.create_model( - model_name=model_name, + options_sql = ml_sql.options(**options) + transform_sql = ml_sql.transform(*transforms) if transforms is not None else None + sql = ml_sql.create_model( + model_name=_create_temp_model_name(), source_sql=source_sql, transform_sql=transform_sql, options_sql=options_sql, ) - # fit the model, synchronously - session.bqclient.query(sql).result() - - model = session.bqclient.get_model(model_name) - return BqmlModel(session, model) + return _create_bqml_model_with_sql(session=session, sql=sql) def create_bqml_time_series_model( - train_X: bigframes.dataframe.DataFrame, - train_y: bigframes.dataframe.DataFrame, + train_X: bpd.DataFrame, + train_y: bpd.DataFrame, transforms: Optional[Iterable[str]] = None, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> BqmlModel: @@ -258,25 +255,18 @@ def create_bqml_time_series_model( # pickpocket session object from the dataframe session = train_X._get_block().expr._session - model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" source_sql = input_data.sql - options_sql = bigframes.ml.sql.options(**options) + options_sql = ml_sql.options(**options) - transform_sql = ( - bigframes.ml.sql.transform(*transforms) if transforms is not None else None - ) - sql = bigframes.ml.sql.create_model( - model_name=model_name, + transform_sql = ml_sql.transform(*transforms) if transforms is not None else None + sql = ml_sql.create_model( + model_name=_create_temp_model_name(), source_sql=source_sql, transform_sql=transform_sql, options_sql=options_sql, ) - # fit the model, synchronously - session.bqclient.query(sql).result() - - model = session.bqclient.get_model(model_name) - return BqmlModel(session, model) + return _create_bqml_model_with_sql(session=session, sql=sql) def create_bqml_remote_model( @@ -287,25 +277,22 @@ def create_bqml_remote_model( """Create a session-temporary BQML remote model with the CREATE MODEL statement Args: - connection_name: a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection - options: a dict of options to configure the model. Generates a BQML OPTIONS - clause + connection_name: + a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection + options: + a dict of options to configure the model. Generates a BQML OPTIONS clause - Returns: a BqmlModel, wrapping a trained model in BigQuery + Returns: + BqmlModel: a BqmlModel wrapping a trained model in BigQuery """ - model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" - options_sql = bigframes.ml.sql.options(**options) - sql = bigframes.ml.sql.create_remote_model( - model_name=model_name, + options_sql = ml_sql.options(**options) + sql = ml_sql.create_remote_model( + model_name=_create_temp_model_name(), connection_name=connection_name, options_sql=options_sql, ) - # create the model, synchronously - session.bqclient.query(sql).result() - - model = session.bqclient.get_model(model_name) - return BqmlModel(session, model) + return _create_bqml_model_with_sql(session=session, sql=sql) def create_bqml_imported_model( @@ -320,15 +307,25 @@ def create_bqml_imported_model( Returns: a BqmlModel, wrapping a trained model in BigQuery """ - model_name = f"{session._session_dataset_id}.{uuid.uuid4().hex}" - options_sql = bigframes.ml.sql.options(**options) - sql = bigframes.ml.sql.create_imported_model( - model_name=model_name, + options_sql = ml_sql.options(**options) + sql = ml_sql.create_imported_model( + model_name=_create_temp_model_name(), options_sql=options_sql, ) - # create the model, synchronously - session.bqclient.query(sql).result() + return _create_bqml_model_with_sql(session=session, sql=sql) + + +def _create_temp_model_name() -> str: + return uuid.uuid4().hex + + +def _create_bqml_model_with_sql(session: bigframes.Session, sql: str) -> BqmlModel: + # fit the model, synchronously + job = session.bqclient.query(sql) + job.result() - model = session.bqclient.get_model(model_name) + # real model path in the session specific hidden dataset and table prefix + model_name_full = f"{job.destination.dataset_id}.{job.destination.table_id}" + model = session.bqclient.get_model(model_name_full) return BqmlModel(session, model) diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 08551d1cb8c..76b4f9ced6a 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -13,34 +13,32 @@ # limitations under the License. """Matrix Decomposition models. This module is styled after Scikit-Learn's decomposition module: -https://scikit-learn.org/stable/modules/decomposition.html""" +https://scikit-learn.org/stable/modules/decomposition.html.""" from __future__ import annotations -from typing import cast, List, Optional, TYPE_CHECKING +from typing import cast, List, Optional, Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.core +import bigframes +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.decomposition._pca class PCA( third_party.bigframes_vendored.sklearn.decomposition._pca.PCA, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.sklearn.decomposition._pca.PCA.__doc__ def __init__(self, n_components=3): self.n_components = n_components - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> PCA: + @classmethod + def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> PCA: assert model.model_type == "PCA" kwargs = {} @@ -50,17 +48,19 @@ def _from_bq(session: bigframes.Session, model: bigquery.Model) -> PCA: if "numPrincipalComponents" in last_fitting: kwargs["n_components"] = int(last_fitting["numPrincipalComponents"]) - new_pca = PCA(**kwargs) - new_pca._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_pca = cls(**kwargs) + new_pca._bqml_model = core.BqmlModel(session, model) return new_pca def fit( self, - X: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], y=None, transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> PCA: + (X,) = utils.convert_to_dataframe(X) + + self._bqml_model = core.create_bqml_model( train_X=X, transforms=transforms, options={ @@ -68,37 +68,45 @@ def fit( "num_principal_components": self.n_components, }, ) + return self - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: - """Predict the closest cluster for each sample in X. - - Args: - X: a BigQuery DataFrame to predict. - y: ignored for API consistency. - - Returns: predicted BigQuery DataFrames.""" + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, self._bqml_model.predict(X)[ ["principal_component_" + str(i + 1) for i in range(self.n_components)] ], ) def to_gbq(self, model_name: str, replace: bool = False) -> PCA: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + PCA: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") new_model = self._bqml_model.copy(model_name, replace) return new_model.session.read_gbq_model(model_name) + + def score( + self, + X=None, + y=None, + ) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before score") + + # TODO(b/291973741): X param is ignored. Update BQML supports input in ML.EVALUTE. + return self._bqml_model.evaluate() diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index 117759ca1c6..56a0cc3d949 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -12,20 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Linear models. This module is styled after Scikit-Learn's linear_model module: -https://scikit-learn.org/stable/modules/linear_model.html""" +"""Ensemble models. This module is styled after Scikit-Learn's ensemble module: +https://scikit-learn.org/stable/modules/ensemble.html""" from __future__ import annotations -from typing import cast, Dict, List, Literal, Optional, TYPE_CHECKING +from typing import cast, Dict, List, Literal, Optional, Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.core +import bigframes +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.ensemble._forest import third_party.bigframes_vendored.xgboost.sklearn @@ -51,7 +49,7 @@ class XGBRegressor( third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor.__doc__ @@ -96,10 +94,12 @@ def __init__( self.min_rel_progress = min_rel_progress self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBRegressor: + @classmethod + def _from_bq( + cls, session: bigframes.Session, model: bigquery.Model + ) -> XGBRegressor: assert model.model_type == "BOOSTED_TREE_REGRESSOR" kwargs = {} @@ -107,14 +107,14 @@ def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBRegressor: # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun last_fitting = model.training_runs[-1]["trainingOptions"] - dummy_regressor = XGBRegressor() + dummy_regressor = cls() for bf_param, bf_value in dummy_regressor.__dict__.items(): bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) if bqml_param is not None: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) - new_xgb_regressor = XGBRegressor(**kwargs) - new_xgb_regressor._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_xgb_regressor = cls(**kwargs) + new_xgb_regressor._bqml_model = core.BqmlModel(session, model) return new_xgb_regressor @property @@ -145,26 +145,31 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> XGBRegressor: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -175,9 +180,11 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], ): + X, y = utils.convert_to_dataframe(X, y) + if not self._bqml_model: raise RuntimeError("A model must be fitted before score") @@ -187,11 +194,13 @@ def score( return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. Returns: saved model.""" if not self._bqml_model: @@ -203,7 +212,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor: class XGBClassifier( third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier.__doc__ @@ -249,10 +258,12 @@ def __init__( self.min_rel_progress = min_rel_progress self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBClassifier: + @classmethod + def _from_bq( + cls, session: bigframes.Session, model: bigquery.Model + ) -> XGBClassifier: assert model.model_type == "BOOSTED_TREE_CLASSIFIER" kwargs = {} @@ -266,8 +277,8 @@ def _from_bq(session: bigframes.Session, model: bigquery.Model) -> XGBClassifier if bqml_param is not None: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) - new_xgb_classifier = XGBClassifier(**kwargs) - new_xgb_classifier._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_xgb_classifier = cls(**kwargs) + new_xgb_classifier._bqml_model = core.BqmlModel(session, model) return new_xgb_classifier @property @@ -298,26 +309,29 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> XGBClassifier: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -328,25 +342,30 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], ): if not self._bqml_model: raise RuntimeError("A model must be fitted before score") + X, y = utils.convert_to_dataframe(X, y) + input_data = ( X.join(y, how="outer") if (X is not None) and (y is not None) else None ) return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + XGBClassifier: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") @@ -356,7 +375,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier: class RandomForestRegressor( third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = ( @@ -396,11 +415,11 @@ def __init__( self.min_rel_progress = min_rel_progress self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod + @classmethod def _from_bq( - session: bigframes.Session, model: bigquery.Model + cls, session: bigframes.Session, model: bigquery.Model ) -> RandomForestRegressor: assert model.model_type == "RANDOM_FOREST_REGRESSOR" @@ -409,16 +428,14 @@ def _from_bq( # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun last_fitting = model.training_runs[-1]["trainingOptions"] - dummy_model = RandomForestRegressor() + dummy_model = cls() for bf_param, bf_value in dummy_model.__dict__.items(): bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) if bqml_param is not None: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) - new_random_forest_regressor = RandomForestRegressor(**kwargs) - new_random_forest_regressor._bqml_model = bigframes.ml.core.BqmlModel( - session, model - ) + new_random_forest_regressor = cls(**kwargs) + new_random_forest_regressor._bqml_model = core.BqmlModel(session, model) return new_random_forest_regressor @property @@ -446,26 +463,32 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> RandomForestRegressor: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -476,36 +499,41 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], ): """Calculate evaluation metrics of the model. Args: - X: a BigFrames DataFrame as evaluation data. - y: a BigFrames DataFrame as evaluation labels. - - Returns: a BigFrames DataFrame as evaluation result.""" + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame as evaluation data. + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame as evaluation labels. + + Returns: + bigframes.dataframe.DataFrame: The DataFrame as evaluation result. + """ if not self._bqml_model: raise RuntimeError("A model must be fitted before score") - if (X is None) != (y is None): - raise ValueError( - "Either both or neither of test_X and test_y must be specified" - ) + X, y = utils.convert_to_dataframe(X, y) + input_data = ( X.join(y, how="outer") if (X is not None) and (y is not None) else None ) return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegressor: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + RandomForestRegressor: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") @@ -515,7 +543,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegresso class RandomForestClassifier( third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = ( @@ -555,11 +583,11 @@ def __init__( self.min_rel_progress = min_rel_progress self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod + @classmethod def _from_bq( - session: bigframes.Session, model: bigquery.Model + cls, session: bigframes.Session, model: bigquery.Model ) -> RandomForestClassifier: assert model.model_type == "RANDOM_FOREST_CLASSIFIER" @@ -574,10 +602,8 @@ def _from_bq( if bqml_param is not None: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) - new_random_forest_classifier = RandomForestClassifier(**kwargs) - new_random_forest_classifier._bqml_model = bigframes.ml.core.BqmlModel( - session, model - ) + new_random_forest_classifier = cls(**kwargs) + new_random_forest_classifier._bqml_model = core.BqmlModel(session, model) return new_random_forest_classifier @property @@ -605,26 +631,32 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> RandomForestClassifier: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -635,36 +667,41 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], ): """Calculate evaluation metrics of the model. Args: - X: a BigFrames DataFrame as evaluation data. - y: a BigFrames DataFrame as evaluation labels. - - Returns: a BigFrames DataFrame as evaluation result.""" + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame as evaluation data. + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame as evaluation labels. + + Returns: + bigframes.dataframe.DataFrame: The DataFrame as evaluation result. + """ if not self._bqml_model: raise RuntimeError("A model must be fitted before score") - if (X is None) != (y is None): - raise ValueError( - "Either both or neither of test_X and test_y must be specified" - ) + X, y = utils.convert_to_dataframe(X, y) + input_data = ( X.join(y, how="outer") if (X is not None) and (y is not None) else None ) return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestClassifier: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + RandomForestClassifier: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index 24afaad7f26..b7e0553ecb7 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -12,35 +12,35 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Forcasting models.""" + from __future__ import annotations -from typing import cast, Dict, List, Optional, TYPE_CHECKING +from typing import cast, Dict, List, Optional, Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.core +import bigframes +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd _PREDICT_OUTPUT_COLUMNS = ["forecast_timestamp", "forecast_value"] -class ARIMAPlus(bigframes.ml.base.TrainablePredictor): +class ARIMAPlus(base.TrainablePredictor): """Time Series ARIMA Plus model.""" def __init__(self): - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> ARIMAPlus: + @classmethod + def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> ARIMAPlus: assert model.model_type == "ARIMA_PLUS" kwargs: Dict[str, str | int | bool | float | List[str]] = {} - new_arima_plus = ARIMAPlus(**kwargs) - new_arima_plus._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_arima_plus = cls(**kwargs) + new_arima_plus._bqml_model = core.BqmlModel(session, model) return new_arima_plus @property @@ -50,74 +50,90 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, ): - """Fit the model to training data + """Fit the model to training data. Args: - X: A dataframe of training timestamp. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A dataframe of training timestamp. + + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + Target values for training. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. + + Returns: + ARIMAPlus: Fitted estimator. + """ + X, y = utils.convert_to_dataframe(X, y) - y: Target values for training.""" - self._bqml_model = bigframes.ml.core.create_bqml_time_series_model( + self._bqml_model = core.create_bqml_time_series_model( X, y, transforms=transforms, options=self._bqml_options, ) - def predict(self, X=None) -> bigframes.dataframe.DataFrame: + def predict(self, X=None) -> bpd.DataFrame: """Predict the closest cluster for each sample in X. Args: - X: ignored, to be compatible with other APIs. + X (default None): + ignored, to be compatible with other APIs. + Returns: - The predicted BigQuery DataFrames. Which contains 2 columns - "forecast_timestamp" and "forecast_value". + bigframes.dataframe.DataFrame: The predicted DataFrames. Which + contains 2 columns "forecast_timestamp" and "forecast_value". """ if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, self._bqml_model.forecast()[_PREDICT_OUTPUT_COLUMNS], ) - # Unlike regression models, time series forcasting can only evaluate with unseen data. X and y must be providee. def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, - ) -> bigframes.dataframe.DataFrame: + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: """Calculate evaluation metrics of the model. Args: - X: - A BigQuery DataFrames only contains 1 column as + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame only contains 1 column as evaluation timestamp. The timestamp must be within the horizon of the model, which by default is 1000 data points. - y: - A BigQuery DataFrames only contains 1 column as + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + A BigQuery DataFrame only contains 1 column as evaluation numeric values. Returns: - A BigQuery DataFrames as evaluation result. + bigframes.dataframe.DataFrame: A DataFrame as evaluation result. """ if not self._bqml_model: raise RuntimeError("A model must be fitted before score") + X, y = utils.convert_to_dataframe(X, y) input_data = X.join(y, how="outer") return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> ARIMAPlus: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + ARIMAPlus: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py index 581ee2b1e24..89078f8267f 100644 --- a/bigframes/ml/imported.py +++ b/bigframes/ml/imported.py @@ -12,47 +12,68 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Imported models.""" + from __future__ import annotations -from typing import cast, TYPE_CHECKING +from typing import cast, Optional, Union -if TYPE_CHECKING: - import bigframes +from google.cloud import bigquery -import bigframes.ml.base -import bigframes.ml.core +import bigframes +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd -class TensorFlowModel(bigframes.ml.base.Predictor): +class TensorFlowModel(base.Predictor): """Imported TensorFlow model. Args: - session: BQ session to create the model - model_path: GCS path that holds the model files.""" + session (BigQuery Session): + BQ session to create the model + model_path (str): + GCS path that holds the model files.""" - def __init__(self, session: bigframes.Session, model_path: str): + def __init__(self, session: bigframes.Session, model_path: Optional[str] = None): self.session = session self.model_path = model_path - self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + self._bqml_model: Optional[core.BqmlModel] = None def _create_bqml_model(self): options = {"model_type": "TENSORFLOW", "model_path": self.model_path} - return bigframes.ml.core.create_bqml_imported_model( - session=self.session, options=options - ) + return core.create_bqml_imported_model(session=self.session, options=options) + + @classmethod + def _from_bq( + cls, session: bigframes.Session, model: bigquery.Model + ) -> TensorFlowModel: + assert model.model_type == "TENSORFLOW" - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + tf_model = cls(session=session, model_path=None) + tf_model._bqml_model = core.BqmlModel(session, model) + return tf_model + + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: """Predict the result from input DataFrame. Args: - X: Input DataFrame, schema is defined by the model. + X (bigframes.dataframe.DataFrame): + Input DataFrame, schema is defined by the model. + + Returns: + bigframes.dataframe.DataFrame: Output DataFrame, schema is defined by the model.""" + + if not self._bqml_model: + if self.model_path is None: + raise ValueError("Model GCS path must be provided.") + self._bqml_model = self._create_bqml_model() + self._bqml_model = cast(core.BqmlModel, self._bqml_model) + + (X,) = utils.convert_to_dataframe(X) - Returns: Output DataFrame, schema is defined by the model.""" df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -61,37 +82,74 @@ def predict( ], ) + def to_gbq(self, model_name: str, replace: bool = False) -> TensorFlowModel: + """Save the model to BigQuery. + + Args: + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. + + Returns: + TensorFlowModel: saved model.""" + if not self._bqml_model: + if self.model_path is None: + raise ValueError("Model GCS path must be provided.") + self._bqml_model = self._create_bqml_model() + self._bqml_model = cast(core.BqmlModel, self._bqml_model) + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) -class OnnxModel(bigframes.ml.base.BaseEstimator): + +class ONNXModel(base.Predictor): """Imported Open Neural Network Exchange (ONNX) model. Args: - session: BQ session to create the model - model_path: GCS path that holds the model files.""" + session (BigQuery Session): + BQ session to create the model + model_path (str): + Cloud Storage path that holds the model files.""" - def __init__(self, session: bigframes.Session, model_path: str): + def __init__(self, session: bigframes.Session, model_path: Optional[str] = None): self.session = session self.model_path = model_path - self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + self._bqml_model: Optional[core.BqmlModel] = None def _create_bqml_model(self): options = {"model_type": "ONNX", "model_path": self.model_path} - return bigframes.ml.core.create_bqml_imported_model( - session=self.session, options=options - ) + return core.create_bqml_imported_model(session=self.session, options=options) + + @classmethod + def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> ONNXModel: + assert model.model_type == "ONNX" - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + onnx_model = cls(session=session, model_path=None) + onnx_model._bqml_model = core.BqmlModel(session, model) + return onnx_model + + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: """Predict the result from input DataFrame. Args: - X: Input DataFrame, schema is defined by the model. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Input DataFrame or Series, schema is defined by the model. + + Returns: + bigframes.dataframe.DataFrame: Output DataFrame, schema is defined by the model.""" + + if not self._bqml_model: + if self.model_path is None: + raise ValueError("Model GCS path must be provided.") + self._bqml_model = self._create_bqml_model() + self._bqml_model = cast(core.BqmlModel, self._bqml_model) + + (X,) = utils.convert_to_dataframe(X) - Returns: Output DataFrame, schema is defined by the model.""" df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -99,3 +157,23 @@ def predict( ] ], ) + + def to_gbq(self, model_name: str, replace: bool = False) -> ONNXModel: + """Save the model to BigQuery. + + Args: + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. + + Returns: + ONNXModel: saved model.""" + if not self._bqml_model: + if self.model_path is None: + raise ValueError("Model GCS path must be provided.") + self._bqml_model = self._create_bqml_model() + self._bqml_model = cast(core.BqmlModel, self._bqml_model) + + new_model = self._bqml_model.copy(model_name, replace) + return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index 2b2ea5c2afc..0b18db93154 100644 --- a/bigframes/ml/linear_model.py +++ b/bigframes/ml/linear_model.py @@ -12,27 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Linear models. This module is styled after Scikit-Learn's linear_model module: -https://scikit-learn.org/stable/modules/linear_model.html""" +"""Linear models. This module is styled after scikit-learn's linear_model module: +https://scikit-learn.org/stable/modules/linear_model.html.""" from __future__ import annotations -from typing import cast, Dict, List, Optional, TYPE_CHECKING +from typing import cast, Dict, List, Optional, Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes - -import bigframes.ml.base -import bigframes.ml.core +import bigframes +import bigframes.constants as constants +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.linear_model._base import third_party.bigframes_vendored.sklearn.linear_model._logistic class LinearRegression( third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = ( third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression.__doc__ @@ -43,10 +42,12 @@ def __init__( fit_intercept=True, ): self.fit_intercept = fit_intercept - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod - def _from_bq(session: bigframes.Session, model: bigquery.Model) -> LinearRegression: + @classmethod + def _from_bq( + cls, session: bigframes.Session, model: bigquery.Model + ) -> LinearRegression: assert model.model_type == "LINEAR_REGRESSION" # TODO(bmil): construct a standard way to extract these properties @@ -57,8 +58,8 @@ def _from_bq(session: bigframes.Session, model: bigquery.Model) -> LinearRegress if "fitIntercept" in last_fitting: kwargs["fit_intercept"] = last_fitting["fitIntercept"] - new_linear_regression = LinearRegression(**kwargs) - new_linear_regression._bqml_model = bigframes.ml.core.BqmlModel(session, model) + new_linear_regression = cls(**kwargs) + new_linear_regression._bqml_model = core.BqmlModel(session, model) return new_linear_regression @property @@ -72,26 +73,29 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> LinearRegression: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -102,25 +106,28 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, - ) -> bigframes.dataframe.DataFrame: + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before score") - input_data = ( - X.join(y, how="outer") if (X is not None) and (y is not None) else None - ) + X, y = utils.convert_to_dataframe(X, y) + + input_data = X.join(y, how="outer") return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + LinearRegression: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") @@ -130,7 +137,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression: class LogisticRegression( third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression, - bigframes.ml.base.TrainablePredictor, + base.TrainablePredictor, ): __doc__ = ( third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression.__doc__ @@ -144,11 +151,11 @@ def __init__( ): self.fit_intercept = fit_intercept self.auto_class_weights = auto_class_weights - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None - @staticmethod + @classmethod def _from_bq( - session: bigframes.Session, model: bigquery.Model + cls, session: bigframes.Session, model: bigquery.Model ) -> LogisticRegression: assert model.model_type == "LOGISTIC_REGRESSION" @@ -166,10 +173,8 @@ def _from_bq( # if "labelClassWeights" in last_fitting: # kwargs["class_weights"] = last_fitting["labelClassWeights"] - new_logistic_regression = LogisticRegression(**kwargs) - new_logistic_regression._bqml_model = bigframes.ml.core.BqmlModel( - session, model - ) + new_logistic_regression = cls(**kwargs) + new_logistic_regression._bqml_model = core.BqmlModel(session, model) return new_logistic_regression @property @@ -186,26 +191,32 @@ def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: def fit( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, - ): - self._bqml_model = bigframes.ml.core.create_bqml_model( + ) -> LogisticRegression: + X, y = utils.convert_to_dataframe(X, y) + + self._bqml_model = core.create_bqml_model( X, y, transforms=transforms, options=self._bqml_options, ) + return self def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.predict(X) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[ [ cast(str, field.name) @@ -216,32 +227,37 @@ def predict( def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, - ) -> bigframes.dataframe.DataFrame: + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before score") - input_data = ( - X.join(y, how="outer") if (X is not None) and (y is not None) else None - ) + X, y = utils.convert_to_dataframe(X, y) + + input_data = X.join(y, how="outer") return self._bqml_model.evaluate(input_data) def to_gbq(self, model_name: str, replace: bool = False) -> LogisticRegression: - """Save the model to Google Cloud BigQuey. + """Save the model to BigQuery. Args: - model_name: the name of the model. - replace: whether to replace if the model already exists. Default to False. + model_name (str): + the name of the model. + replace (bool, default False): + whether to replace if the model already exists. Default to False. - Returns: saved model.""" + Returns: + LogisticRegression: saved model.""" if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") # TODO(ashleyxu): b/285162045 support auto_class_weights once the API is # fixed and enable the tests. if self.auto_class_weights is True: - raise NotImplementedError("auto_class_weight is not supported yet.") + raise NotImplementedError( + f"auto_class_weight is not supported yet. {constants.FEEDBACK_LINK}" + ) new_model = self._bqml_model.copy(model_name, replace) return new_model.session.read_gbq_model(model_name) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index ef4f28f1f7d..def97b56ff0 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -12,80 +12,92 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""LLM models.""" + from __future__ import annotations -from typing import cast +from typing import cast, Union import bigframes +import bigframes.constants as constants from bigframes.core import blocks -import bigframes.ml.base -import bigframes.ml.core +from bigframes.ml import base, core, utils +import bigframes.pandas as bpd _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" _TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE = "CLOUD_AI_TEXT_EMBEDDING_MODEL_V1" -_EMBED_TEXT_RESULT_COLUMN = "ml_embed_text_embedding" +_EMBED_TEXT_RESULT_COLUMN = "text_embedding" -class PaLM2TextGenerator(bigframes.ml.base.Predictor): +class PaLM2TextGenerator(base.Predictor): """PaLM2 text generator LLM model. Args: - session: BQ session to create the model - connection_name: connection to connect with remote service. str of the format ..""" + session (BigQuery Session): + BQ session to create the model + connection_name (str): + connection to connect with remote service. str of the format ..""" def __init__(self, session: bigframes.Session, connection_name: str): self.session = session self.connection_name = connection_name - self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): options = { "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, } - return bigframes.ml.core.create_bqml_remote_model( + return core.create_bqml_remote_model( session=self.session, connection_name=self.connection_name, options=options ) def predict( self, - X: bigframes.dataframe.DataFrame, + X: Union[bpd.DataFrame, bpd.Series], temperature: float = 0.0, max_output_tokens: int = 128, top_k: int = 40, top_p: float = 0.95, - ) -> bigframes.dataframe.DataFrame: + ) -> bpd.DataFrame: """Predict the result from input DataFrame. Args: - X: Input DataFrame, which needs to contain a column with name "prompt". Only the column will be used as input. Prompts can include preamble, questions, suggestions, instructions, or examples. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Input DataFrame or Series, which needs to contain a column with name "prompt". Only the column will be used as input. + Prompts can include preamble, questions, suggestions, instructions, or examples. - temperature: The temperature is used for sampling during the response generation, which occurs when topP and topK are applied. + temperature (float, default 0.0): + The temperature is used for sampling during the response generation, which occurs when topP and topK are applied. Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that expect a true or correct response, while higher temperatures can lead to more diverse or unexpected results. A temperature of 0 is deterministic: the highest probability token is always selected. For most use cases, try starting with a temperature of 0.2. Default 0. - max_output_tokens: Maximum number of tokens that can be generated in the response. Specify a lower value for shorter responses and a higher value for longer responses. + max_output_tokens (int, default 128): + Maximum number of tokens that can be generated in the response. Specify a lower value for shorter responses and a higher value for longer responses. A token may be smaller than a word. A token is approximately four characters. 100 tokens correspond to roughly 60-80 words. Default 128. - top_k: Top-k changes how the model selects tokens for output. A top-k of 1 means the selected token is the most probable among all tokens + top_k (int, default 40): + Top-k changes how the model selects tokens for output. A top-k of 1 means the selected token is the most probable among all tokens in the model’s vocabulary (also called greedy decoding), while a top-k of 3 means that the next token is selected from among the 3 most probable tokens (using temperature). For each token selection step, the top K tokens with the highest probabilities are sampled. Then tokens are further filtered based on topP with the final token selected using temperature sampling. Specify a lower value for less random responses and a higher value for more random responses. Default 40. - top_p: Top-p changes how the model selects tokens for output. Tokens are selected from most K (see topK parameter) probable to least until the sum of their probabilities equals the top-p value. + top_p (float, default 0.95):: + Top-p changes how the model selects tokens for output. Tokens are selected from most K (see topK parameter) probable to least until the sum of their probabilities equals the top-p value. For example, if tokens A, B, and C have a probability of 0.3, 0.2, and 0.1 and the top-p value is 0.5, then the model will select either A or B as the next token (using temperature) and not consider C at all. Specify a lower value for less random responses and a higher value for more random responses. Default 0.95. - Returns: Output DataFrame with only 1 column as the output text results.""" + Returns: + bigframes.dataframe.DataFrame: Output DataFrame with only 1 column as the output text results.""" # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models if temperature < 0.0 or temperature > 1.0: @@ -98,8 +110,13 @@ def predict( raise ValueError(f"top_k must be [1, 40], but is {top_k}.") if top_p < 0.0 or top_p > 1.0: raise ValueError(f"top_p must be [0.0, 1.0], but is {top_p}.") + + (X,) = utils.convert_to_dataframe(X) + if len(X.columns) != 1: - raise ValueError("Only support one column as input.") + raise ValueError( + f"Only support one column as input. {constants.FEEDBACK_LINK}" + ) # BQML identified the column by name col_label = cast(blocks.Label, X.columns[0]) @@ -114,45 +131,52 @@ def predict( } df = self._bqml_model.generate_text(X, options) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[[_TEXT_GENERATE_RESULT_COLUMN]], ) -class PaLM2EmbeddingGenerator(bigframes.ml.base.Predictor): - """PaLM2 embedding generator LLM model. +class PaLM2TextEmbeddingGenerator(base.Predictor): + """PaLM2 text embedding generator LLM model. Args: - session: BQ session to create the model - connection_name: connection to connect with remote service. str of the format ..""" + session (BigQuery Session): + BQ session to create the model + connection_name (str): + connection to connect with remote service. str of the format ..""" def __init__(self, session: bigframes.Session, connection_name: str): self.session = session self.connection_name = connection_name - self._bqml_model: bigframes.ml.core.BqmlModel = self._create_bqml_model() + self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): options = { "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, } - return bigframes.ml.core.create_bqml_remote_model( + return core.create_bqml_remote_model( session=self.session, connection_name=self.connection_name, options=options ) - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: """Predict the result from input DataFrame. Args: - X: Input DataFrame, which needs to contain a column with name "content". Only the column will be used as input. Content can include preamble, questions, suggestions, instructions, or examples. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Input DataFrame, which needs to contain a column with name "content". Only the column will be used as input. Content can include preamble, questions, suggestions, instructions, or examples. - Returns: Output DataFrame with only 1 column as the output embedding results.""" + Returns: + bigframes.dataframe.DataFrame: Output DataFrame with only 1 column as the output embedding results + """ # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models + (X,) = utils.convert_to_dataframe(X) + if len(X.columns) != 1: - raise ValueError("Only support one column as input.") + raise ValueError( + f"Only support one column as input. {constants.FEEDBACK_LINK}" + ) # BQML identified the column by name col_label = cast(blocks.Label, X.columns[0]) @@ -161,8 +185,8 @@ def predict( options = { "flatten_json_output": True, } - df = self._bqml_model.embed_text(X, options) + df = self._bqml_model.generate_text_embedding(X, options) return cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[[_EMBED_TEXT_RESULT_COLUMN]], ) diff --git a/bigframes/ml/loader.py b/bigframes/ml/loader.py index cc43b47698e..805747c49ba 100644 --- a/bigframes/ml/loader.py +++ b/bigframes/ml/loader.py @@ -14,61 +14,81 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Union +from types import MappingProxyType +from typing import Union from google.cloud import bigquery -if TYPE_CHECKING: - import bigframes +import bigframes +import bigframes.constants as constants +from bigframes.ml import ( + cluster, + decomposition, + ensemble, + forecasting, + imported, + linear_model, + pipeline, +) -import bigframes.ml.cluster -import bigframes.ml.decomposition -import bigframes.ml.ensemble -import bigframes.ml.forecasting -import bigframes.ml.linear_model +_BQML_MODEL_TYPE_MAPPING = MappingProxyType( + { + "LINEAR_REGRESSION": linear_model.LinearRegression, + "LOGISTIC_REGRESSION": linear_model.LogisticRegression, + "KMEANS": cluster.KMeans, + "PCA": decomposition.PCA, + "BOOSTED_TREE_REGRESSOR": ensemble.XGBRegressor, + "BOOSTED_TREE_CLASSIFIER": ensemble.XGBClassifier, + "ARIMA_PLUS": forecasting.ARIMAPlus, + "RANDOM_FOREST_REGRESSOR": ensemble.RandomForestRegressor, + "RANDOM_FOREST_CLASSIFIER": ensemble.RandomForestClassifier, + "TENSORFLOW": imported.TensorFlowModel, + "ONNX": imported.ONNXModel, + } +) def from_bq( - session: bigframes.Session, model: bigquery.Model + session: bigframes.Session, bq_model: bigquery.Model ) -> Union[ - bigframes.ml.decomposition.PCA, - bigframes.ml.cluster.KMeans, - bigframes.ml.linear_model.LinearRegression, - bigframes.ml.linear_model.LogisticRegression, - bigframes.ml.ensemble.XGBRegressor, - bigframes.ml.ensemble.XGBClassifier, - bigframes.ml.forecasting.ARIMAPlus, - bigframes.ml.ensemble.RandomForestRegressor, - bigframes.ml.ensemble.RandomForestClassifier, + decomposition.PCA, + cluster.KMeans, + linear_model.LinearRegression, + linear_model.LogisticRegression, + ensemble.XGBRegressor, + ensemble.XGBClassifier, + forecasting.ARIMAPlus, + ensemble.RandomForestRegressor, + ensemble.RandomForestClassifier, + imported.TensorFlowModel, + imported.ONNXModel, + pipeline.Pipeline, ]: """Load a BQML model to BigQuery DataFrames ML. Args: session: a BigQuery DataFrames session. - model: a BigQuery model. + bq_model: a BigQuery model. Returns: A BigQuery DataFrames ML model object. """ - if model.model_type == "LINEAR_REGRESSION": - return bigframes.ml.linear_model.LinearRegression._from_bq(session, model) - elif model.model_type == "KMEANS": - return bigframes.ml.cluster.KMeans._from_bq(session, model) - elif model.model_type == "PCA": - return bigframes.ml.decomposition.PCA._from_bq(session, model) - elif model.model_type == "LOGISTIC_REGRESSION": - return bigframes.ml.linear_model.LogisticRegression._from_bq(session, model) - elif model.model_type == "BOOSTED_TREE_REGRESSOR": - return bigframes.ml.ensemble.XGBRegressor._from_bq(session, model) - elif model.model_type == "BOOSTED_TREE_CLASSIFIER": - return bigframes.ml.ensemble.XGBClassifier._from_bq(session, model) - elif model.model_type == "ARIMA_PLUS": - return bigframes.ml.forecasting.ARIMAPlus._from_bq(session, model) - elif model.model_type == "RANDOM_FOREST_REGRESSOR": - return bigframes.ml.ensemble.RandomForestRegressor._from_bq(session, model) - elif model.model_type == "RANDOM_FOREST_CLASSIFIER": - return bigframes.ml.ensemble.RandomForestClassifier._from_bq(session, model) - else: - raise NotImplementedError( - f"Model type {model.model_type} is not yet supported by BigQuery DataFrames." + if _is_bq_model_pipeline(bq_model): + return pipeline.Pipeline._from_bq(session, bq_model) + + return _model_from_bq(session, bq_model) + + +def _model_from_bq(session: bigframes.Session, bq_model: bigquery.Model): + if bq_model.model_type in _BQML_MODEL_TYPE_MAPPING: + return _BQML_MODEL_TYPE_MAPPING[bq_model.model_type]._from_bq( # type: ignore + session=session, model=bq_model ) + + raise NotImplementedError( + f"Model type {bq_model.model_type} is not yet supported by BigQuery DataFrames. {constants.FEEDBACK_LINK}" + ) + + +def _is_bq_model_pipeline(bq_model: bigquery.Model) -> bool: + return "transformColumns" in bq_model._properties diff --git a/bigframes/ml/metrics.py b/bigframes/ml/metrics.py index 861e1f02d10..3bcb621f745 100644 --- a/bigframes/ml/metrics.py +++ b/bigframes/ml/metrics.py @@ -13,17 +13,18 @@ # limitations under the License. """Metrics functions for evaluating models. This module is styled after -Scikit-Learn's metrics module: https://scikit-learn.org/stable/modules/metrics.html""" +Scikit-Learn's metrics module: https://scikit-learn.org/stable/modules/metrics.html.""" import inspect import typing -from typing import Tuple +from typing import Tuple, Union import numpy as np import pandas as pd import sklearn.metrics as sklearn_metrics # type: ignore -import bigframes.core.blocks as blocks +import bigframes.constants as constants +from bigframes.ml import utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.metrics._classification as vendored_mertics_classification import third_party.bigframes_vendored.sklearn.metrics._ranking as vendored_mertics_ranking @@ -31,33 +32,20 @@ def r2_score( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], force_finite=True, ) -> float: - # TODO(bmil): support multioutput - if len(y_true.columns) > 1 or len(y_pred.columns) > 1: - raise NotImplementedError( - "Only one labels column, one predictions column is supported" - ) - - y_true_series = typing.cast( - bpd.Series, y_true[typing.cast(str, y_true.columns.tolist()[0])] - ) - y_pred_series = typing.cast( - bpd.Series, y_pred[typing.cast(str, y_pred.columns.tolist()[0])] - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) # total sum of squares # (dataframe, scalar) binops - # TODO(bmil): remove multiply by self when bigframes supports pow() # TODO(tbergeron): These stats are eagerly evaluated. Move to lazy representation once scalar subqueries supported. delta_from_mean = y_true_series - y_true_series.mean() ss_total = (delta_from_mean * delta_from_mean).sum() # residual sum of squares # (scalar, scalar) binops - # TODO(bmil): remove multiply by self when bigframes supports pow() delta_from_pred = y_true_series - y_pred_series ss_res = (delta_from_pred * delta_from_pred).sum() @@ -71,22 +59,12 @@ def r2_score( def accuracy_score( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], normalize=True, ) -> float: # TODO(ashleyxu): support sample_weight as the parameter - if len(y_true.columns) != 1 or len(y_pred.columns) != 1: - raise NotImplementedError( - "Only one labels column, one predictions column is supported" - ) - - y_true_series = typing.cast( - bpd.Series, y_true[typing.cast(str, y_true.columns.tolist()[0])] - ) - y_pred_series = typing.cast( - bpd.Series, y_pred[typing.cast(str, y_pred.columns.tolist()[0])] - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) # Compute accuracy for each possible representation # TODO(ashleyxu): add multilabel classification support where y_type @@ -103,73 +81,86 @@ def accuracy_score( def roc_curve( - y_true: bpd.DataFrame, - y_score: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_score: Union[bpd.DataFrame, bpd.Series], drop_intermediate: bool = True, ) -> Tuple[bpd.Series, bpd.Series, bpd.Series]: # TODO(bmil): Add multi-class support # TODO(bmil): Add multi-label support - if len(y_true.columns) > 1 or len(y_score.columns) > 1: - raise NotImplementedError("Only binary classification is supported") # TODO(bmil): Implement drop_intermediate if drop_intermediate: - raise NotImplementedError("drop_intermediate is not yet implemented") + raise NotImplementedError( + f"drop_intermediate is not yet implemented. {constants.FEEDBACK_LINK}" + ) - # TODO(bmil): remove this once bigframes supports the necessary operations - session = y_true._block.expr._session - pd_y_true = y_true.to_pandas() - pd_y_score = y_score.to_pandas() + y_true_series, y_score_series = utils.convert_to_series(y_true, y_score) + + session = y_true_series._block.expr._session # We operate on rows, so, remove the index if there is one # TODO(bmil): check that the indexes are equivalent before removing - pd_y_true = pd_y_true.reset_index(drop=True) - pd_y_score = pd_y_score.reset_index(drop=True) - pd_df = pd.DataFrame( + y_true_series = typing.cast(bpd.Series, y_true_series.reset_index(drop=True)) + y_score_series = typing.cast(bpd.Series, y_score_series.reset_index(drop=True)) + + df = bpd.DataFrame( { - "y_true": pd_y_true[pd_y_true.columns[0]], - "y_score": pd_y_score[pd_y_score.columns[0]], + "y_true": y_true_series, + "y_score": y_score_series, } ) - total_positives = pd_df.y_true.sum() - total_negatives = len(pd_df) - total_positives + total_positives = y_true_series.sum() + total_negatives = y_true_series.count() - total_positives - pd_df = pd_df.sort_values(by="y_score", ascending=False) - pd_df["cum_tp"] = pd_df.y_true.cumsum() - pd_df["cum_fp"] = (~pd_df.y_true.astype(bool)).cumsum() + df = df.sort_values(by="y_score", ascending=False) + df["cum_tp"] = df["y_true"].cumsum() + # have to astype("Int64") as not supported boolean cumsum yet. + df["cum_fp"] = ( + (~typing.cast(bpd.Series, df["y_true"].astype("boolean"))) + .astype("Int64") + .cumsum() + ) # produce just one data point per y_score - pd_df = pd_df.groupby("y_score", as_index=False).last() - pd_df = pd_df.sort_values(by="y_score", ascending=False) + df = df.drop_duplicates(subset="y_score", keep="last") + df = df.sort_values(by="y_score", ascending=False) - pd_df["tpr"] = pd_df.cum_tp / total_positives - pd_df["fpr"] = pd_df.cum_fp / total_negatives - pd_df["thresholds"] = pd_df.y_score + df["tpr"] = typing.cast(bpd.Series, df["cum_tp"]) / total_positives + df["fpr"] = typing.cast(bpd.Series, df["cum_fp"]) / total_negatives + df["thresholds"] = typing.cast(bpd.Series, df["y_score"].astype("Float64")) # sklearn includes an extra datapoint for the origin with threshold np.inf - pd_origin = pd.DataFrame({"tpr": [0.0], "fpr": [0.0], "thresholds": np.inf}) - pd_df = pd.concat([pd_origin, pd_df]) + # having problems with concating inline + df_origin = session.read_pandas( + pd.DataFrame({"tpr": [0.0], "fpr": [0.0], "thresholds": np.inf}) + ) + df = typing.cast(bpd.DataFrame, bpd.concat([df_origin, df], ignore_index=True)) + df = df.reset_index(drop=True) - df = session.read_pandas(pd_df) - return df.fpr, df.tpr, df.thresholds + return ( + typing.cast(bpd.Series, df["fpr"]), + typing.cast(bpd.Series, df["tpr"]), + typing.cast(bpd.Series, df["thresholds"]), + ) roc_curve.__doc__ = inspect.getdoc(vendored_mertics_ranking.roc_curve) -def roc_auc_score(y_true: bpd.DataFrame, y_score: bpd.DataFrame) -> float: +def roc_auc_score( + y_true: Union[bpd.DataFrame, bpd.Series], y_score: Union[bpd.DataFrame, bpd.Series] +) -> float: # TODO(bmil): Add multi-class support # TODO(bmil): Add multi-label support - if len(y_true.columns) > 1 or len(y_score.columns) > 1: - raise NotImplementedError("Only binary classification is supported") + y_true_series, y_score_series = utils.convert_to_series(y_true, y_score) - fpr, tpr, _ = roc_curve(y_true, y_score, drop_intermediate=False) + fpr, tpr, _ = roc_curve(y_true_series, y_score_series, drop_intermediate=False) # TODO(bmil): remove this once bigframes supports the necessary operations - pd_fpr = fpr.compute() - pd_tpr = tpr.compute() + pd_fpr = fpr.to_pandas() + pd_tpr = tpr.to_pandas() # Use the trapezoid rule to compute the area under the ROC curve width_diff = pd_fpr.diff().iloc[1:].reset_index(drop=True) @@ -181,14 +172,13 @@ def roc_auc_score(y_true: bpd.DataFrame, y_score: bpd.DataFrame) -> float: def auc( - x: bpd.DataFrame, - y: bpd.DataFrame, + x: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], ) -> float: - if len(x.columns) != 1 or len(y.columns) != 1: - raise ValueError("Only 1-D data structure is supported") + x_series, y_series = utils.convert_to_series(x, y) # TODO(b/286410053) Support ML exceptions and error handling. - auc = sklearn_metrics.auc(x.to_pandas(), y.to_pandas()) + auc = sklearn_metrics.auc(x_series.to_pandas(), y_series.to_pandas()) return auc @@ -196,30 +186,24 @@ def auc( def confusion_matrix( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], ) -> pd.DataFrame: # TODO(ashleyxu): support labels and sample_weight parameters - # TODO(ashleyxu): support bpd.Series as input type - if len(y_true.columns) != 1 or len(y_pred.columns) != 1: - raise NotImplementedError( - "Only one labels column, one predictions column is supported" - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) - y_true_column = typing.cast(blocks.Label, y_true.columns[0]) - y_pred_series = typing.cast( - bpd.Series, - y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], - ) - confusion_df = y_true.assign(y_pred=y_pred_series) + y_true_series = y_true_series.rename("y_true") + confusion_df = y_true_series.to_frame().assign(y_pred=y_pred_series) confusion_df = confusion_df.assign(dummy=0) groupby_count = ( - confusion_df.groupby(by=[y_true_column, "y_pred"], as_index=False) + confusion_df.groupby(by=["y_true", "y_pred"], as_index=False) .count() .to_pandas() ) - unique_values = sorted(set(groupby_count.y_true).union(set(groupby_count.y_pred))) + unique_values = sorted( + set(groupby_count["y_true"]).union(set(groupby_count["y_pred"])) + ) confusion_matrix = pd.DataFrame( 0, index=pd.Index(unique_values), columns=pd.Index(unique_values), dtype=int @@ -241,28 +225,17 @@ def confusion_matrix( def recall_score( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], average: str = "binary", ) -> pd.Series: # TODO(ashleyxu): support more average type, default to "binary" - # TODO(ashleyxu): support bpd.Series as input type - if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + if average is not None: raise NotImplementedError( - "Only one labels column, one predictions column is supported" + f"Only average=None is supported. {constants.FEEDBACK_LINK}" ) - if average is not None: - raise NotImplementedError("Only average=None is supported") - - y_true_series = typing.cast( - bpd.Series, - y_true[typing.cast(blocks.Label, y_true.columns.tolist()[0])], - ) - y_pred_series = typing.cast( - bpd.Series, - y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) is_accurate = y_true_series == y_pred_series unique_labels = ( @@ -275,7 +248,7 @@ def recall_score( recall = ( is_accurate.groupby(y_true_series).sum() / is_accurate.groupby(y_true_series).count() - ).compute() + ).to_pandas() recall_score = pd.Series(0, index=index) for i in recall_score.index: @@ -288,28 +261,17 @@ def recall_score( def precision_score( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], average: str = "binary", ) -> pd.Series: # TODO(ashleyxu): support more average type, default to "binary" - # TODO(ashleyxu): support bpd.Series as input type - if len(y_true.columns) != 1 or len(y_pred.columns) != 1: + if average is not None: raise NotImplementedError( - "Only one labels column, one predictions column is supported" + f"Only average=None is supported. {constants.FEEDBACK_LINK}" ) - if average is not None: - raise NotImplementedError("Only average=None is supported") - - y_true_series = typing.cast( - bpd.Series, - y_true[typing.cast(blocks.Label, y_true.columns.tolist()[0])], - ) - y_pred_series = typing.cast( - bpd.Series, - y_pred[typing.cast(blocks.Label, y_pred.columns.tolist()[0])], - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) is_accurate = y_true_series == y_pred_series unique_labels = ( @@ -322,7 +284,7 @@ def precision_score( precision = ( is_accurate.groupby(y_pred_series).sum() / is_accurate.groupby(y_pred_series).count() - ).compute() + ).to_pandas() precision_score = pd.Series(0, index=index) for i in precision.index: @@ -337,22 +299,20 @@ def precision_score( def f1_score( - y_true: bpd.DataFrame, - y_pred: bpd.DataFrame, + y_true: Union[bpd.DataFrame, bpd.Series], + y_pred: Union[bpd.DataFrame, bpd.Series], average: str = "binary", ) -> pd.Series: # TODO(ashleyxu): support more average type, default to "binary" - # TODO(ashleyxu): support bpd.Series as input type - if len(y_true.columns) != 1 or len(y_pred.columns) != 1: - raise NotImplementedError( - "Only one labels column, one predictions column is supported" - ) + y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred) if average is not None: - raise NotImplementedError("Only average=None is supported") + raise NotImplementedError( + f"Only average=None is supported. {constants.FEEDBACK_LINK}" + ) - recall = recall_score(y_true, y_pred, average=None) - precision = precision_score(y_true, y_pred, average=None) + recall = recall_score(y_true_series, y_pred_series, average=None) + precision = precision_score(y_true_series, y_pred_series, average=None) f1_score = pd.Series(0, index=recall.index) for index in recall.index: diff --git a/bigframes/ml/model_selection.py b/bigframes/ml/model_selection.py index 73a2fd76a9a..110cbcf4934 100644 --- a/bigframes/ml/model_selection.py +++ b/bigframes/ml/model_selection.py @@ -14,40 +14,40 @@ """Functions for test/train split and model tuning. This module is styled after Scikit-Learn's model_selection module: -https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection""" +https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection.""" from typing import List, Union -import bigframes -import bigframes.dataframe +from bigframes.ml import utils +import bigframes.pandas as bpd def train_test_split( - *dataframes: bigframes.dataframe.DataFrame, + *arrays: Union[bpd.DataFrame, bpd.Series], test_size: Union[float, None] = None, train_size: Union[float, None] = None, random_state: Union[int, None] = None, -) -> List[bigframes.dataframe.DataFrame]: - """Splits dataframes into random train and test subsets +) -> List[Union[bpd.DataFrame, bpd.Series]]: + """Splits dataframes or series into random train and test subsets. Args: - *dataframes: - A sequence of BigQuery DataFrames that can be joined on + *arrays (bigframes.dataframe.DataFrame or bigframes.series.Series): + A sequence of BigQuery DataFrames or Series that can be joined on their indexes - test_size: + test_size (default None): The proportion of the dataset to include in the test split. If None, this will default to the complement of train_size. If both are none, it will be set to 0.25. - train_size: + train_size (default None): The proportion of the dataset to include in the train split. If None, this will default to the complement of test_size. - random_state: + random_state (default None): A seed to use for randomly choosing the rows of the split. If not set, a random split will be generated each time. Returns: - A list of BigQuery DataFrames. + List[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]]: A list of BigQuery DataFrames or Series. """ # TODO(garrettwu): Scikit-Learn throws an error when the dataframes don't have the same @@ -73,14 +73,22 @@ def train_test_split( f"The sum of train_size and test_size exceeds 1.0. train_size: {train_size}. test_size: {test_size}" ) - results = dataframes[0]._split( - fracs=(train_size, test_size), random_state=random_state - ) - train_index = results[0].index - test_index = results[1].index + dfs = list(utils.convert_to_dataframe(*arrays)) - results += [ - df.loc[index] for df in dataframes[1:] for index in (train_index, test_index) + split_dfs = dfs[0]._split(fracs=(train_size, test_size), random_state=random_state) + train_index = split_dfs[0].index + test_index = split_dfs[1].index + + split_dfs += [ + df.loc[index] for df in dfs[1:] for index in (train_index, test_index) ] + # convert back to Series. + results: List[Union[bpd.DataFrame, bpd.Series]] = [] + for i, array in enumerate(arrays): + if isinstance(array, bpd.Series): + results += utils.convert_to_series(split_dfs[2 * i], split_dfs[2 * i + 1]) + else: + results += (split_dfs[2 * i], split_dfs[2 * i + 1]) + return results diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index 022ba0148d7..bfd03925265 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -13,15 +13,19 @@ # limitations under the License. """For composing estimators together. This module is styled after Scikit-Learn's -pipeline module: https://scikit-learn.org/stable/modules/pipeline.html""" +pipeline module: https://scikit-learn.org/stable/modules/pipeline.html.""" from __future__ import annotations -from typing import List, Optional, Tuple +from typing import cast, List, Optional, Tuple, Union + +from google.cloud import bigquery import bigframes -from bigframes.ml import base, cluster, compose, decomposition, preprocessing +import bigframes.constants as constants +from bigframes.ml import base, compose, loader, preprocessing, utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.pipeline @@ -36,7 +40,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): if len(steps) != 2: raise NotImplementedError( - "Currently only two step (transform, estimator) pipelines are supported" + f"Currently only two step (transform, estimator) pipelines are supported. {constants.FEEDBACK_LINK}" ) transform, estimator = steps[0][1], steps[1][1] @@ -51,7 +55,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): self._transform = transform else: raise NotImplementedError( - f"Transform {transform} is not yet supported by Pipeline" + f"Transform {transform} is not yet supported by Pipeline. {constants.FEEDBACK_LINK}" ) if not isinstance( @@ -59,44 +63,134 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): base.TrainablePredictor, ): raise NotImplementedError( - f"Estimator {estimator} is not supported by Pipeline" + f"Estimator {estimator} is not supported by Pipeline. {constants.FEEDBACK_LINK}" ) self._transform = transform self._estimator = estimator + @classmethod + def _from_bq(cls, session: bigframes.Session, bq_model: bigquery.Model) -> Pipeline: + col_transformer = _extract_as_column_transformer(bq_model) + transform = _merge_column_transformer(bq_model, col_transformer) + + estimator = loader._model_from_bq(session, bq_model) + return cls([("transform", transform), ("estimator", estimator)]) + def fit( self, - X: bigframes.dataframe.DataFrame, - y: Optional[bigframes.dataframe.DataFrame] = None, - ): + X: Union[bpd.DataFrame, bpd.Series], + y: Optional[Union[bpd.DataFrame, bpd.Series]] = None, + ) -> Pipeline: + (X,) = utils.convert_to_dataframe(X) + compiled_transforms = self._transform._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] if y is not None: # If labels columns are present, they should pass through un-transformed + (y,) = utils.convert_to_dataframe(y) transform_sqls.extend(y.columns.tolist()) self._estimator.fit(X=X, y=y, transforms=transform_sqls) + return self - def predict( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: return self._estimator.predict(X) def score( self, - X: bigframes.dataframe.DataFrame, - y: bigframes.dataframe.DataFrame, - ): - if isinstance(self._estimator, (cluster.KMeans, decomposition.PCA)): - raise NotImplementedError("KMeans/PCA haven't supported score method.") - - # TODO(b/289280565): remove type ignore after updating KMeans and PCA - return self._estimator.score(X=X, y=y) # type: ignore - - def to_gbq(self, model_name: str, replace: bool = False): - self._estimator.to_gbq(model_name, replace) + X: Union[bpd.DataFrame, bpd.Series], + y: Optional[Union[bpd.DataFrame, bpd.Series]] = None, + ) -> bpd.DataFrame: + (X,) = utils.convert_to_dataframe(X) + if y is not None: + (y,) = utils.convert_to_dataframe(y) + + return self._estimator.score(X=X, y=y) + + def to_gbq(self, model_name: str, replace: bool = False) -> Pipeline: + """Save the pipeline to BigQuery. + + Args: + model_name (str): + the name of the model(pipeline). + replace (bool, default False): + whether to replace if the model(pipeline) already exists. Default to False. + + Returns: + Pipeline: saved model(pipeline).""" + if not self._estimator._bqml_model: + raise RuntimeError("A model must be fitted before it can be saved") + + new_model = self._estimator._bqml_model.copy(model_name, replace) + + return new_model.session.read_gbq_model(model_name) + + +def _extract_as_column_transformer( + bq_model: bigquery.Model, +) -> compose.ColumnTransformer: + """Extract transformers as ColumnTransformer obj from a BQ Model.""" + assert "transformColumns" in bq_model._properties + + transformers: List[ + Tuple[ + str, + Union[preprocessing.OneHotEncoder, preprocessing.StandardScaler], + Union[str, List[str]], + ] + ] = [] + for transform_col in bq_model._properties["transformColumns"]: + # pass the columns that are not transformed + if "transformSql" not in transform_col: + continue + + transform_sql: str = cast(dict, transform_col)["transformSql"] + if transform_sql.startswith("ML.STANDARD_SCALER"): + transformers.append( + ( + "standard_scaler", + *preprocessing.StandardScaler._parse_from_sql(transform_sql), + ) + ) + elif transform_sql.startswith("ML.ONE_HOT_ENCODER"): + transformers.append( + ( + "ont_hot_encoder", + *preprocessing.OneHotEncoder._parse_from_sql(transform_sql), + ) + ) + else: + raise NotImplementedError( + f"Unsupported transformer type. {constants.FEEDBACK_LINK}" + ) - # TODO: should instead load from GBQ, but loading pipelines is not implemented yet - return self + return compose.ColumnTransformer(transformers=transformers) + + +def _merge_column_transformer( + bq_model: bigquery.Model, column_transformer: compose.ColumnTransformer +) -> Union[ + compose.ColumnTransformer, + preprocessing.StandardScaler, + preprocessing.OneHotEncoder, +]: + """Try to merge the column transformer to a simple transformer.""" + transformers = column_transformer.transformers_ + + assert len(transformers) > 0 + _, transformer_0, column_0 = transformers[0] + columns = [column_0] + for _, transformer, column in transformers[1:]: + # all transformers are the same + if transformer != transformer_0: + return column_transformer + columns.append(column) + # all feature columns are transformed + if sorted( + [cast(str, feature_column.name) for feature_column in bq_model.feature_columns] + ) == sorted(columns): + return transformer_0 + + return column_transformer diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 59d2cc2ae9e..500a9fcb244 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -13,29 +13,35 @@ # limitations under the License. """Transformers that prepare data for other estimators. This module is styled after -Scikit-Learn's preprocessing module: https://scikit-learn.org/stable/modules/preprocessing.html""" +Scikit-Learn's preprocessing module: https://scikit-learn.org/stable/modules/preprocessing.html.""" +from __future__ import annotations import typing -from typing import List, Optional, Tuple +from typing import Any, cast, List, Literal, Optional, Tuple, Union -import bigframes -import bigframes.ml -import bigframes.ml.sql +from bigframes.ml import base, core +from bigframes.ml import sql as ml_sql +from bigframes.ml import utils +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.preprocessing._data import third_party.bigframes_vendored.sklearn.preprocessing._encoder class StandardScaler( third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, - bigframes.ml.base.BaseEstimator, + base.BaseEstimator, ): __doc__ = ( third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler.__doc__ ) def __init__(self): - self._bqml_model: Optional[bigframes.ml.core.BqmlModel] = None + self._bqml_model: Optional[core.BqmlModel] = None + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return type(other) is StandardScaler and self._bqml_model == other._bqml_model def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in @@ -47,20 +53,35 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: Returns: a list of tuples of (sql_expression, output_name)""" return [ ( - bigframes.ml.sql.ml_standard_scaler(column, f"scaled_{column}"), + ml_sql.ml_standard_scaler(column, f"scaled_{column}"), f"scaled_{column}", ) for column in columns ] + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[StandardScaler, str]: + """Parse SQL to tuple(StandardScaler, column_label). + + Args: + sql: SQL string of format "ML.STANDARD_SCALER({col_label}) OVER()" + + Returns: + tuple(StandardScaler, column_label)""" + col_label = sql[sql.find("(") + 1 : sql.find(")")] + return cls(), col_label + def fit( self, - X: bigframes.dataframe.DataFrame, - ): + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> StandardScaler: + (X,) = utils.convert_to_dataframe(X) + compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = bigframes.ml.core.create_bqml_model( + self._bqml_model = core.create_bqml_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, @@ -68,57 +89,124 @@ def fit( # The schema of TRANSFORM output is not available in the model API, so save it during fitting self._output_names = [name for _, name in compiled_transforms] + return self - def transform( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("Must be fitted before transform") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.transform(X) return typing.cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[self._output_names], ) class OneHotEncoder( third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, - bigframes.ml.base.BaseEstimator, + base.BaseEstimator, ): + # BQML max value https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder#syntax + TOP_K_DEFAULT = 1000000 + FREQUENCY_THRESHOLD_DEFAULT = 0 + __doc__ = ( third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder.__doc__ ) # All estimators must implement __init__ to document their parameters, even # if they don't have any - def __init__(self): - pass + def __init__( + self, + drop: Optional[Literal["most_frequent"]] = None, + min_frequency: Optional[int] = None, + max_categories: Optional[int] = None, + ): + if max_categories is not None and max_categories < 2: + raise ValueError( + f"max_categories has to be larger than or equal to 2, input is {max_categories}." + ) + self.drop = drop + self.min_frequency = min_frequency + self.max_categories = max_categories + self._bqml_model: Optional[core.BqmlModel] = None + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return ( + type(other) is OneHotEncoder + and self._bqml_model == other._bqml_model + and self.drop == other.drop + and self.min_frequency == other.min_frequency + and self.max_categories == other.max_categories + ) def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: - columns: a list of column names to transform + columns: + a list of column names to transform Returns: a list of tuples of (sql_expression, output_name)""" + + drop = self.drop if self.drop is not None else "none" + # minus one here since BQML's inplimentation always includes index 0, and top_k is on top of that. + top_k = ( + (self.max_categories - 1) + if self.max_categories is not None + else OneHotEncoder.TOP_K_DEFAULT + ) + frequency_threshold = ( + self.min_frequency + if self.min_frequency is not None + else OneHotEncoder.FREQUENCY_THRESHOLD_DEFAULT + ) return [ ( - bigframes.ml.sql.ml_one_hot_encoder(column, f"onehotencoded_{column}"), + ml_sql.ml_one_hot_encoder( + column, drop, top_k, frequency_threshold, f"onehotencoded_{column}" + ), f"onehotencoded_{column}", ) for column in columns ] + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[OneHotEncoder, str]: + """Parse SQL to tuple(OneHotEncoder, column_label). + + Args: + sql: SQL string of format "ML.ONE_HOT_ENCODER({col_label}, '{drop}', {top_k}, {frequency_threshold}) OVER() " + + Returns: + tuple(OneHotEncoder, column_label)""" + s = sql[sql.find("(") + 1 : sql.find(")")] + col_label, drop_str, top_k, frequency_threshold = s.split(", ") + drop = ( + cast(Literal["most_frequent"], "most_frequent") + if drop_str.lower() == "'most_frequent'" + else None + ) + max_categories = int(top_k) + 1 + min_frequency = int(frequency_threshold) + + return cls(drop, min_frequency, max_categories), col_label + def fit( self, - X: bigframes.dataframe.DataFrame, - ): + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> OneHotEncoder: + (X,) = utils.convert_to_dataframe(X) + compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = bigframes.ml.core.create_bqml_model( + self._bqml_model = core.create_bqml_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, @@ -126,15 +214,16 @@ def fit( # The schema of TRANSFORM output is not available in the model API, so save it during fitting self._output_names = [name for _, name in compiled_transforms] + return self - def transform( - self, X: bigframes.dataframe.DataFrame - ) -> bigframes.dataframe.DataFrame: + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("Must be fitted before transform") + (X,) = utils.convert_to_dataframe(X) + df = self._bqml_model.transform(X) return typing.cast( - bigframes.dataframe.DataFrame, + bpd.DataFrame, df[self._output_names], ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 7017b2a4cd1..b8d9e2c6737 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -18,6 +18,8 @@ from typing import Iterable, Optional, Union +import bigframes.constants as constants + def _encode_value(v: Union[str, int, float, Iterable[str]]) -> str: """Encode a parameter value for SQL""" @@ -29,7 +31,7 @@ def _encode_value(v: Union[str, int, float, Iterable[str]]) -> str: inner = ", ".join([_encode_value(x) for x in v]) return f"[{inner}]" else: - raise ValueError("Unexpected value type") + raise ValueError(f"Unexpected value type. {constants.FEEDBACK_LINK}") def _build_param_Iterable(**kwargs: Union[str, int, float, Iterable[str]]) -> str: @@ -77,9 +79,12 @@ def ml_standard_scaler(numeric_expr_sql: str, name: str) -> str: return f"""ML.STANDARD_SCALER({numeric_expr_sql}) OVER() AS {name}""" -def ml_one_hot_encoder(numeric_expr_sql: str, name: str) -> str: - """Encode ML.ONE_HOT_ENCODER for BQML""" - return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}) OVER() AS {name}""" +def ml_one_hot_encoder( + numeric_expr_sql: str, drop: str, top_k: int, frequency_threshold: int, name: str +) -> str: + """Encode ML.ONE_HOT_ENCODER for BQML. + https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder for params.""" + return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}, '{drop}', {top_k}, {frequency_threshold}) OVER() AS {name}""" def create_model( @@ -88,9 +93,8 @@ def create_model( transform_sql: Optional[str] = None, options_sql: Optional[str] = None, ) -> str: - """Encode the CREATE MODEL statement for BQML""" - # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed - parts = [f"CREATE MODEL `{model_name}`"] + """Encode the CREATE TEMP MODEL statement for BQML""" + parts = [f"CREATE TEMP MODEL `{model_name}`"] if transform_sql: parts.append(transform_sql) if options_sql: @@ -104,9 +108,8 @@ def create_remote_model( connection_name: str, options_sql: Optional[str] = None, ) -> str: - """Encode the CREATE MODEL statement for BQML""" - # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed - parts = [f"CREATE MODEL `{model_name}`"] + """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + parts = [f"CREATE TEMP MODEL `{model_name}`"] parts.append(connection(connection_name)) if options_sql: parts.append(options_sql) @@ -117,9 +120,8 @@ def create_imported_model( model_name: str, options_sql: Optional[str] = None, ) -> str: - """Encode the CREATE MODEL statement for BQML""" - # TODO(garrettwu): This should be CREATE TEMP MODEL after b/145824779 is fixed - parts = [f"CREATE MODEL `{model_name}`"] + """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + parts = [f"CREATE TEMP MODEL `{model_name}`"] if options_sql: parts.append(options_sql) return "\n".join(parts) @@ -162,9 +164,11 @@ def ml_generate_text(model_name: str, source_sql: str, struct_options: str) -> s ({source_sql}), {struct_options})""" -def ml_embed_text(model_name: str, source_sql: str, struct_options: str) -> str: - """Encode ML.EMBED_TEXT for BQML""" - return f"""SELECT * FROM ML.EMBED_TEXT(MODEL `{model_name}`, +def ml_generate_text_embedding( + model_name: str, source_sql: str, struct_options: str +) -> str: + """Encode ML.GENERATE_TEXT_EMBEDDING for BQML""" + return f"""SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `{model_name}`, ({source_sql}), {struct_options})""" diff --git a/bigframes/ml/utils.py b/bigframes/ml/utils.py new file mode 100644 index 00000000000..299282d3337 --- /dev/null +++ b/bigframes/ml/utils.py @@ -0,0 +1,58 @@ +# 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 typing +from typing import Iterable, Union + +import bigframes.constants as constants +from bigframes.core import blocks +import bigframes.pandas as bpd + +# Internal type alias +ArrayType = Union[bpd.DataFrame, bpd.Series] + + +def convert_to_dataframe(*input: ArrayType) -> Iterable[bpd.DataFrame]: + return (_convert_to_dataframe(frame) for frame in input) + + +def _convert_to_dataframe(frame: ArrayType) -> bpd.DataFrame: + if isinstance(frame, bpd.DataFrame): + return frame + if isinstance(frame, bpd.Series): + return frame.to_frame() + raise ValueError( + f"Unsupported type {type(frame)} to convert to DataFrame. {constants.FEEDBACK_LINK}" + ) + + +def convert_to_series(*input: ArrayType) -> Iterable[bpd.Series]: + return (_convert_to_series(frame) for frame in input) + + +def _convert_to_series(frame: ArrayType) -> bpd.Series: + if isinstance(frame, bpd.DataFrame): + if len(frame.columns) != 1: + raise ValueError( + "To convert into Series, DataFrames can only contain one column. " + f"Try input with only one column. {constants.FEEDBACK_LINK}" + ) + + label = typing.cast(blocks.Label, frame.columns.tolist()[0]) + return typing.cast(bpd.Series, frame[label]) + if isinstance(frame, bpd.Series): + return frame + raise ValueError( + f"Unsupported type {type(frame)} to convert to Series. {constants.FEEDBACK_LINK}" + ) diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 781cc493397..13063af75f9 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -23,11 +23,14 @@ import ibis.expr.operations.generic import ibis.expr.types as ibis_types import numpy as np +import pandas as pd +import bigframes.constants as constants import bigframes.dtypes import bigframes.dtypes as dtypes _ZERO = typing.cast(ibis_types.NumericValue, ibis_types.literal(0)) +_NAN = typing.cast(ibis_types.NumericValue, ibis_types.literal(np.nan)) _INF = typing.cast(ibis_types.NumericValue, ibis_types.literal(np.inf)) BinaryOp = typing.Callable[[ibis_types.Value, ibis_types.Value], ibis_types.Value] @@ -39,7 +42,9 @@ ### Unary Ops class UnaryOp: def _as_ibis(self, x): - raise NotImplementedError("Base class UnaryOp has no implementation.") + raise NotImplementedError( + f"Base class UnaryOp has no implementation. {constants.FEEDBACK_LINK}" + ) @property def is_windowed(self): @@ -113,6 +118,79 @@ def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.StringValue, x).capitalize() +class ContainsStringOp(UnaryOp): + def __init__(self, pat: str, case: bool = True): + self._pat = pat + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).contains(self._pat) + + +class ContainsRegexOp(UnaryOp): + def __init__(self, pat: str): + self._pat = pat + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_search(self._pat) + + +class ReplaceStringOp(UnaryOp): + def __init__(self, pat: str, repl: str): + self._pat = pat + self._repl = repl + + def _as_ibis(self, x: ibis_types.Value): + pat_str_value = typing.cast( + ibis_types.StringValue, ibis_types.literal(self._pat) + ) + repl_str_value = typing.cast( + ibis_types.StringValue, ibis_types.literal(self._pat) + ) + + return typing.cast(ibis_types.StringValue, x).replace( + pat_str_value, repl_str_value + ) + + +class ReplaceRegexOp(UnaryOp): + def __init__(self, pat: str, repl: str): + self._pat = pat + self._repl = repl + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_replace(self._pat, self._repl) + + +class StartsWithOp(UnaryOp): + def __init__(self, pat: typing.Sequence[str]): + self._pat = pat + + def _as_ibis(self, x: ibis_types.Value): + any_match = None + for pat in self._pat: + pat_match = typing.cast(ibis_types.StringValue, x).startswith(pat) + if any_match is not None: + any_match = any_match | pat_match + else: + any_match = pat_match + return any_match if any_match is not None else ibis_types.literal(False) + + +class EndsWithOp(UnaryOp): + def __init__(self, pat: typing.Sequence[str]): + self._pat = pat + + def _as_ibis(self, x: ibis_types.Value): + any_match = None + for pat in self._pat: + pat_match = typing.cast(ibis_types.StringValue, x).endswith(pat) + if any_match is not None: + any_match = any_match | pat_match + else: + any_match = pat_match + return any_match if any_match is not None else ibis_types.literal(False) + + class HashOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.IntegerValue, x).hash() @@ -192,6 +270,15 @@ def _as_ibis(self, x: ibis_types.Value): ) +class ExtractOp(UnaryOp): + def __init__(self, pat: str, n: int = 1): + self._pat = pat + self._n = n + + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_extract(self._pat, self._n) + + class SliceOp(UnaryOp): def __init__(self, start, stop): self._start = start @@ -201,6 +288,20 @@ def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.StringValue, x)[self._start : self._stop] +class IsInOp(UnaryOp): + def __init__(self, values, match_nulls: bool = True): + self._values = values + self._match_nulls = match_nulls + + def _as_ibis(self, x: ibis_types.Value): + if self._match_nulls and any(is_null(value) for value in self._values): + return x.isnull() | x.isin( + [val for val in self._values if not is_null(val)] + ) + else: + return x.isin(self._values) + + class BinopPartialRight(UnaryOp): def __init__(self, binop: BinaryOp, right_scalar: typing.Any): self._binop = binop @@ -231,7 +332,7 @@ class RemoteFunctionOp(UnaryOp): def __init__(self, func: typing.Callable, apply_on_null=True): if not hasattr(func, "bigframes_remote_function"): raise TypeError( - "only a bigframes remote function is supported as a callable" + f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}" ) self._func = func @@ -330,13 +431,6 @@ def or_op( ) -def isin_op( - x: ibis_types.Value, - y: ibis_types.Value, -): - return x.isin(y) - - @short_circuit_nulls() def add_op( x: ibis_types.Value, @@ -431,27 +525,43 @@ def floordiv_op( ) +def _is_float(x: ibis_types.Value): + return isinstance(x, (ibis_types.FloatingColumn, ibis_types.FloatingScalar)) + + @short_circuit_nulls() def mod_op( x: ibis_types.Value, y: ibis_types.Value, ): - # TODO(tbergeron): fully support floats, including when mixed with integer - # Pandas has inconsitency about whether N mod 0. Most conventions have this be NAN. - # For some dtypes, the result is 0 instead. This implementation results in NA always. - x_numeric = typing.cast(ibis_types.NumericValue, x) - y_numeric = typing.cast(ibis_types.NumericValue, y) + is_result_float = _is_float(x) | _is_float(y) + x_numeric = typing.cast( + ibis_types.NumericValue, + x.cast(ibis_dtypes.Decimal(precision=38, scale=9, nullable=True)) + if is_result_float + else x, + ) + y_numeric = typing.cast( + ibis_types.NumericValue, + y.cast(ibis_dtypes.Decimal(precision=38, scale=9, nullable=True)) + if is_result_float + else y, + ) # Hacky short-circuit to avoid passing zero-literal to sql backend, evaluate locally instead to null. op = y.op() if isinstance(op, ibis.expr.operations.generic.Literal) and op.value == 0: return ibis_types.null().cast(x.type()) bq_mod = x_numeric % y_numeric # Bigquery will maintain x sign here + if is_result_float: + bq_mod = typing.cast(ibis_types.NumericValue, bq_mod.cast(ibis_dtypes.float64)) + # In BigQuery returned value has the same sign as X. In pandas, the sign of y is used, so we need to flip the result if sign(x) != sign(y) return ( ibis.case() .when( - y_numeric == _ZERO, _ZERO * x_numeric + y_numeric == _ZERO, + _NAN * x_numeric if is_result_float else _ZERO * x_numeric, ) # Dummy op to propogate nulls and type from x arg .when( (y_numeric < _ZERO) & (bq_mod > _ZERO), (y_numeric + bq_mod) @@ -544,3 +654,8 @@ def clip_op( .else_(original) .end() ) + + +def is_null(value) -> bool: + # float NaN/inf should be treated as distinct from 'true' null values + return typing.cast(bool, pd.isna(value)) and not isinstance(value, float) diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py index 028ed4b6065..1687f705a1e 100644 --- a/bigframes/operations/aggregations.py +++ b/bigframes/operations/aggregations.py @@ -20,6 +20,9 @@ import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types +import bigframes.constants as constants +import third_party.bigframes_vendored.ibis.expr.operations as vendored_ibis_ops + class WindowOp: def _as_ibis(self, value: ibis_types.Column, window=None): @@ -37,6 +40,8 @@ def handles_ties(self): class AggregateOp(WindowOp): + name = "abstract_aggregate" + def _as_ibis(self, value: ibis_types.Column, window=None): raise NotImplementedError("Base class AggregateOp has no implementaiton.") @@ -51,13 +56,15 @@ def constrained_op(op, column: ibis_types.Column, window=None): return operation(op, column, window) else: raise ValueError( - f"Numeric operation cannot be applied to type {column.type()}" + f"Numeric operation cannot be applied to type {column.type()}. {constants.FEEDBACK_LINK}" ) return constrained_op class SumOp(AggregateOp): + name = "sum" + @numeric_op def _as_ibis( self, column: ibis_types.NumericColumn, window=None @@ -69,7 +76,50 @@ def _as_ibis( ) +class MedianOp(AggregateOp): + name = "median" + + @numeric_op + def _as_ibis( + self, column: ibis_types.NumericColumn, window=None + ) -> ibis_types.NumericValue: + # PERCENTILE_CONT has very few allowed windows. For example, "window + # framing clause is not allowed for analytic function percentile_cont". + if window is not None: + raise NotImplementedError( + f"Median with windowing is not supported. {constants.FEEDBACK_LINK}" + ) + + # TODO(swast): Allow switching between exact and approximate median. + # For now, the best we can do is an approximate median when we're doing + # an aggregation, as PERCENTILE_CONT is only an analytic function. + return typing.cast(ibis_types.NumericValue, column.approx_median()) + + +class ApproxQuartilesOp(AggregateOp): + def __init__(self, quartile: int): + self.name = f"{quartile*25}%" + self._quartile = quartile + + @numeric_op + def _as_ibis( + self, column: ibis_types.NumericColumn, window=None + ) -> ibis_types.NumericValue: + # PERCENTILE_CONT has very few allowed windows. For example, "window + # framing clause is not allowed for analytic function percentile_cont". + if window is not None: + raise NotImplementedError( + f"Approx Quartiles with windowing is not supported. {constants.FEEDBACK_LINK}" + ) + value = vendored_ibis_ops.ApproximateMultiQuantile( + column, num_bins=4 # type: ignore + ).to_expr()[self._quartile] + return typing.cast(ibis_types.NumericValue, value) + + class MeanOp(AggregateOp): + name = "mean" + @numeric_op def _as_ibis( self, column: ibis_types.NumericColumn, window=None @@ -78,6 +128,8 @@ def _as_ibis( class ProductOp(AggregateOp): + name = "product" + @numeric_op def _as_ibis( self, column: ibis_types.NumericColumn, window=None @@ -117,16 +169,22 @@ def _as_ibis( class MaxOp(AggregateOp): + name = "max" + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: return _apply_window_if_present(column.max(), window) class MinOp(AggregateOp): + name = "min" + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: return _apply_window_if_present(column.min(), window) class StdOp(AggregateOp): + name = "std" + @numeric_op def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: return _apply_window_if_present( @@ -135,6 +193,8 @@ def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: class VarOp(AggregateOp): + name = "var" + @numeric_op def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: return _apply_window_if_present( @@ -143,6 +203,8 @@ def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: class CountOp(AggregateOp): + name = "count" + def _as_ibis( self, column: ibis_types.Column, window=None ) -> ibis_types.IntegerValue: @@ -153,7 +215,32 @@ def skips_nulls(self): return False +class CutOp(WindowOp): + def __init__(self, bins: int): + self._bins = bins + + def _as_ibis(self, x: ibis_types.Column, window=None): + col_min = _apply_window_if_present(x.min(), window) + col_max = _apply_window_if_present(x.max(), window) + bin_width = (col_max - col_min) / self._bins + out = ibis.case() + for bin in range(self._bins - 1): + out = out.when(x <= (col_min + (bin + 1) * bin_width), bin) + out = out.when(x.notnull(), self._bins - 1) + return out.end() + + @property + def skips_nulls(self): + return False + + @property + def handles_ties(self): + return True + + class NuniqueOp(AggregateOp): + name = "nunique" + def _as_ibis( self, column: ibis_types.Column, window=None ) -> ibis_types.IntegerValue: @@ -165,6 +252,8 @@ def skips_nulls(self): class RankOp(WindowOp): + name = "rank" + def _as_ibis( self, column: ibis_types.Column, window=None ) -> ibis_types.IntegerValue: @@ -230,6 +319,8 @@ def _as_ibis( class AnyOp(AggregateOp): + name = "any" + def _as_ibis( self, column: ibis_types.Column, window=None ) -> ibis_types.BooleanValue: @@ -274,6 +365,7 @@ def _map_to_literal( sum_op = SumOp() mean_op = MeanOp() +median_op = MedianOp() product_op = ProductOp() max_op = MaxOp() min_op = MinOp() @@ -286,3 +378,26 @@ def _map_to_literal( all_op = AllOp() any_op = AnyOp() first_op = FirstOp() + + +# TODO: Alternative names and lookup from numpy function objects +AGGREGATIONS_LOOKUP: dict[str, AggregateOp] = { + op.name: op + for op in [ + sum_op, + mean_op, + median_op, + product_op, + max_op, + min_op, + std_op, + var_op, + count_op, + all_op, + any_op, + nunique_op, + ApproxQuartilesOp(1), + ApproxQuartilesOp(2), + ApproxQuartilesOp(3), + ] +} diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index 29017d2bbe3..caef33919b3 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -19,6 +19,7 @@ import ibis.expr.types as ibis_types import pandas as pd +import bigframes.constants as constants import bigframes.core.blocks as blocks import bigframes.core.scalar as scalars import bigframes.dtypes @@ -29,7 +30,9 @@ # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. # TODO(tbergeron): Convert to bytes-based limit -MAX_INLINE_SERIES_SIZE = 5000 +# TODO(swast): Address issues with string escaping and empty tables before +# re-enabling inline data (ibis.memtable) feature. +MAX_INLINE_SERIES_SIZE = -1 class SeriesMethods: @@ -47,7 +50,9 @@ def __init__( ): block = None if copy is not None and not copy: - raise ValueError("Series constructor only supports copy=True") + raise ValueError( + f"Series constructor only supports copy=True. {constants.FEEDBACK_LINK}" + ) if isinstance(data, blocks.Block): assert len(data.value_columns) == 1 assert len(data.column_labels) == 1 @@ -60,12 +65,12 @@ def __init__( if name: if not isinstance(name, str): raise NotImplementedError( - "BigQuery DataFrames only supports string series names." + f"BigQuery DataFrames only supports string series names. {constants.FEEDBACK_LINK}" ) block = block.with_column_labels([name]) if index: raise NotImplementedError( - "Series 'index' constructor parameter not supported when passing BigQuery-backed objects" + f"Series 'index' constructor parameter not supported when passing BigQuery-backed objects. {constants.FEEDBACK_LINK}" ) if dtype: block = block.multi_apply_unary_op( @@ -138,7 +143,7 @@ def _apply_binary_op( if isinstance(other, pd.Series): # TODO: Convert to BigQuery DataFrames series raise NotImplementedError( - "Pandas series not supported supported as operand." + f"Pandas series not supported supported as operand. {constants.FEEDBACK_LINK}" ) if isinstance(other, series.Series): (left, right, block) = self._align(other, how=alignment) diff --git a/bigframes/operations/strings.py b/bigframes/operations/strings.py index 57029ede416..a16ecb0d328 100644 --- a/bigframes/operations/strings.py +++ b/bigframes/operations/strings.py @@ -14,13 +14,23 @@ from __future__ import annotations -from typing import Literal, Optional, Union +import re +from typing import cast, Literal, Optional, Union +import bigframes.constants as constants +import bigframes.dataframe as df import bigframes.operations as ops import bigframes.operations.base import bigframes.series as series import third_party.bigframes_vendored.pandas.core.strings.accessor as vendorstr +# Maps from python to re2 +REGEXP_FLAGS = { + re.IGNORECASE: "i", + re.MULTILINE: "m", + re.DOTALL: "s", +} + class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods): __doc__ = vendorstr.StringMethods.__doc__ @@ -72,6 +82,84 @@ def repeat(self, repeats: int) -> series.Series: def capitalize(self) -> series.Series: return self._apply_unary_op(ops.capitalize_op) + def contains( + self, pat, case: bool = True, flags: int = 0, *, regex: bool = True + ) -> series.Series: + if not case: + return self.contains(pat, flags=flags | re.IGNORECASE, regex=True) + if regex: + re2flags = _parse_flags(flags) + if re2flags: + pat = re2flags + pat + return self._apply_unary_op(ops.ContainsRegexOp(pat)) + else: + return self._apply_unary_op(ops.ContainsStringOp(pat)) + + def extract(self, pat: str, flags: int = 0): + re2flags = _parse_flags(flags) + if re2flags: + pat = re2flags + pat + compiled = re.compile(pat) + if compiled.groups == 0: + raise ValueError("No capture groups in 'pat'") + + results: list[str] = [] + block = self._block + for i in range(compiled.groups): + labels = [ + label + for label, groupn in compiled.groupindex.items() + if i + 1 == groupn + ] + label = labels[0] if labels else str(i) + block, id = block.apply_unary_op( + self._value_column, ops.ExtractOp(pat, i + 1), result_label=label + ) + results.append(id) + block = block.select_columns(results) + return df.DataFrame(block) + + def replace( + self, + pat: Union[str, re.Pattern], + repl: str, + *, + case: Optional[bool] = None, + flags: int = 0, + regex: bool = False, + ) -> series.Series: + is_compiled = isinstance(pat, re.Pattern) + patstr = cast(str, pat.pattern if is_compiled else pat) # type: ignore + if case is False: + return self.replace(pat, repl, flags=flags | re.IGNORECASE, regex=True) + if regex: + re2flags = _parse_flags(flags) + if re2flags: + patstr = re2flags + patstr + return self._apply_unary_op(ops.ReplaceRegexOp(patstr, repl)) + else: + if is_compiled: + raise ValueError( + "Must set 'regex'=True if using compiled regex pattern." + ) + return self._apply_unary_op(ops.ReplaceStringOp(patstr, repl)) + + def startswith( + self, + pat: Union[str, tuple[str, ...]], + ) -> series.Series: + if not isinstance(pat, tuple): + pat = (pat,) + return self._apply_unary_op(ops.StartsWithOp(pat)) + + def endswith( + self, + pat: Union[str, tuple[str, ...]], + ) -> series.Series: + if not isinstance(pat, tuple): + pat = (pat,) + return self._apply_unary_op(ops.EndsWithOp(pat)) + def cat( self, others: Union[str, series.Series], @@ -79,3 +167,22 @@ def cat( join: Literal["outer", "left"] = "left", ) -> series.Series: return self._apply_binary_op(others, ops.concat_op, alignment=join) + + +def _parse_flags(flags: int) -> Optional[str]: + re2flags = [] + for reflag, re2flag in REGEXP_FLAGS.items(): + if flags & flags: + re2flags.append(re2flag) + flags = flags ^ reflag + + # Remaining flags couldn't be mapped to re2 engine + if flags: + raise NotImplementedError( + f"Could not handle RegexFlag: {flags}. {constants.FEEDBACK_LINK}" + ) + + if re2flags: + return "(?" + "".join(re2flags) + ")" + else: + return None diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index e7c14351513..cc8b4e5cc4d 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -46,40 +46,94 @@ import bigframes.series import bigframes.session import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat +import third_party.bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile + +# Support pandas dtype attribute +NA = pandas.NA +BooleanDtype = pandas.BooleanDtype +Float64Dtype = pandas.Float64Dtype +Int64Dtype = pandas.Int64Dtype +StringDtype = pandas.StringDtype +ArrowDtype = pandas.ArrowDtype # Include method definition so that the method appears in our docs for # bigframes.pandas general functions. @typing.overload def concat( - objs: Iterable[bigframes.dataframe.DataFrame], *, join, ignore_index + objs: Iterable[bigframes.series.Series], + *, + axis: typing.Literal["index", 0] = ..., + join=..., + ignore_index=..., +) -> bigframes.series.Series: + ... + + +@typing.overload +def concat( + objs: Iterable[bigframes.dataframe.DataFrame], + *, + axis: typing.Literal["index", 0] = ..., + join=..., + ignore_index=..., ) -> bigframes.dataframe.DataFrame: ... @typing.overload def concat( - objs: Iterable[bigframes.series.Series], *, join, ignore_index -) -> bigframes.series.Series: + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], + *, + axis: typing.Literal["columns", 1], + join=..., + ignore_index=..., +) -> bigframes.dataframe.DataFrame: ... +@typing.overload def concat( - objs: Union[ - Iterable[bigframes.dataframe.DataFrame], Iterable[bigframes.series.Series] - ], + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], *, + axis=..., + join=..., + ignore_index=..., +) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: + ... + + +def concat( + objs: Iterable[Union[bigframes.dataframe.DataFrame, bigframes.series.Series]], + *, + axis: typing.Union[str, int] = 0, join: Literal["inner", "outer"] = "outer", ignore_index: bool = False, ) -> Union[bigframes.dataframe.DataFrame, bigframes.series.Series]: return bigframes.core.reshape.concat( - objs=objs, join=join, ignore_index=ignore_index + objs=objs, axis=axis, join=join, ignore_index=ignore_index ) concat.__doc__ = vendored_pandas_concat.concat.__doc__ +def cut( + x: bigframes.series.Series, + bins: int, + *, + labels: Optional[bool] = None, +) -> bigframes.series.Series: + return bigframes.core.reshape.cut( + x, + bins, + labels=labels, + ) + + +cut.__doc__ = vendored_pandas_tile.cut.__doc__ + + options = config.options """Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames.""" @@ -88,9 +142,12 @@ def concat( def reset_session() -> None: - """Start a fresh session next time a function requires a session. + """Start a fresh session the next time a function requires a session. Closes the current session if it was already started. + + Returns: + None """ global _global_session @@ -124,15 +181,22 @@ def _with_default_session(func: Callable[..., _T], *args, **kwargs) -> _T: def _set_default_session_location_if_possible(query): - # If the default session has not started yet and this is the first API user - # is calling, then set the default location as per the query. + # Set the location as per the query if this is the first query the user is + # running and: + # (1) Default session has not started yet, and + # (2) Location is not set yet, and + # (3) Use of regional endpoints is not set. # If query is a table name, then it would be the location of the table. # If query is a SQL with a table, then it would be table's location. # If query is a SQL with no table, then it would be the BQ default location. - if options.bigquery._session_started or options.bigquery.use_regional_endpoints: + if ( + options.bigquery._session_started + or options.bigquery.location + or options.bigquery.use_regional_endpoints + ): return - bqclient, _, _ = bigframes.session._create_bq_clients( + bqclient, _, _, _ = bigframes.session._create_cloud_clients( project=options.bigquery.project, location=options.bigquery.location, use_regional_endpoints=options.bigquery.use_regional_endpoints, @@ -320,6 +384,16 @@ def remote_function( remote_function.__doc__ = inspect.getdoc(bigframes.session.Session.remote_function) +def read_gbq_function(function_name: str): + return _with_default_session( + bigframes.session.Session.read_gbq_function, + function_name=function_name, + ) + + +read_gbq_function.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_function) + + # Other aliases DataFrame = bigframes.dataframe.DataFrame Index = bigframes.core.indexes.Index @@ -332,6 +406,7 @@ def remote_function( "options", "read_csv", "read_gbq", + "read_gbq_function", "read_gbq_model", "read_pandas", "remote_function", diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 3fe4e38d785..5924941cc54 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -27,20 +27,28 @@ import tempfile import textwrap import time -import typing +from typing import List, NamedTuple, Optional, Sequence, TYPE_CHECKING -if typing.TYPE_CHECKING: +if TYPE_CHECKING: from bigframes.session import Session import cloudpickle import google.api_core.exceptions from google.cloud import bigquery, bigquery_connection_v1, functions_v2 +from google.cloud.bigquery.routine import Routine +from google.cloud.bigquery.standard_sql import StandardSqlTypeNames from ibis.backends.bigquery.compiler import compiles from ibis.backends.bigquery.datatypes import BigQueryType +from ibis.expr.datatypes.core import boolean +from ibis.expr.datatypes.core import DataType as IbisDataType from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type +from ibis.expr.datatypes.core import float64, int64 +from ibis.expr.datatypes.core import string as ibis_string import ibis.expr.operations as ops import ibis.expr.rules as rlz +import bigframes.constants as constants + # TODO(shobs): Change the min log level to INFO after the development stabilizes # before June 2023 logging.basicConfig( @@ -52,10 +60,14 @@ # https://docs.python.org/3/library/pickle.html#data-stream-format _pickle_protocol_version = 4 -# Input and output python types supported by BigQuery DataFrames remote functions. +# Input and output types supported by BigQuery DataFrames remote functions. # TODO(shobs): Extend the support to all types supported by BQ remote functions # https://cloud.google.com/bigquery/docs/remote-functions#limitations -_supported_io_types = set((bool, float, int, str)) +_supported_io_ibis_types = {boolean, float64, int64, ibis_string} +TYPE_ERROR_MESSAGE_FORMAT = ( + f"Type {{}} not supported, supported types are {_supported_io_ibis_types}. " + f"{constants.FEEDBACK_LINK}" +) def get_remote_function_locations(bq_location): @@ -94,9 +106,20 @@ def _run_system_command(command): if exit_code: raise RuntimeError( f"Command: {command}\nOutput: {stdout.decode()}\nError: {stderr.decode()}" + f"{constants.FEEDBACK_LINK}" ) +def routine_ref_to_string_for_query(routine_ref: bigquery.RoutineReference) -> str: + return f"`{routine_ref.project}.{routine_ref.dataset_id}`.{routine_ref.routine_id}" + + +class IbisSignature(NamedTuple): + parameter_names: List[str] + input_types: List[IbisDataType] + output_type: IbisDataType + + def get_cloud_function_name(def_, uniq_suffix=None): """Get the name of the cloud function.""" cf_name = _get_hash(def_) @@ -123,6 +146,7 @@ def __init__( self, gcp_project_id, cloud_function_region, + cloud_functions_client, bq_location, bq_dataset, bq_client, @@ -131,6 +155,7 @@ def __init__( ): self._gcp_project_id = gcp_project_id self._cloud_function_region = cloud_function_region + self._cloud_functions_client = cloud_functions_client self._bq_location = bq_location self._bq_dataset = bq_dataset self._bq_client = bq_client @@ -184,21 +209,18 @@ def create_bq_remote_function( f"{name} {BigQueryType.from_ibis(input_types[idx])}" ) create_function_ddl = f""" - CREATE OR REPLACE FUNCTION `{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name}({','.join(bq_function_args)}) - RETURNS {bq_function_return_type} - REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}` - OPTIONS ( - endpoint = "{endpoint}" - )""" + CREATE OR REPLACE FUNCTION `{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name}({','.join(bq_function_args)}) + RETURNS {bq_function_return_type} + REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}` + OPTIONS ( + endpoint = "{endpoint}" + )""" logger.info(f"Creating BQ remote function: {create_function_ddl}") + # TODO: Use session._start_query() so we get progress bar query_job = self._bq_client.query(create_function_ddl) # Make an API request. query_job.result() # Wait for the job to complete. logger.info(f"Created remote function {query_job.ddl_target_routine}") - def get_remote_function_fully_qualified_name(self, name): - "Get the fully qualilfied name for a BQ remote function." - return "{}.{}.{}".format(self._gcp_project_id, self._bq_dataset, name) - def get_cloud_function_fully_qualified_name(self, name): "Get the fully qualilfied name for a cloud function." return "projects/{}/locations/{}/functions/{}".format( @@ -207,10 +229,11 @@ def get_cloud_function_fully_qualified_name(self, name): def get_cloud_function_endpoint(self, name): """Get the http endpoint of a cloud function if it exists.""" - client = functions_v2.FunctionServiceClient() fully_qualified_name = self.get_cloud_function_fully_qualified_name(name) try: - response = client.get_function(name=fully_qualified_name) + response = self._cloud_functions_client.get_function( + name=fully_qualified_name + ) return response.service_config.uri except google.api_core.exceptions.NotFound: pass @@ -233,6 +256,34 @@ def create_bq_connection(self): def check_bq_connection_exists(self): """Check if the BigQuery Connection exists.""" client = self._bq_connection_client + if self._bq_connection_id.count(".") == 1: + bq_location, bq_connection_id = self._bq_connection_id.split(".") + if bq_location != self._bq_location: + logger.info( + f"Reset location {self._bq_location} to match the" + + f"location in connection name: {bq_location}" + ) + self._bq_location = bq_location + self._bq_connection_id = bq_connection_id + elif self._bq_connection_id.count(".") == 2: + ( + gcp_project_id, + bq_location, + bq_connection_id, + ) = self._bq_connection_id.split(".") + if gcp_project_id != self._gcp_project_id: + raise ValueError( + "The project_id does not match BigQuery connection gcp_project_id: " + f"{self._gcp_project_id}." + ) + if bq_location != self._bq_location: + logger.info( + f"Reset location {self._bq_location} to match the" + + f"location in connection name: {bq_location}" + ) + self._gcp_project_id = gcp_project_id + self._bq_location = bq_location + self._bq_connection_id = bq_connection_id request = bigquery_connection_v1.GetConnectionRequest( name=client.connection_path( self._gcp_project_id, self._bq_location, self._bq_connection_id @@ -299,7 +350,6 @@ def generate_cloud_function_main_code(self, def_, dir): def {handler_func_name}(request): request_json = request.get_json(silent=True) - print("[debug] received json request: " + str(request_json)) calls = request_json["calls"] replies = [] for call in calls: @@ -402,7 +452,9 @@ def create_cloud_function(self, def_, cf_name): # Fetch the endpoint of the just created function endpoint = self.get_cloud_function_endpoint(cf_name) if not endpoint: - raise ValueError("Couldn't fetch the http endpoint") + raise ValueError( + f"Couldn't fetch the http endpoint. {constants.FEEDBACK_LINK}" + ) logger.info( f"Successfully created cloud function {cf_name} with uri ({endpoint})" @@ -466,7 +518,8 @@ def check_cloud_function_tools_and_permissions(self): # cloud function and BigQuery remote function respectively if not shutil.which("gcloud"): raise ValueError( - "gcloud tool not installed, install it from https://cloud.google.com/sdk/docs/install" + "gcloud tool not installed, install it from https://cloud.google.com/sdk/docs/install. " + f"{constants.FEEDBACK_LINK}" ) # TODO(shobs): Check for permissions too @@ -481,110 +534,231 @@ def check_cloud_function_tools_and_permissions(self): # `cloudasset.googleapis.com` +def remote_function_node( + routine_ref: bigquery.RoutineReference, ibis_signature: IbisSignature +): + """Creates an Ibis node representing a remote function call.""" + + fields = { + name: rlz.value(type_) + for name, type_ in zip( + ibis_signature.parameter_names, ibis_signature.input_types + ) + } + + try: + fields["output_type"] = rlz.shape_like("args", dtype=ibis_signature.output_type) # type: ignore + except TypeError: + fields["output_dtype"] = property(lambda _: ibis_signature.output_type) + fields["output_shape"] = rlz.shape_like("args") + + node = type(routine_ref_to_string_for_query(routine_ref), (ops.ValueOp,), fields) # type: ignore + + @compiles(node) + def compile_node(t, op): + return "{}({})".format(node.__name__, ", ".join(map(t.translate, op.args))) + + def f(*args, **kwargs): + return node(*args, **kwargs).to_expr() + + f.bigframes_remote_function = str(routine_ref) # type: ignore + + return f + + +def ibis_type_from_python_type(t: type) -> IbisDataType: + ibis_type = python_type_to_bigquery_type(t) + assert ibis_type in _supported_io_ibis_types, TYPE_ERROR_MESSAGE_FORMAT.format( + ibis_type + ) + return ibis_type + + +def ibis_type_from_type_kind(tk: StandardSqlTypeNames) -> IbisDataType: + ibis_type = BigQueryType.to_ibis(tk) + assert ibis_type in _supported_io_ibis_types, TYPE_ERROR_MESSAGE_FORMAT.format( + ibis_type + ) + return ibis_type + + +def ibis_signature_from_python_signature( + signature: inspect.Signature, + input_types: Sequence[type], + output_type: type, +) -> IbisSignature: + return IbisSignature( + parameter_names=list(signature.parameters.keys()), + input_types=[ibis_type_from_python_type(t) for t in input_types], + output_type=ibis_type_from_python_type(output_type), + ) + + +def ibis_signature_from_routine( + routine: Routine, +) -> IbisSignature: + return IbisSignature( + parameter_names=[arg.name for arg in routine.arguments], + input_types=[ + ibis_type_from_type_kind(arg.data_type.type_kind) + for arg in routine.arguments + ], + output_type=ibis_type_from_type_kind(routine.return_type.type_kind), + ) + + +class DatasetMissingError(ValueError): + pass + + +def get_routine_reference( + routine_ref_str: str, + bigquery_client: bigquery.Client, + session: Optional[Session], +) -> bigquery.RoutineReference: + try: + # Handle cases ".." and + # ".". + return bigquery.RoutineReference.from_string( + routine_ref_str, + default_project=bigquery_client.project, + ) + except ValueError: + # Handle case of "". + if not session: + raise DatasetMissingError + + dataset_ref = bigquery.DatasetReference( + bigquery_client.project, session._session_dataset_id + ) + return dataset_ref.routine(routine_ref_str) + + # Inspired by @udf decorator implemented in ibis-bigquery package # https://github.com/ibis-project/ibis-bigquery/blob/main/ibis_bigquery/udf/__init__.py # which has moved as @js to the ibis package # https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/udf/__init__.py def remote_function( - input_types: typing.Sequence[type], + input_types: Sequence[type], output_type: type, - session: typing.Optional[Session] = None, - bigquery_client: typing.Optional[bigquery.Client] = None, - bigquery_connection_client: typing.Optional[ + session: Optional[Session] = None, + bigquery_client: Optional[bigquery.Client] = None, + bigquery_connection_client: Optional[ bigquery_connection_v1.ConnectionServiceClient ] = None, - dataset: typing.Optional[str] = None, - bigquery_connection: typing.Optional[str] = None, + cloud_functions_client: Optional[functions_v2.FunctionServiceClient] = None, + dataset: Optional[str] = None, + bigquery_connection: Optional[str] = None, reuse: bool = True, ): """Decorator to turn a user defined function into a BigQuery remote function. .. deprecated:: 0.0.1 - Use :func:`bigframes.pandas.remote_function` instead. + This is an internal method. Please use :func:`bigframes.pandas.remote_function` instead. + + .. note:: + Please make sure following is setup before using this API: + + 1. Have the below APIs enabled for your project: + + * BigQuery Connection API + * Cloud Functions API + * Cloud Run API + * Cloud Build API + * Artifact Registry API + * Cloud Resource Manager API + + This can be done from the cloud console (change `PROJECT_ID` to yours): + https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID + + Or from the gcloud CLI: + + `$ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com` + + 2. Have following IAM roles enabled for you: + + * BigQuery Data Editor (roles/bigquery.dataEditor) + * BigQuery Connection Admin (roles/bigquery.connectionAdmin) + * Cloud Functions Developer (roles/cloudfunctions.developer) + * Service Account User (roles/iam.serviceAccountUser) + * Storage Object Viewer (roles/storage.objectViewer) + * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) + + 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: + + 1. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection + 2. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + + Alternatively, the IAM could also be setup via the gcloud CLI: + + `$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`. Args: - input_types : list(type). + input_types list(type): List of input data types in the user defined function. - output_type : type. + output_type type: Data type of the output in the user defined function. - session : bigframes.Session, Optional + session (bigframes.Session, Optional): BigQuery DataFrames session to use for getting default project, dataset and BigQuery connection. - bigquery_client : google.cloud.bigquery.Client, Optional + bigquery_client (google.cloud.bigquery.Client, Optional): Client to use for BigQuery operations. If this param is not provided then bigquery client from the session would be used. - bigquery_connection_client : google.cloud.bigquery_connection_v1.ConnectionServiceClient, Optional + bigquery_connection_client (google.cloud.bigquery_connection_v1.ConnectionServiceClient, Optional): + Client to use for cloud functions operations. If this param is not + provided then functions client from the session would be used. + cloud_functions_client (google.cloud.functions_v2.FunctionServiceClient, Optional): Client to use for BigQuery connection operations. If this param is not provided then bigquery connection client from the session would be used. - dataset : str, Optional. - Dataset to use to create a BigQuery function. It should be in + dataset (str, Optional.): + Dataset in which to create a BigQuery remote function. It should be in `.` or `` format. If this - param is not provided then session dataset id would be used. - bigquery_connection : str, Optional. - Name of the BigQuery connection. If this param is not provided then - the bigquery connection from the session would be used. If it is pre - created in the same location as the `bigquery_client.location` then - it would be used, otherwise it would be created dynamically using - the `bigquery_connection_client` assuming the user has necessary - priviliges. - reuse : bool, Optional. - Reuse the remote function if already exists. - `True` by default, which will result in reusing an existing remote + parameter is not provided then session dataset id is used. + bigquery_connection (str, Optional): + Name of the BigQuery connection in the form of `CONNECTION_ID` or + `LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID`. + If this param is not provided then the bigquery connection from the session + would be used. If it is pre created in the same location as the + `bigquery_client.location` then it would be used, otherwise it is created + dynamically using the `bigquery_connection_client` assuming the user has necessary + priviliges. The PROJECT_ID should be the same as the BigQuery connection project. + reuse (bool, Optional): + Reuse the remote function if is already exists. + `True` by default, which results in reusing an existing remote function (if any) that was previously created for the same udf. - Setting it to false would force creating a unique remote function. + Setting it to false forces the creation of creating a unique remote function. If the required remote function does not exist then it would be created irrespective of this param. - Notes: - Please make sure following is setup before using this API: - - 1. Have the below APIs enabled for your project: - a. BigQuery Connection API - b. Cloud Functions API - c. Cloud Run API - d. Cloud Build API - e. Artifact Registry API - f. Cloud Resource Manager API - - This can be done from the cloud console (change PROJECT_ID to yours): - https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID - Or from the gcloud CLI: - $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com - - 2. Have following IAM roles enabled for you: - a. BigQuery Data Editor (roles/bigquery.dataEditor) - b. BigQuery Connection Admin (roles/bigquery.connectionAdmin) - c. Cloud Functions Developer (roles/cloudfunctions.developer) - d. Service Account User (roles/iam.serviceAccountUser) - e. Storage Object Viewer (roles/storage.objectViewer) - f. Project IAM Admin (roles/resourcemanager.projectIamAdmin) - (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) - - 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: - a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection - b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function - Alternatively, the IAM could also be setup via the gcloud CLI: - $ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker" - """ # A BigQuery client is required to perform BQ operations - if not bigquery_client: - if session: - bigquery_client = session.bqclient + if not bigquery_client and session: + bigquery_client = session.bqclient if not bigquery_client: raise ValueError( - "A bigquery client must be provided, either directly or via session" + "A bigquery client must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" ) # A BigQuery connection client is required to perform BQ connection operations + if not bigquery_connection_client and session: + bigquery_connection_client = session.bqconnectionclient if not bigquery_connection_client: + raise ValueError( + "A bigquery connection client must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" + ) + + # A cloud functions client is required to perform cloud functions operations + if not cloud_functions_client: if session: - bigquery_connection_client = session.bqconnectionclient - if not bigquery_connection_client: + cloud_functions_client = session.cloudfunctionsclient + if not cloud_functions_client: raise ValueError( - "A bigquery connection client must be provided, either directly or via session" + "A functions connection client must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" ) # BQ remote function must be persisted, for which we need a dataset @@ -593,16 +767,15 @@ def remote_function( dataset_ref = bigquery.DatasetReference.from_string( dataset, default_project=bigquery_client.project ) - gcp_project_id = dataset_ref.project - bq_dataset = dataset_ref.dataset_id + elif session: + dataset_ref = bigquery.DatasetReference.from_string( + session._session_dataset_id, default_project=bigquery_client.project + ) else: - gcp_project_id = bigquery_client.project - if session: - bq_dataset = session._session_dataset_id - if not gcp_project_id: - raise ValueError("Project must be provided, either directly or via session") - if not bq_dataset: - raise ValueError("Dataset must be provided, either directly or via session") + raise ValueError( + "Project and dataset must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" + ) bq_location, cloud_function_region = get_remote_function_locations( bigquery_client.location @@ -614,7 +787,8 @@ def remote_function( bigquery_connection = session._remote_udf_connection # type: ignore if not bigquery_connection: raise ValueError( - "BigQuery connection must be provided, either directly or via session" + "BigQuery connection must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" ) uniq_suffix = None @@ -628,70 +802,65 @@ def wrapper(f): raise TypeError("f must be callable, got {}".format(f)) signature = inspect.signature(f) - parameter_names = signature.parameters.keys() - - # Check supported python datatypes and convert to ibis datatypes - type_error_message_format = ( - "type {{}} not supported, supported types are {}.".format( - ", ".join([type_.__name__ for type_ in _supported_io_types]) - ) - ) - for type_ in input_types: - assert type_ in _supported_io_types, type_error_message_format.format(type_) - assert output_type in _supported_io_types, type_error_message_format.format( - output_type + ibis_signature = ibis_signature_from_python_signature( + signature, input_types, output_type ) - input_types_ibis = [ - python_type_to_bigquery_type(type_) for type_ in input_types - ] - output_type_ibis = python_type_to_bigquery_type(output_type) - - rf_node_fields = { - name: rlz.value(type) - for name, type in zip(parameter_names, input_types_ibis) - } - - try: - rf_node_fields["output_type"] = rlz.shape_like( - "args", dtype=output_type_ibis - ) - except TypeError: - rf_node_fields["output_dtype"] = property(lambda _: output_type_ibis) - rf_node_fields["output_shape"] = rlz.shape_like("args") remote_function_client = RemoteFunctionClient( - gcp_project_id, + dataset_ref.project, cloud_function_region, + cloud_functions_client, bq_location, - bq_dataset, + dataset_ref.dataset_id, bigquery_client, bigquery_connection_client, bigquery_connection, ) rf_name, cf_name = remote_function_client.provision_bq_remote_function( - f, input_types_ibis, output_type_ibis, uniq_suffix + f, ibis_signature.input_types, ibis_signature.output_type, uniq_suffix ) - rf_fully_qualified_name = f"`{gcp_project_id}.{bq_dataset}`.{rf_name}" - rf_node = type(rf_fully_qualified_name, (ops.ValueOp,), rf_node_fields) - - @compiles(rf_node) - def compiles_rf_node(t, op): - return "{}({})".format( - rf_node.__name__, ", ".join(map(t.translate, op.args)) - ) - @functools.wraps(f) - def wrapped(*args, **kwargs): - node = rf_node(*args, **kwargs) - return node.to_expr() + node = remote_function_node(dataset_ref.routine(rf_name), ibis_signature) - wrapped.__signature__ = signature - wrapped.bigframes_remote_function = ( - remote_function_client.get_remote_function_fully_qualified_name(rf_name) - ) - wrapped.bigframes_cloud_function = ( + node = functools.wraps(f)(node) + node.__signature__ = signature + node.bigframes_cloud_function = ( remote_function_client.get_cloud_function_fully_qualified_name(cf_name) ) - return wrapped + + return node return wrapper + + +def read_gbq_function( + function_name: str, + session: Optional[Session] = None, + bigquery_client: Optional[bigquery.Client] = None, +): + """ + Read an existing BigQuery function and prepare it for use in future queries. + """ + + # A BigQuery client is required to perform BQ operations + if not bigquery_client and session: + bigquery_client = session.bqclient + if not bigquery_client: + raise ValueError( + "A bigquery client must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" + ) + + try: + routine_ref = get_routine_reference(function_name, bigquery_client, session) + except DatasetMissingError: + raise ValueError( + "Project and dataset must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" + ) + + # Find the routine and get its arguments. + routine = bigquery_client.get_routine(routine_ref) + ibis_signature = ibis_signature_from_routine(routine) + + return remote_function_node(routine_ref, ibis_signature) diff --git a/bigframes/series.py b/bigframes/series.py index f008e0fb060..537991ed004 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -16,9 +16,10 @@ from __future__ import annotations +import numbers import textwrap import typing -from typing import Any, Optional, Union +from typing import Any, Mapping, Optional, Tuple, Union import google.cloud.bigquery as bigquery import ibis.expr.types as ibis_types @@ -27,6 +28,7 @@ import pandas.core.dtypes.common import typing_extensions +import bigframes.constants as constants import bigframes.core from bigframes.core import WindowSpec import bigframes.core.block_transforms as block_ops @@ -34,11 +36,16 @@ import bigframes.core.groupby as groupby import bigframes.core.indexers import bigframes.core.indexes as indexes -from bigframes.core.ordering import OrderingColumnReference, OrderingDirection +from bigframes.core.ordering import ( + OrderingColumnReference, + OrderingDirection, + STABLE_SORTS, +) import bigframes.core.scalar as scalars import bigframes.core.window import bigframes.dataframe import bigframes.dtypes +import bigframes.formatting_helpers as formatter import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops import bigframes.operations.base @@ -46,7 +53,8 @@ import bigframes.operations.strings as strings import third_party.bigframes_vendored.pandas.core.series as vendored_pandas_series -LevelsType = typing.Union[str, int, typing.Sequence[typing.Union[str, int]]] +LevelType = typing.Union[str, int] +LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]] class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series): @@ -100,21 +108,68 @@ def values(self) -> numpy.ndarray: @property def query_job(self) -> Optional[bigquery.QueryJob]: + """BigQuery job metadata for the most recent query. + + Returns: + The most recent `QueryJob + `_. + """ + if self._query_job is None: + self._set_internal_query_job(self._compute_dry_run()) return self._query_job + def _set_internal_query_job(self, query_job: bigquery.QueryJob): + self._query_job = query_job + def __len__(self): return self.shape[0] def copy(self) -> Series: return Series(self._block) - def rename(self, index: Optional[str], **kwargs) -> Series: + def rename( + self, index: Union[blocks.Label, Mapping[Any, Any]] = None, **kwargs + ) -> Series: if len(kwargs) != 0: raise NotImplementedError( - "rename does not currently support any keyword arguments." + f"rename does not currently support any keyword arguments. {constants.FEEDBACK_LINK}" ) - block = self._block.with_column_labels([index]) - return Series(block) + + # rename the Series name + if index is None or isinstance( + index, str + ): # Python 3.9 doesn't allow isinstance of Optional + index = typing.cast(Optional[str], index) + block = self._block.with_column_labels([index]) + return Series(block) + + # rename the index + if isinstance(index, Mapping): + index = typing.cast(Mapping[Any, Any], index) + block = self._block + for k, v in index.items(): + new_idx_ids = [] + for idx_id, idx_dtype in zip(block.index_columns, block.index_dtypes): + # Will throw if key type isn't compatible with index type, which leads to invalid SQL. + block.create_constant(k, dtype=idx_dtype) + + # Will throw if value type isn't compatible with index type. + block, const_id = block.create_constant(v, dtype=idx_dtype) + block, cond_id = block.apply_unary_op( + idx_id, ops.BinopPartialRight(ops.ne_op, k) + ) + block, new_idx_id = block.apply_ternary_op( + idx_id, cond_id, const_id, ops.where_op + ) + + new_idx_ids.append(new_idx_id) + block = block.drop_columns([const_id, cond_id]) + + block = block.set_index(new_idx_ids, index_labels=block.index_labels) + + return Series(block) + + raise ValueError(f"Unsupported type of parameter index: {type(index)}") def rename_axis( self, @@ -123,7 +178,7 @@ def rename_axis( ) -> Series: if len(kwargs) != 0: raise NotImplementedError( - "rename_axis does not currently support any keyword arguments." + f"rename_axis does not currently support any keyword arguments. {constants.FEEDBACK_LINK}" ) # limited implementation: the new index name is simply the 'mapper' parameter if _is_list_like(mapper): @@ -151,8 +206,15 @@ def __repr__(self) -> str: # maybe we just print the job metadata that we have so far? # TODO(swast): Avoid downloading the whole series by using job # metadata, like we do with DataFrame. - preview = self.compute() - return repr(preview) + opts = bigframes.options.display + max_results = opts.max_rows + if opts.repr_mode == "deferred": + return formatter.repr_query_job(self.query_job) + + pandas_df, _, query_job = self._block.retrieve_repr_request_results(max_results) + self._set_internal_query_job(query_job) + + return repr(pandas_df.iloc[:, 0]) def _to_ibis_expr(self): """Creates an Ibis table expression representing the Series.""" @@ -168,31 +230,77 @@ def astype( ) -> Series: return self._apply_unary_op(bigframes.operations.AsTypeOp(dtype)) - def compute(self) -> pandas.Series: - """Executes deferred operations and downloads the results.""" - df, query_job = self._block.compute((self._value_column,)) - self._query_job = query_job + def to_pandas( + self, + max_download_size: Optional[int] = None, + sampling_method: Optional[str] = None, + random_state: Optional[int] = None, + ) -> pandas.Series: + """Writes Series to pandas Series. + + Args: + max_download_size (int, default None): + Download size threshold in MB. If max_download_size is exceeded when downloading data + (e.g., to_pandas()), the data will be downsampled if + bigframes.options.sampling.enable_downsampling is True, otherwise, an error will be + raised. If set to a value other than None, this will supersede the global config. + sampling_method (str, default None): + Downsampling algorithms to be chosen from, the choices are: "head": This algorithm + returns a portion of the data from the beginning. It is fast and requires minimal + computations to perform the downsampling; "uniform": This algorithm returns uniform + random samples of the data. If set to a value other than None, this will supersede + the global config. + random_state (int, default None): + The seed for the uniform downsampling algorithm. If provided, the uniform method may + take longer to execute and require more computation. If set to a value other than + None, this will supersede the global config. + + Returns: + pandas.Series: A pandas Series with all rows of this Series if the data_sampling_threshold_mb + is not exceeded; otherwise, a pandas Series with downsampled rows of the DataFrame. + """ + df, query_job = self._block.to_pandas( + (self._value_column,), + max_download_size=max_download_size, + sampling_method=sampling_method, + random_state=random_state, + ) + self._set_internal_query_job(query_job) series = df[self._value_column] series.name = self._name return series - def drop(self, labels: blocks.Label | typing.Sequence[blocks.Label] = None): - block = self._block - index_column = block.index_columns[0] + def _compute_dry_run(self) -> bigquery.QueryJob: + return self._block._compute_dry_run((self._value_column,)) + + def drop( + self, + labels: typing.Any = None, + *, + axis: typing.Union[int, str] = 0, + index: typing.Any = None, + columns: Union[blocks.Label, typing.Iterable[blocks.Label]] = None, + level: typing.Optional[LevelType] = None, + ) -> Series: + if labels and index: + raise ValueError("Must specify exacly one of 'labels' or 'index'") + index = labels or index + # ignore axis, columns params + block = self._block + level_id = self._resolve_levels(level or 0)[0] if _is_list_like(labels): block, inverse_condition_id = block.apply_unary_op( - index_column, ops.partial_right(ops.isin_op, labels) + level_id, ops.IsInOp(index, match_nulls=True) ) block, condition_id = block.apply_unary_op( inverse_condition_id, ops.invert_op ) - else: block, condition_id = block.apply_unary_op( - index_column, ops.partial_right(ops.ne_op, labels) + level_id, ops.partial_right(ops.ne_op, labels) ) - block = block.filter(condition_id) + block = block.filter(condition_id, keep_null=True) block = block.drop_columns([condition_id]) return Series(block.select_column(self._value_column)) @@ -248,6 +356,11 @@ def cummin(self) -> Series: agg_ops.min_op, bigframes.core.WindowSpec(following=0) ) + def cumprod(self) -> Series: + return self._apply_window_op( + agg_ops.product_op, bigframes.core.WindowSpec(following=0) + ) + def shift(self, periods: int = 1) -> Series: window = bigframes.core.WindowSpec( preceding=periods if periods > 0 else None, @@ -255,8 +368,8 @@ def shift(self, periods: int = 1) -> Series: ) return self._apply_window_op(agg_ops.ShiftOp(periods), window) - def diff(self) -> Series: - return self - self.shift(1) + def diff(self, periods: int = 1) -> Series: + return self - self.shift(periods=periods) def rank( self, @@ -337,64 +450,66 @@ def notna(self) -> "Series": notnull = notna - def __and__(self, other: bool | int | Series | pandas.Series) -> Series: + def __and__(self, other: bool | int | Series) -> Series: return self._apply_binary_op(other, ops.and_op) __rand__ = __and__ - def __or__(self, other: bool | int | Series | pandas.Series) -> Series: + def __or__(self, other: bool | int | Series) -> Series: return self._apply_binary_op(other, ops.or_op) __ror__ = __or__ - def __add__(self, other: float | int | Series | pandas.Series) -> Series: + def __add__(self, other: float | int | Series) -> Series: return self.add(other) - def __radd__(self, other: float | int | Series | pandas.Series) -> Series: + def __radd__(self, other: float | int | Series) -> Series: return self.radd(other) - def add(self, other: float | int | Series | pandas.Series) -> Series: + def add(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.add_op) - def radd(self, other: float | int | Series | pandas.Series) -> Series: + def radd(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.add_op)) - def __sub__(self, other: float | int | Series | pandas.Series) -> Series: + def __sub__(self, other: float | int | Series) -> Series: return self.sub(other) - def __rsub__(self, other: float | int | Series | pandas.Series) -> Series: + def __rsub__(self, other: float | int | Series) -> Series: return self.rsub(other) - def sub(self, other: float | int | Series | pandas.Series) -> Series: + def sub(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.sub_op) - def rsub(self, other: float | int | Series | pandas.Series) -> Series: + def rsub(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.sub_op)) - def __mul__(self, other: float | int | Series | pandas.Series) -> Series: + subtract = sub + + def __mul__(self, other: float | int | Series) -> Series: return self.mul(other) - def __rmul__(self, other: float | int | Series | pandas.Series) -> Series: + def __rmul__(self, other: float | int | Series) -> Series: return self.rmul(other) - def mul(self, other: float | int | Series | pandas.Series) -> Series: + def mul(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.mul_op) - def rmul(self, other: float | int | Series | pandas.Series) -> Series: + def rmul(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.mul_op)) multiply = mul - def __truediv__(self, other: float | int | Series | pandas.Series) -> Series: + def __truediv__(self, other: float | int | Series) -> Series: return self.truediv(other) - def __rtruediv__(self, other: float | int | Series | pandas.Series) -> Series: + def __rtruediv__(self, other: float | int | Series) -> Series: return self.rtruediv(other) - def truediv(self, other: float | int | Series | pandas.Series) -> Series: + def truediv(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.div_op) - def rtruediv(self, other: float | int | Series | pandas.Series) -> Series: + def rtruediv(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.div_op)) div = truediv @@ -403,22 +518,22 @@ def rtruediv(self, other: float | int | Series | pandas.Series) -> Series: rdiv = rtruediv - def __floordiv__(self, other: float | int | Series | pandas.Series) -> Series: + def __floordiv__(self, other: float | int | Series) -> Series: return self.floordiv(other) - def __rfloordiv__(self, other: float | int | Series | pandas.Series) -> Series: + def __rfloordiv__(self, other: float | int | Series) -> Series: return self.rfloordiv(other) - def floordiv(self, other: float | int | Series | pandas.Series) -> Series: + def floordiv(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.floordiv_op) - def rfloordiv(self, other: float | int | Series | pandas.Series) -> Series: + def rfloordiv(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.floordiv_op)) - def __lt__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + def __lt__(self, other: float | int | Series) -> Series: # type: ignore return self.lt(other) - def __le__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + def __le__(self, other: float | int | Series) -> Series: # type: ignore return self.le(other) def lt(self, other) -> Series: @@ -427,10 +542,10 @@ def lt(self, other) -> Series: def le(self, other) -> Series: return self._apply_binary_op(other, ops.le_op) - def __gt__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + def __gt__(self, other: float | int | Series) -> Series: # type: ignore return self.gt(other) - def __ge__(self, other: float | int | Series | pandas.Series) -> Series: # type: ignore + def __ge__(self, other: float | int | Series) -> Series: # type: ignore return self.ge(other) def gt(self, other) -> Series: @@ -451,6 +566,16 @@ def mod(self, other) -> Series: # type: ignore def rmod(self, other) -> Series: # type: ignore return self._apply_binary_op(other, ops.reverse(ops.mod_op)) + def divmod(self, other) -> Tuple[Series, Series]: # type: ignore + # TODO(huanc): when self and other both has dtype int and other contains zeros, + # the output should be dtype float, both floordiv and mod returns dtype int in this case. + return (self.floordiv(other), self.mod(other)) + + def rdivmod(self, other) -> Tuple[Series, Series]: # type: ignore + # TODO(huanc): when self and other both has dtype int and self contains zeros, + # the output should be dtype float, both floordiv and mod returns dtype int in this case. + return (self.rfloordiv(other), self.rmod(other)) + def __matmul__(self, other): return (self * other).sum() @@ -503,11 +628,48 @@ def _central_moment(self, n: int) -> float: delta_power = delta_power * mean_deltas return delta_power.mean() - def kurt(self) -> float: - # TODO(tbergeron): Cache intermediate count/moment/etc. statistics at block level + def agg(self, func: str | typing.Sequence[str]) -> scalars.Scalar | Series: + if _is_list_like(func): + if self.dtype not in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES: + raise NotImplementedError( + f"Multiple aggregations only supported on numeric series. {constants.FEEDBACK_LINK}" + ) + aggregations = [agg_ops.AGGREGATIONS_LOOKUP[f] for f in func] + return Series( + self._block.summarize( + [self._value_column], + aggregations, + ) + ) + else: + + return self._apply_aggregation( + agg_ops.AGGREGATIONS_LOOKUP[typing.cast(str, func)] + ) + + def skew(self): + count = self.count() + if count < 3: + return pandas.NA + + moment3 = self._central_moment(3) + moment2 = self.var() * (count - 1) / count # Convert sample var to pop var + + # See G1 estimator: + # https://en.wikipedia.org/wiki/Skewness#Sample_skewness + numerator = moment3 + denominator = moment2 ** (3 / 2) + adjustment = (count * (count - 1)) ** 0.5 / (count - 2) + + return (numerator / denominator) * adjustment + + def kurt(self): count = self.count() + if count < 4: + return pandas.NA + moment4 = self._central_moment(4) - moment2 = self._central_moment(2) # AKA: Population Variance + moment2 = self.var() * (count - 1) / count # Convert sample var to pop var # Kurtosis is often defined as the second standardize moment: moment(4)/moment(2)**2 # Pandas however uses Fisher’s estimator, implemented below @@ -523,8 +685,8 @@ def mode(self) -> Series: block = self._block # Approach: Count each value, return each value for which count(x) == max(counts)) block, agg_ids = block.aggregate( - [self._value_column], - ((self._value_column, agg_ops.count_op),), + by_column_ids=[self._value_column], + aggregations=((self._value_column, agg_ops.count_op),), as_index=False, ) value_count_col_id = agg_ids[0] @@ -551,6 +713,13 @@ def mode(self) -> Series: def mean(self) -> float: return typing.cast(float, self._apply_aggregation(agg_ops.mean_op)) + def median(self, *, exact: bool = False) -> float: + if exact: + raise NotImplementedError( + f"Only approximate median is supported. {constants.FEEDBACK_LINK}" + ) + return typing.cast(float, self._apply_aggregation(agg_ops.median_op)) + def sum(self) -> float: return typing.cast(float, self._apply_aggregation(agg_ops.sum_op)) @@ -622,12 +791,26 @@ def argmin(self) -> scalars.Scalar: scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] ) - def __getitem__(self, indexer: Series): + def __getitem__(self, indexer): # TODO: enforce stricter alignment, should fail if indexer is missing any keys. - (left, right, block) = self._align(indexer, "left") - block = block.filter(right) - block = block.select_column(left) - return Series(block) + use_iloc = ( + isinstance(indexer, slice) + and all( + isinstance(x, numbers.Integral) or (x is None) + for x in [indexer.start, indexer.stop, indexer.step] + ) + ) or ( + isinstance(indexer, numbers.Integral) + and not isinstance(self._block.index.dtypes[0], pandas.Int64Dtype) + ) + if use_iloc: + return self.iloc[indexer] + if isinstance(indexer, Series): + (left, right, block) = self._align(indexer, "left") + block = block.filter(right) + block = block.select_column(left) + return Series(block) + return self.loc[indexer] def __getattr__(self, key: str): if hasattr(pandas.Series, key): @@ -635,11 +818,7 @@ def __getattr__(self, key: str): textwrap.dedent( f""" BigQuery DataFrames has not yet implemented an equivalent to - 'pandas.Series.{key}'. Please check - https://github.com/googleapis/python-bigquery-dataframes/issues for - existing feature requests, or file your own. - Please include information about your use case, as well as - relevant code snippets. + 'pandas.Series.{key}'. {constants.FEEDBACK_LINK} """ ) ) @@ -652,12 +831,7 @@ def _align3(self, other1: Series | scalars.Scalar, other2: Series | scalars.Scal return (values[0], values[1], values[2], index) def _apply_aggregation(self, op: agg_ops.AggregateOp) -> Any: - aggregation_result = typing.cast( - ibis_types.Scalar, op._as_ibis(self[self.notnull()]._to_ibis_expr()) - ) - return bigframes.core.scalar.DeferredScalar( - aggregation_result, self._block._expr._session - ).compute() + return self._block.get_stat(self._value_column, op) def _apply_window_op( self, @@ -687,7 +861,9 @@ def value_counts( ) return Series(block) - def sort_values(self, *, axis=0, ascending=True, na_position="last") -> Series: + def sort_values( + self, *, axis=0, ascending=True, kind: str = "quicksort", na_position="last" + ) -> Series: if na_position not in ["first", "last"]: raise ValueError("Param na_position must be one of 'first' or 'last'") direction = OrderingDirection.ASC if ascending else OrderingDirection.DESC @@ -698,7 +874,8 @@ def sort_values(self, *, axis=0, ascending=True, na_position="last") -> Series: direction=direction, na_last=(na_position == "last"), ) - ] + ], + stable=kind in STABLE_SORTS, ) return Series(block) @@ -750,6 +927,8 @@ def groupby( raise ValueError("as_index=False only valid with DataFrame") if axis: raise ValueError("No axis named {} for object type Series".format(level)) + if not as_index: + raise ValueError("'as_index'=False only applies to DataFrame") if by is not None: return self._groupby_values(by, dropna) if level is not None: @@ -765,7 +944,7 @@ def _groupby_level( return groupby.SeriesGroupBy( self._block, self._value_column, - self._resolve_levels(level), + by_col_ids=self._resolve_levels(level), value_name=self.name, dropna=dropna, ) @@ -805,14 +984,14 @@ def _groupby_values( matches = block.index_name_to_col_id.get(key, []) if len(matches) != 1: raise ValueError( - f"GroupBy key {key} does not map to unambiguous index level" + f"GroupBy key {key} does not match a unique index level. BigQuery DataFrames only interprets lists of strings as index level names, not directly as per-row group assignments." ) grouping_cols = [*grouping_cols, matches[0]] return groupby.SeriesGroupBy( block, value_col, - grouping_cols, + by_col_ids=grouping_cols, value_name=self.name, dropna=dropna, ) @@ -851,7 +1030,8 @@ def mask(self, cond, other=None) -> Series: if not isinstance(cond, Series): raise TypeError( - f"Only bigframes series condition is supported, received {type(cond).__name__}" + f"Only bigframes series condition is supported, received {type(cond).__name__}. " + f"{constants.FEEDBACK_LINK}" ) return self.where(~cond, other) @@ -862,13 +1042,13 @@ def to_frame(self) -> bigframes.dataframe.DataFrame: def to_csv(self, path_or_buf=None, **kwargs) -> typing.Optional[str]: # TODO(b/280651142): Implement version that leverages bq export native csv support to bypass local pandas step. - return self.compute().to_csv(path_or_buf, **kwargs) + return self.to_pandas().to_csv(path_or_buf, **kwargs) def to_dict(self, into: type[dict] = dict) -> typing.Mapping: - return typing.cast(dict, self.compute().to_dict(into)) + return typing.cast(dict, self.to_pandas().to_dict(into)) def to_excel(self, excel_writer, sheet_name="Sheet1", **kwargs) -> None: - return self.compute().to_excel(excel_writer, sheet_name, **kwargs) + return self.to_pandas().to_excel(excel_writer, sheet_name, **kwargs) def to_json( self, @@ -879,17 +1059,17 @@ def to_json( **kwargs, ) -> typing.Optional[str]: # TODO(b/280651142): Implement version that leverages bq export native csv support to bypass local pandas step. - return self.compute().to_json(path_or_buf, **kwargs) + return self.to_pandas().to_json(path_or_buf, **kwargs) def to_latex( self, buf=None, columns=None, header=True, index=True, **kwargs ) -> typing.Optional[str]: - return self.compute().to_latex( + return self.to_pandas().to_latex( buf, columns=columns, header=header, index=index, **kwargs ) def tolist(self) -> list: - return self.compute().to_list() + return self.to_pandas().to_list() to_list = tolist @@ -900,17 +1080,17 @@ def to_markdown( index: bool = True, **kwargs, ) -> typing.Optional[str]: - return self.compute().to_markdown(buf, mode=mode, index=index, **kwargs) # type: ignore + return self.to_pandas().to_markdown(buf, mode=mode, index=index, **kwargs) # type: ignore def to_numpy( self, dtype=None, copy=False, na_value=None, **kwargs ) -> numpy.ndarray: - return self.compute().to_numpy(dtype, copy, na_value, **kwargs) + return self.to_pandas().to_numpy(dtype, copy, na_value, **kwargs) __array__ = to_numpy def to_pickle(self, path, **kwargs) -> None: - return self.compute().to_pickle(path, **kwargs) + return self.to_pandas().to_pickle(path, **kwargs) def to_string( self, @@ -925,7 +1105,7 @@ def to_string( max_rows=None, min_rows=None, ) -> typing.Optional[str]: - return self.compute().to_string( + return self.to_pandas().to_string( buf, na_rep, float_format, @@ -939,7 +1119,7 @@ def to_string( ) def to_xarray(self): - return self.compute().to_xarray() + return self.to_pandas().to_xarray() # Keep this at the bottom of the Series class to avoid # confusing type checker by overriding str diff --git a/bigframes/session.py b/bigframes/session.py index 9590ecb8c77..28a38f93073 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -46,6 +46,7 @@ import google.cloud.bigquery as bigquery import google.cloud.bigquery_connection_v1 import google.cloud.bigquery_storage_v1 +import google.cloud.functions_v2 import google.cloud.storage as storage # type: ignore import ibis import ibis.backends.bigquery as ibis_bigquery @@ -56,15 +57,21 @@ import pydata_google_auth import bigframes._config.bigquery_options as bigquery_options +import bigframes.constants as constants import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.guid as guid -from bigframes.core.ordering import OrderingColumnReference +import bigframes.core.io as bigframes_io +from bigframes.core.ordering import IntegerEncoding, OrderingColumnReference import bigframes.dataframe as dataframe import bigframes.formatting_helpers as formatting_helpers -import bigframes.ml.loader +from bigframes.remote_function import read_gbq_function as bigframes_rgf from bigframes.remote_function import remote_function as bigframes_rf import bigframes.version + +# Even though the ibis.backends.bigquery.registry import is unused, it's needed +# to register new and replacement ops with the Ibis BigQuery backend. +import third_party.bigframes_vendored.ibis.backends.bigquery.registry # noqa import third_party.bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq import third_party.bigframes_vendored.pandas.io.parquet as third_party_pandas_parquet import third_party.bigframes_vendored.pandas.io.parsers.readers as third_party_pandas_readers @@ -81,13 +88,10 @@ _BIGQUERYCONNECTION_REGIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com" _BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com" -# TODO(swast): Need to connect to regional endpoints when performing remote -# functions operations (BQ Connection API, Cloud Run / Cloud Functions). +_MAX_CLUSTER_COLUMNS = 4 -# pydata-google-auth credentials in case auth credentials are not available -# otherwise -_pydata_google_auth_credentials: Optional[google.auth.credentials.Credentials] = None -_pydata_google_auth_project: Optional[str] = None +# TODO(swast): Need to connect to regional endpoints when performing remote +# functions operations (BQ Connection IAM, Cloud Run / Cloud Functions). logger = logging.getLogger(__name__) @@ -97,53 +101,11 @@ def _is_query(query_or_table: str) -> bool: return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None -# TODO(shobs): Remove it after the same is available via pydata-google-auth -# after https://github.com/pydata/pydata-google-auth/pull/71 is merged, released -# and upgraded in the google colab image. -def _ensure_application_default_credentials_in_colab_environment(): - # This is a special handling for google colab environment where we want to - # use the colab specific authentication flow - # https://github.com/googlecolab/colabtools/blob/3c8772efd332289e1c6d1204826b0915d22b5b95/google/colab/auth.py#L209 - try: - from google.colab import auth - - auth.authenticate_user() - except Exception: - # We are catching a broad exception class here because we want to be - # agnostic to anything that could internally go wrong in the google - # colab auth. Some of the known exception we want to pass on are: - # - # ModuleNotFoundError: No module named 'google.colab' - # ImportError: cannot import name 'auth' from 'google.cloud' - # MessageError: Error: credential propagation was unsuccessful - # - # The MessageError happens on Vertex Colab when it fails to resolve auth - # from the Compute Engine Metadata server. - pass - - -pydata_google_auth.auth._ensure_application_default_credentials_in_colab_environment = ( - _ensure_application_default_credentials_in_colab_environment -) - - def _get_default_credentials_with_project(): - global _pydata_google_auth_credentials, _pydata_google_auth_project - if not _pydata_google_auth_credentials or not _pydata_google_auth_credentials.valid: - # We want to initiate auth via a non-local web server which - # particularly helps in a cloud notebook environment where the - # machine running the notebook UI and the VM running the notebook - # runtime are not the same. - # TODO(shobs, b/278903498): Use BigQuery DataFrames's own client id - # and secret - ( - _pydata_google_auth_credentials, - _pydata_google_auth_project, - ) = pydata_google_auth.default(_SCOPES, use_local_webserver=False) - return _pydata_google_auth_credentials, _pydata_google_auth_project + return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False) -def _create_bq_clients( +def _create_cloud_clients( project: Optional[str], location: Optional[str], use_regional_endpoints: Optional[bool], @@ -152,6 +114,7 @@ def _create_bq_clients( bigquery.Client, google.cloud.bigquery_connection_v1.ConnectionServiceClient, google.cloud.bigquery_storage_v1.BigQueryReadClient, + google.cloud.functions_v2.FunctionServiceClient, ]: """Create and initialize BigQuery client objects.""" @@ -170,7 +133,10 @@ def _create_bq_clients( ) if not project: - raise ValueError("Project must be set to initialize BigQuery client.") + raise ValueError( + "Project must be set to initialize BigQuery client. " + "Try setting `bigframes.options.bigquery.project` first." + ) if use_regional_endpoints: bq_options = google.api_core.client_options.ClientOptions( @@ -193,6 +159,7 @@ def _create_bq_clients( client_options=bq_options, credentials=credentials, project=project, + location=location, ) bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo( @@ -213,7 +180,15 @@ def _create_bq_clients( credentials=credentials, ) - return bqclient, bqconnectionclient, bqstorageclient + functions_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + cloudfunctionsclient = google.cloud.functions_v2.FunctionServiceClient( + client_info=functions_info, + credentials=credentials, + ) + + return bqclient, bqconnectionclient, bqstorageclient, cloudfunctionsclient class Session( @@ -242,7 +217,8 @@ def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): self.bqclient, self.bqconnectionclient, self.bqstorageclient, - ) = _create_bq_clients( + self.cloudfunctionsclient, + ) = _create_cloud_clients( project=context.project, location=self._location, use_regional_endpoints=context.use_regional_endpoints, @@ -310,8 +286,23 @@ def close(self): 24 hours of inactivity or after 7 days.""" if self._session_id is not None and self.bqclient is not None: abort_session_query = "CALL BQ.ABORT_SESSION('{}')".format(self._session_id) - query_job = self.bqclient.query(abort_session_query) - query_job.result() # blocks until finished + try: + query_job = self.bqclient.query(abort_session_query) + query_job.result() # blocks until finished + except google.api_core.exceptions.BadRequest as e: + # Ignore the exception when the BQ session itself has expired + # https://cloud.google.com/bigquery/docs/sessions-terminating#auto-terminate_a_session + if not e.message.startswith( + f"Session {self._session_id} has expired and is no longer available." + ): + raise + except google.auth.exceptions.RefreshError: + # The refresh token may itself have been invalidated or expired + # https://developers.google.com/identity/protocols/oauth2#expiration + # Don't raise the exception in this case while closing the + # BigFrames session, so that the end user has a path for getting + # out of a bad session due to unusable credentials. + pass self._session_id = None def read_gbq( @@ -367,16 +358,10 @@ def read_gbq_query( else: index_cols = list(index_col) - # Make sure we cluster by the index column so that subsequent - # operations are as speedy as they can be. - if index_cols: - destination: bigquery.Table | bigquery.TableReference = ( - self._query_to_session_table(query, index_cols) - ) - else: - _, query_job = self._start_query(query) - query_job.result() # Wait for job to finish. - destination = query_job.destination + # Can't cluster since don't know if index_cols are clusterable data types + # TODO(tbergeron): Maybe use dryrun to determine types of index_cols to see if can cluster + _, query_job = self._start_query(query) + destination = query_job.destination # If there was no destination table, that means the query must have # been DDL or DML. Return some job metadata, instead. @@ -411,6 +396,9 @@ def read_gbq_table( See also: :meth:`Session.read_gbq`. """ + if max_results and max_results <= 0: + raise ValueError("`max_results` should be a positive number.") + # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so # these docstrings are inline. # TODO(swast): Can we re-use the temp table from other reads in the @@ -425,7 +413,6 @@ def read_gbq_table( f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" ) else: - # TODO(swast): Read from a table snapshot so that reads are consistent. table_expression = self.ibis_client.table( table_ref.table_id, database=f"{table_ref.project}.{table_ref.dataset_id}", @@ -442,6 +429,8 @@ def read_gbq_table( else: index_cols = list(index_col) + hidden_cols: typing.Sequence[str] = () + for key in index_cols: if key not in table_expression.columns: raise ValueError( @@ -466,7 +455,7 @@ def read_gbq_table( SELECT (SELECT COUNT(*) FROM full_table) AS total_count, (SELECT COUNT(*) FROM distinct_table) AS distinct_count """ - results, _ = self._start_query(is_unique_sql) + results, query_job = self._start_query(is_unique_sql) row = next(iter(results)) total_count = row["total_count"] @@ -476,9 +465,24 @@ def read_gbq_table( ordering_value_columns=[ core.OrderingColumnReference(column_id) for column_id in index_cols ], + total_ordering_columns=frozenset(index_cols), ) - if not is_total_ordering: + # We have a total ordering, so query via "time travel" so that + # the underlying data doesn't mutate. + if is_total_ordering: + + # Get the timestamp from the job metadata rather than the query + # text so that the query for determining uniqueness of the ID + # columns can be cached. + current_timestamp = query_job.started + + # The job finished, so we should have a start time. + assert current_timestamp is not None + table_expression = self.ibis_client.sql( + bigframes_io.create_snapshot_sql(table_ref, current_timestamp) + ) + else: # Make sure when we generate an ordering, the row_number() # coresponds to the index columns. table_expression = table_expression.order_by(index_cols) @@ -491,26 +495,37 @@ def read_gbq_table( """, ) ) + + # When ordering by index columns, apply limit after ordering to + # make limit more predictable. + if max_results is not None: + table_expression = table_expression.limit(max_results) else: + if max_results is not None: + # Apply limit before generating rownums and creating temp table + # This makes sure the offsets are valid and limits the number of + # rows for which row numbers must be generated + table_expression = table_expression.limit(max_results) table_expression, ordering = self._create_sequential_ordering( table_expression ) - ordering_id_column = ordering.ordering_id - assert ordering_id_column is not None + hidden_cols = ( + (ordering.total_order_col.column_id,) + if ordering.total_order_col + else () + ) + assert len(ordering.ordering_value_columns) > 0 is_total_ordering = True - index_cols = [ordering_id_column] - index_labels = [None] - - if max_results is not None: - if max_results <= 0: - raise ValueError("`max_results` should be a positive number.") - table_expression = table_expression.limit(max_results) + # Block constructor will generate default index if passed empty + index_cols = [] + index_labels = [] return self._read_gbq_with_ordering( table_expression=table_expression, col_order=col_order, index_cols=index_cols, index_labels=index_labels, + hidden_cols=hidden_cols, ordering=ordering, is_total_ordering=is_total_ordering, ) @@ -522,16 +537,23 @@ def _read_gbq_with_ordering( col_order: Iterable[str] = (), index_cols: Sequence[str] = (), index_labels: Sequence[Optional[str]] = (), + hidden_cols: Sequence[str] = (), ordering: core.ExpressionOrdering, is_total_ordering: bool = False, ) -> dataframe.DataFrame: """Internal helper method that loads DataFrame from Google BigQuery given an ordering column. Args: - table_expression: an ibis table expression to be executed in BigQuery. - col_order: List of BigQuery column names in the desired order for results DataFrame. - index_cols: List of column names to use as the index or multi-index. - ordering: Column name to be used for ordering. If not supplied, a default ordering is generated. + table_expression: + an ibis table expression to be executed in BigQuery. + col_order: + List of BigQuery column names in the desired order for results DataFrame. + index_cols: + List of column names to use as the index or multi-index. + hidden_cols: + Columns that should be hidden. Ordering columns may (not always) be hidden + ordering: + Column name to be used for ordering. If not supplied, a default ordering is generated. Returns: A DataFrame representing results of the query or table. @@ -542,30 +564,23 @@ def _read_gbq_with_ordering( f"Got {len(index_labels)}, expected {len(index_cols)}." ) - if not index_cols: - raise ValueError("Need at least 1 index column.") - # Logic: # no total ordering, index -> create sequential order, ordered by index, use for both ordering and index # total ordering, index -> use ordering as ordering, index as index # This code block ensures the existence of a total ordering. + column_keys = list(col_order) + if len(column_keys) == 0: + non_value_columns = set([*index_cols, *hidden_cols]) + column_keys = [ + key for key in table_expression.columns if key not in non_value_columns + ] if not is_total_ordering: # Rows are not ordered, we need to generate a default ordering and materialize it table_expression, ordering = self._create_sequential_ordering( table_expression, index_cols ) - index_col_values = [table_expression[index_id] for index_id in index_cols] - - column_keys = list(col_order) - if len(column_keys) == 0: - non_columns = set(index_cols) - if ordering.ordering_id is not None: - non_columns.add(ordering.ordering_id) - column_keys = [ - key for key in table_expression.columns if key not in non_columns - ] return self._read_ibis( table_expression, index_col_values, @@ -589,7 +604,7 @@ def _read_bigquery_load_job( index_cols = list(index_col) if not job_config.clustering_fields and index_cols: - job_config.clustering_fields = index_cols + job_config.clustering_fields = index_cols[:_MAX_CLUSTER_COLUMNS] if isinstance(filepath_or_buffer, str): if filepath_or_buffer.startswith("gs://"): @@ -606,7 +621,7 @@ def _read_bigquery_load_job( filepath_or_buffer, table, job_config=job_config ) - load_job.result() # Wait for the job to complete + self._start_generic_job(load_job) # The BigQuery REST API for tables.get doesn't take a session ID, so we # can't get the schema for a temp table that way. @@ -622,12 +637,9 @@ def _read_ibis( index_cols: Sequence[ibis_types.Value], index_labels: Sequence[Optional[str]], column_keys: Sequence[str], - ordering: Optional[core.ExpressionOrdering] = None, - ): + ordering: core.ExpressionOrdering, + ) -> dataframe.DataFrame: """Turns a table expression (plus index column) into a DataFrame.""" - hidden_ordering_columns = None - if ordering is not None and ordering.ordering_id is not None: - hidden_ordering_columns = (table_expression[ordering.ordering_id],) columns = list(index_cols) for key in column_keys: @@ -635,6 +647,12 @@ def _read_ibis( raise ValueError(f"Column '{key}' not found in this table.") columns.append(table_expression[key]) + non_hidden_ids = [col.get_name() for col in columns] + hidden_ordering_columns = [] + for ref in ordering.all_ordering_columns: + if ref.column_id not in non_hidden_ids: + hidden_ordering_columns.append(table_expression[ref.column_id]) + block = blocks.Block( core.ArrayValue( self, table_expression, columns, hidden_ordering_columns, ordering @@ -646,16 +664,19 @@ def _read_ibis( return dataframe.DataFrame(block) def read_gbq_model(self, model_name: str): - """Loads a BQML model from Google BigQuery. + """Loads a BigQuery ML model from BigQuery. Args: - model_name : the model's name in BigQuery in the format - `project_id.dataset_id.model_id`, or just `dataset_id.model_id` - to load from the default project. + model_name (str): + the model's name in BigQuery in the format + `project_id.dataset_id.model_id`, or just `dataset_id.model_id` + to load from the default project. Returns: A bigframes.ml Model wrapping the model. """ + import bigframes.ml.loader + model_ref = bigquery.ModelReference.from_string( model_name, default_project=self.bqclient.project ) @@ -663,16 +684,17 @@ def read_gbq_model(self, model_name: str): return bigframes.ml.loader.from_bq(self, model) def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame: - """Loads DataFrame from a Pandas DataFrame. + """Loads DataFrame from a pandas DataFrame. - The Pandas DataFrame will be persisted as a temporary BigQuery table, which can be + The pandas DataFrame will be persisted as a temporary BigQuery table, which can be automatically recycled after the Session is closed. Args: - pandas_dataframe: a Pandas DataFrame object to be loaded. + pandas_dataframe (pandas.DataFrame): + a pandas DataFrame object to be loaded. Returns: - A BigQuery DataFrames. + bigframes.dataframe.DataFrame: The BigQuery DataFrame. """ # Add order column to pandas DataFrame to preserve order in BigQuery ordering_col = "rowid" @@ -686,7 +708,7 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame pandas_dataframe_copy[ordering_col] = np.arange(pandas_dataframe_copy.shape[0]) # Specify the datetime dtypes, which is auto-detected as timestamp types. - schema = [] + schema: list[bigquery.SchemaField] = [] for column, dtype in zip(pandas_dataframe.columns, pandas_dataframe.dtypes): if dtype == "timestamp[us][pyarrow]": schema.append( @@ -699,11 +721,13 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame filter(lambda name: name is not None, pandas_dataframe_copy.index.names) ) index_labels = typing.cast(List[Optional[str]], index_cols) - cluster_cols = index_cols + [ordering_col] + + # Clustering probably not needed anyways as pandas tables are small + cluster_cols = [ordering_col] if len(index_cols) == 0: - index_cols = [ordering_col] - index_labels = [None] + # Block constructor will implicitly build default index + pass job_config = bigquery.LoadJobConfig(schema=schema) job_config.clustering_fields = cluster_cols @@ -718,10 +742,12 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame load_table_destination, job_config=job_config, ) - load_job.result() # Wait for the job to complete + self._start_generic_job(load_job) ordering = core.ExpressionOrdering( - ordering_id_column=OrderingColumnReference(ordering_col), is_sequential=True + ordering_value_columns=[OrderingColumnReference(ordering_col)], + total_ordering_columns=frozenset([ordering_col]), + integer_encoding=IntegerEncoding(True, is_sequential=True), ) table_expression = self.ibis_client.sql( f"SELECT * FROM `{load_table_destination.table_id}`" @@ -731,6 +757,7 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame table_expression=table_expression, index_cols=index_cols, index_labels=index_labels, + hidden_cols=(ordering_col,), ordering=ordering, is_total_ordering=True, ) @@ -771,14 +798,16 @@ def read_csv( if any(param is not None for param in (dtype, names)): not_supported = ("dtype", "names") raise NotImplementedError( - f"BigQuery engine does not support these arguments: {not_supported}" + f"BigQuery engine does not support these arguments: {not_supported}. " + f"{constants.FEEDBACK_LINK}" ) if index_col is not None and ( not index_col or not isinstance(index_col, str) ): raise NotImplementedError( - "BigQuery engine only supports a single column name for `index_col`." + "BigQuery engine only supports a single column name for `index_col`. " + f"{constants.FEEDBACK_LINK}" ) # None value for index_col cannot be passed to read_gbq @@ -794,13 +823,15 @@ def read_csv( col_order = tuple(col for col in usecols) else: raise NotImplementedError( - "BigQuery engine only supports an iterable of strings for `usecols`." + "BigQuery engine only supports an iterable of strings for `usecols`. " + f"{constants.FEEDBACK_LINK}" ) valid_encodings = {"UTF-8", "ISO-8859-1"} if encoding is not None and encoding not in valid_encodings: raise NotImplementedError( - f"BigQuery engine only supports the following encodings: {valid_encodings}" + f"BigQuery engine only supports the following encodings: {valid_encodings}. " + f"{constants.FEEDBACK_LINK}" ) job_config = bigquery.LoadJobConfig() @@ -830,7 +861,8 @@ def read_csv( else: if any(arg in kwargs for arg in ("chunksize", "iterator")): raise NotImplementedError( - "'chunksize' and 'iterator' arguments are not supported." + "'chunksize' and 'iterator' arguments are not supported. " + f"{constants.FEEDBACK_LINK}" ) if isinstance(filepath_or_buffer, str): @@ -904,20 +936,32 @@ def _create_sequential_ordering( ibis.row_number().cast(ibis_dtypes.int64).name(default_ordering_name) ) table = table.mutate(**{default_ordering_name: default_ordering_col}) + clusterable_index_cols = [ + col for col in index_cols if _can_cluster(table[col].type()) + ] + cluster_cols = (clusterable_index_cols + [default_ordering_name])[ + :_MAX_CLUSTER_COLUMNS + ] table_ref = self._query_to_session_table( self.ibis_client.compile(table), - cluster_cols=list(index_cols) + [default_ordering_name], + cluster_cols=cluster_cols, ) table = self.ibis_client.sql(f"SELECT * FROM `{table_ref.table_id}`") ordering_reference = core.OrderingColumnReference(default_ordering_name) ordering = core.ExpressionOrdering( - ordering_id_column=ordering_reference, is_sequential=True + ordering_value_columns=[ordering_reference], + total_ordering_columns=frozenset([default_ordering_name]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), ) return table, ordering def _query_to_session_table( self, query_text: str, cluster_cols: Iterable[str] ) -> bigquery.TableReference: + if len(list(cluster_cols)) > _MAX_CLUSTER_COLUMNS: + raise ValueError( + f"Too many cluster columns: {list(cluster_cols)}, max {_MAX_CLUSTER_COLUMNS} allowed." + ) # Can't set a table in _SESSION as destination via query job API, so we # run DDL, instead. table = self._create_session_table() @@ -929,9 +973,8 @@ def _query_to_session_table( CLUSTER BY {cluster_cols_sql} AS {query_text} """ - query_job = self.bqclient.query(ddl_text) try: - query_job.result() # Wait for the job to complete + self._start_query(ddl_text) # Wait for the job to complete except google.api_core.exceptions.Conflict: # Allow query retry to succeed. pass @@ -947,21 +990,58 @@ def remote_function( ): """Decorator to turn a user defined function into a BigQuery remote function. + .. note:: + Please make sure following is setup before using this API: + + 1. Have the below APIs enabled for your project: + + * BigQuery Connection API + * Cloud Functions API + * Cloud Run API + * Cloud Build API + * Artifact Registry API + * Cloud Resource Manager API + + This can be done from the cloud console (change `PROJECT_ID` to yours): + https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID + + Or from the gcloud CLI: + + `$ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com` + + 2. Have following IAM roles enabled for you: + + * BigQuery Data Editor (roles/bigquery.dataEditor) + * BigQuery Connection Admin (roles/bigquery.connectionAdmin) + * Cloud Functions Developer (roles/cloudfunctions.developer) + * Service Account User (roles/iam.serviceAccountUser) + * Storage Object Viewer (roles/storage.objectViewer) + * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) + + 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: + + 1. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection + 2. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + + Alternatively, the IAM could also be setup via the gcloud CLI: + + `$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`. + Args: input_types (list(type)): List of input data types in the user defined function. output_type (type): Data type of the output in the user defined function. dataset (str, Optional): - Dataset to use to create a BigQuery function. It should be in + Dataset in which to create a BigQuery remote function. It should be in `.` or `` format. If this - param is not provided then session dataset id would be used. + parameter is not provided then session dataset id is used. bigquery_connection (str, Optional): - Name of the BigQuery connection. If it is pre created in the same - location as the `bigquery_client.location` then it would be used, - otherwise it would be created dynamically assuming the user has - necessary priviliges. If this param is not provided then the - bigquery connection from the session would be used. + Name of the BigQuery connection. You should either have the + connection already created in the `location` you have chosen, or + you should have the Project IAM Admin role to enable the service + to create the connection for you if you need it.If this parameter is + not provided then the BigQuery connection from the session is used. reuse (bool, Optional): Reuse the remote function if already exists. `True` by default, which will result in reusing an existing remote @@ -969,38 +1049,14 @@ def remote_function( Setting it to false would force creating a unique remote function. If the required remote function does not exist then it would be created irrespective of this param. + Returns: + callable: A remote function object pointing to the cloud assets created + in the background to support the remote execution. The cloud assets can be + located through the following properties set in the object: - Notes: - Please make sure following is setup before using this API: - - 1. Have the below APIs enabled for your project: - a. BigQuery Connection API - b. Cloud Functions API - c. Cloud Run API - d. Cloud Build API - e. Artifact Registry API - f. Cloud Resource Manager API - - This can be done from the cloud console (change PROJECT_ID to yours): - https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID - Or from the gcloud CLI: - $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com - - 2. Have following IAM roles enabled for you: - a. BigQuery Data Editor (roles/bigquery.dataEditor) - b. BigQuery Connection Admin (roles/bigquery.connectionAdmin) - c. Cloud Functions Developer (roles/cloudfunctions.developer) - d. Service Account User (roles/iam.serviceAccountUser) - e. Storage Object Viewer (roles/storage.objectViewer) - f. Project IAM Admin (roles/resourcemanager.projectIamAdmin) - (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) - - 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set: - a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection - b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function - Alternatively, the IAM could also be setup via the gcloud CLI: - $ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker" + `bigframes_cloud_function` - The google cloud function deployed for the user defined code. + `bigframes_remote_function` - The bigquery remote function capable of calling into `bigframes_cloud_function`. """ return bigframes_rf( input_types, @@ -1011,34 +1067,62 @@ def remote_function( reuse=reuse, ) + def read_gbq_function( + self, + function_name: str, + ): + """Loads a BigQuery function from BigQuery. + + Then it can be applied to a DataFrame or Series. + + Args: + function_name (str): + the function's name in BigQuery in the format + `project_id.dataset_id.function_name`, or + `dataset_id.function_name` to load from the default project, or + `function_name` to load from the default project and the dataset + associated with the current session. + + Returns: + callable: A function object pointing to the BigQuery function read + from BigQuery. + + The object is similar to the one created by the `remote_function` + decorator, including the `bigframes_remote_function` property, but + not including the `bigframes_cloud_function` property. + """ + + return bigframes_rgf( + function_name=function_name, + session=self, + ) + def _start_query( self, sql: str, job_config: Optional[bigquery.job.QueryJobConfig] = None, max_results: Optional[int] = None, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: + """ + Starts query job and waits for results + """ if job_config is not None: query_job = self.bqclient.query(sql, job_config=job_config) else: query_job = self.bqclient.query(sql) opts = bigframes.options.display - if opts.progress_bar is not None: - results_iterator = formatting_helpers.wait_for_job( + if opts.progress_bar is not None and not query_job.configuration.dry_run: + results_iterator = formatting_helpers.wait_for_query_job( query_job, max_results, opts.progress_bar ) else: results_iterator = query_job.result(max_results=max_results) return results_iterator, query_job - def _extract_table(self, source_table, destination_uris, job_config): - extract_job = self.bqclient.extract_table( - source=source_table, - destination_uris=destination_uris, - job_config=job_config, - ) - extract_job.result() - return extract_job + def _get_table_size(self, destination_table): + table = self.bqclient.get_table(destination_table) + return table.num_bytes def _rows_to_dataframe( self, row_iterator: bigquery.table.RowIterator @@ -1050,6 +1134,27 @@ def _rows_to_dataframe( string_dtype=pandas.StringDtype(storage="pyarrow"), ) + def _start_generic_job(self, job: formatting_helpers.GenericJob): + if bigframes.options.display.progress_bar is not None: + formatting_helpers.wait_for_job( + job, bigframes.options.display.progress_bar + ) # Wait for the job to complete + else: + job.result() + def connect(context: Optional[bigquery_options.BigQueryOptions] = None) -> Session: return Session(context) + + +def _can_cluster(ibis_type: ibis_dtypes.DataType): + # https://cloud.google.com/bigquery/docs/clustered-tables + # Notably, float is excluded + return ( + ibis_type.is_integer() + or ibis_type.is_string() + or ibis_type.is_decimal() + or ibis_type.is_date() + or ibis_type.is_timestamp() + or ibis_type.is_boolean() + ) diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst deleted file mode 100644 index cde2a70a3fb..00000000000 --- a/docs/getting_started/index.rst +++ /dev/null @@ -1,27 +0,0 @@ -Getting started -=============== - -Prerequisites --------------- - -* Install the ``bigframes`` package. -* Create a GCP project and billing account. -* When running locally, authenticate with application default credentials. See - the `gcloud auth application-default login - `_ - reference. - -Code sample ------------ - -.. literalinclude:: ../samples/snippets/quickstart.py - :language: python - :dedent: 4 - :start-after: [START bigquery_bigframes_quickstart] - :end-before: [END bigquery_bigframes_quickstart] - -Learn more ----------- - -* See :ref:`user_guide` for more samples. -* See :ref:`bigframes_ml` for more ML samples. diff --git a/docs/index.rst b/docs/index.rst index eb721fbf293..ff1cd09eb7d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,9 +1,11 @@ .. include:: README.rst +API reference +------------- + .. toctree:: :maxdepth: 3 - getting_started/index user_guide/index reference/index diff --git a/docs/reference/bigframes.ml/cluster.rst b/docs/reference/bigframes.ml/cluster.rst index 5456b728b67..e91a28c0511 100644 --- a/docs/reference/bigframes.ml/cluster.rst +++ b/docs/reference/bigframes.ml/cluster.rst @@ -3,3 +3,5 @@ bigframes.ml.cluster .. automodule:: bigframes.ml.cluster :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/compose.rst b/docs/reference/bigframes.ml/compose.rst index e41f17a33b1..9992728362f 100644 --- a/docs/reference/bigframes.ml/compose.rst +++ b/docs/reference/bigframes.ml/compose.rst @@ -3,3 +3,5 @@ bigframes.ml.compose .. automodule:: bigframes.ml.compose :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/decomposition.rst b/docs/reference/bigframes.ml/decomposition.rst index 122dbc2b168..ec804ac8cdc 100644 --- a/docs/reference/bigframes.ml/decomposition.rst +++ b/docs/reference/bigframes.ml/decomposition.rst @@ -3,3 +3,5 @@ bigframes.ml.decomposition .. automodule:: bigframes.ml.decomposition :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/ensemble.rst b/docs/reference/bigframes.ml/ensemble.rst new file mode 100644 index 00000000000..2652ab5aa4d --- /dev/null +++ b/docs/reference/bigframes.ml/ensemble.rst @@ -0,0 +1,7 @@ +bigframes.ml.ensemble +===================== + +.. automodule:: bigframes.ml.ensemble + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/forecasting.rst b/docs/reference/bigframes.ml/forecasting.rst index 037bd7e82ef..04015c99117 100644 --- a/docs/reference/bigframes.ml/forecasting.rst +++ b/docs/reference/bigframes.ml/forecasting.rst @@ -3,3 +3,5 @@ bigframes.ml.forecasting .. automodule:: bigframes.ml.forecasting :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/imported.rst b/docs/reference/bigframes.ml/imported.rst index 54120534c02..c151cbda6f1 100644 --- a/docs/reference/bigframes.ml/imported.rst +++ b/docs/reference/bigframes.ml/imported.rst @@ -3,3 +3,5 @@ bigframes.ml.imported .. automodule:: bigframes.ml.imported :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/index.rst b/docs/reference/bigframes.ml/index.rst index e17edc2ee98..f3cbe1174a7 100644 --- a/docs/reference/bigframes.ml/index.rst +++ b/docs/reference/bigframes.ml/index.rst @@ -13,6 +13,8 @@ API Reference decomposition + ensemble + forecasting imported diff --git a/docs/reference/bigframes.ml/linear_model.rst b/docs/reference/bigframes.ml/linear_model.rst index 5936858e64b..8c6c2765b12 100644 --- a/docs/reference/bigframes.ml/linear_model.rst +++ b/docs/reference/bigframes.ml/linear_model.rst @@ -3,3 +3,5 @@ bigframes.ml.linear_model .. automodule:: bigframes.ml.linear_model :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/llm.rst b/docs/reference/bigframes.ml/llm.rst index 4debef74ffb..20ae7793e73 100644 --- a/docs/reference/bigframes.ml/llm.rst +++ b/docs/reference/bigframes.ml/llm.rst @@ -3,3 +3,5 @@ bigframes.ml.llm .. automodule:: bigframes.ml.llm :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/metrics.rst b/docs/reference/bigframes.ml/metrics.rst index dc951ebed42..aca11f7e9fc 100644 --- a/docs/reference/bigframes.ml/metrics.rst +++ b/docs/reference/bigframes.ml/metrics.rst @@ -3,3 +3,5 @@ bigframes.ml.metrics .. automodule:: bigframes.ml.metrics :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/model_selection.rst b/docs/reference/bigframes.ml/model_selection.rst index b0d4e5ed3ee..d662285f990 100644 --- a/docs/reference/bigframes.ml/model_selection.rst +++ b/docs/reference/bigframes.ml/model_selection.rst @@ -3,3 +3,5 @@ bigframes.ml.model_selection .. automodule:: bigframes.ml.model_selection :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/pipeline.rst b/docs/reference/bigframes.ml/pipeline.rst index 239ee59c431..22e877dc5b3 100644 --- a/docs/reference/bigframes.ml/pipeline.rst +++ b/docs/reference/bigframes.ml/pipeline.rst @@ -3,3 +3,5 @@ bigframes.ml.pipeline .. automodule:: bigframes.ml.pipeline :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/preprocessing.rst b/docs/reference/bigframes.ml/preprocessing.rst index 4d2279fc8ea..eac72da1730 100644 --- a/docs/reference/bigframes.ml/preprocessing.rst +++ b/docs/reference/bigframes.ml/preprocessing.rst @@ -3,3 +3,5 @@ bigframes.ml.preprocessing .. automodule:: bigframes.ml.preprocessing :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes/options.rst b/docs/reference/bigframes/options.rst index 914e4976d54..d831a519fe6 100644 --- a/docs/reference/bigframes/options.rst +++ b/docs/reference/bigframes/options.rst @@ -10,3 +10,5 @@ Options and settings .. autoclass:: bigframes._config.bigquery_options.BigQueryOptions .. autoclass:: bigframes._config.display_options.DisplayOptions + +.. autoclass:: bigframes._config.sampling_options.SamplingOptions diff --git a/docs/user_guide/bigframes.pandas/remote_functions.rst b/docs/user_guide/bigframes.pandas/remote_functions.rst index 3a1a9d342f0..abfe6a10caa 100644 --- a/docs/user_guide/bigframes.pandas/remote_functions.rst +++ b/docs/user_guide/bigframes.pandas/remote_functions.rst @@ -9,32 +9,125 @@ find more details on it via `help` command. .. code-block:: python - import bigframes.pandas as pd - help(pd.remote_function) + import bigframes.pandas as bpd + help(bpd.remote_function) Read a table and inspect the column of interest. .. code-block:: python - df = pd.read_gbq("bigquery-public-data.ml_datasets.penguins") + 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. +APIs and IAM permissions are setup in the project. In our examples we would be +using a pre-created connection named `bigframes-rf-conn`. 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. .. code-block:: python - @pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn') + @bpd.remote_function([float], str, bigquery_connection='bigframes-rf-conn') def get_bucket(num): if not num: return "NA" boundary = 4000 return "at_or_above_4000" if num >= boundary else "below_4000" -Run the custom function on the column of interest to create a new column. +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. .. code-block:: python 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. + +.. code-block:: python + df[['body_mass_g', 'body_mass_bucket']].head(10) + +This 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. + +.. warning:: + The deployed cloud function may be visible to other users with sufficient + privilege in the project. The user should be careful about having any + sensitive data in the code that will be deployed as a remote function. + +The cloud function can be located from a property set in the remote function object. + +.. code-block:: python + + get_bucket.bigframes_cloud_function + +and then there is a BigQuery remote function created configured to call into the +cloud function via the BigQuery connection. That can also be located from +another property set in the remote function object. + +.. code-block:: python + + get_bucket.bigframes_remote_function + +The cloud assets created are persistant and the user can manage them directy +from the Google Cloud Console. + +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 decorated it as a remote function: + +.. code-block:: python + + @bpd.remote_function([str], str, bigquery_connection='bigframes-rf-conn') + def get_hash(input): + import hashlib + # handle missing value + if input is None: + input = "" + encoded_input = input.encode() + hash = hashlib.md5(encoded_input) + return hash.hexdigest() + +We can use this remote function in another `pandas`-like API `map` that can be +applied on a DataFrame: + +.. code-block:: python + + df_redacted = df[["species", "island", "sex"]].map(get_hash) + df_redacted.head(10). + +Using Existing Functions +======================== + +If you have already defined a custom function in BigQuery, either in the +BigQuery Studio or with the `remote_function` decorator above or otherwise, you +may use it with BigQuery DataFrames with the `read_gbq_function` method. + +More details are available via the `help` command: + +.. code-block:: python + + import bigframes.pandas as pd + help(pd.read_gbq_function) + +Here is an example of using `read_gbq_function` to load an existing function +named `get_bucket`: + +.. code-block:: python + + import bigframes.pandas as pd + + df = pd.read_gbq("bigquery-public-data.ml_datasets.penguins") + get_bucket = pd.read_gbq_function("get_bucket") + + df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket)) + df.head(10) + +Note: As mentioned above, if a function is created using the `remote_function` +decorator, its generated name (including project and dataset) is accessible +immediately afterward in the function's `bigframes_remote_function` attribute. +The same string can be passed to `read_gbq_function` later in another context. diff --git a/notebooks/01 - Getting Started.ipynb b/notebooks/01 - Getting Started.ipynb index 132cdf8db0d..473bdd8ceae 100644 --- a/notebooks/01 - Getting Started.ipynb +++ b/notebooks/01 - Getting Started.ipynb @@ -28,7 +28,22 @@ "metadata": { "tags": [] }, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "5c480b60490940d3a45fa6b9ca2cecdb", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 411d90c8-8b22-40b0-ad42-04f9e38c074e is DONE. 0 Bytes processed. body_mass_g\n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", "
sex
tag_number
11720Dream32.115.5188.03050.036.618.4184.03475.0FEMALE
1371Biscoe37.716.0183.03075.0FEMALE1Dream39.819.1184.04650.0MALE
1417Torgersen38.617.0188.02900.02Dream40.918.9184.03900.0MALE
4Dream37.316.8192.03000.0FEMALE
12045Dream43.218.5192.04100.0MALE
9Dream40.220.1200.03975.0MALE
10Dream40.717.0190.03725.040.818.9208.04300.0MALE
1251Biscoe37.617.011Dream39.018.7185.03600.03650.0MALE
12Dream37.016.9185.03000.0FEMALE
1422Torgersen35.717.0189.03350.014Dream34.017.1185.03400.0FEMALE
1394Torgersen40.217.0176.03450.015Dream37.016.5185.03400.0FEMALE
116318Dream39.717.9193.04250.0MALE
19Dream37.818.1193.03750.0MALE
22Dream36.417.0195.03325.040.217.1193.03400.0FEMALE
1329Biscoe38.117.0181.03175.023Dream36.818.5193.03500.0FEMALE
1406Torgersen44.118.0210.026Dream41.518.5201.04000.0MALE
119631Dream36.518.0182.03150.033.116.1178.02900.0FEMALE
1228Biscoe41.618.0192.03950.032Dream37.218.1178.03900.0MALE
1412Torgersen40.318.0195.033Dream39.516.7178.03250.0FEMALE
114235Dream35.718.0202.03550.036.018.5186.03100.0FEMALE
1430Torgersen33.519.0190.03600.0FEMALE36Dream39.618.1186.04450.0MALE
1333Biscoe43.219.0197.04775.038Dream41.320.3194.03550.0MALE
1414Torgersen38.719.0195.03450.041Dream35.718.0202.03550.0FEMALE
119751Dream41.119.0182.038.117.6187.03425.0MALE
1443Torgersen40.619.0199.04000.0MALEFEMALE
1295Biscoe41.020.0203.04725.0MALE53Dream36.017.1187.03700.0FEMALE
\n", + "

25 rows × 6 columns

\n", "[146 rows x 6 columns in total]" ], "text/plain": [ - " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", - "tag_number \n", - "1172 Dream 32.1 15.5 188.0 \n", - "1371 Biscoe 37.7 16.0 183.0 \n", - "1417 Torgersen 38.6 17.0 188.0 \n", - "1204 Dream 40.7 17.0 190.0 \n", - "1251 Biscoe 37.6 17.0 185.0 \n", - "1422 Torgersen 35.7 17.0 189.0 \n", - "1394 Torgersen 40.2 17.0 176.0 \n", - "1163 Dream 36.4 17.0 195.0 \n", - "1329 Biscoe 38.1 17.0 181.0 \n", - "1406 Torgersen 44.1 18.0 210.0 \n", - "1196 Dream 36.5 18.0 182.0 \n", - "1228 Biscoe 41.6 18.0 192.0 \n", - "1412 Torgersen 40.3 18.0 195.0 \n", - "1142 Dream 35.7 18.0 202.0 \n", - "1430 Torgersen 33.5 19.0 190.0 \n", - "1333 Biscoe 43.2 19.0 197.0 \n", - "1414 Torgersen 38.7 19.0 195.0 \n", - "1197 Dream 41.1 19.0 182.0 \n", - "1443 Torgersen 40.6 19.0 199.0 \n", - "1295 Biscoe 41.0 20.0 203.0 \n", - "1207 Dream 38.8 20.0 190.0 \n", - "1349 Biscoe 38.2 20.0 190.0 \n", - "1350 Biscoe 37.8 20.0 190.0 \n", - "1351 Biscoe 38.1 16.5 198.0 \n", - "1116 Dream 37.0 16.5 185.0 \n", + " island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", + "0 Dream 36.6 18.4 184.0 3475.0 \n", + "1 Dream 39.8 19.1 184.0 4650.0 \n", + "2 Dream 40.9 18.9 184.0 3900.0 \n", + "4 Dream 37.3 16.8 192.0 3000.0 \n", + "5 Dream 43.2 18.5 192.0 4100.0 \n", + "9 Dream 40.2 20.1 200.0 3975.0 \n", + "10 Dream 40.8 18.9 208.0 4300.0 \n", + "11 Dream 39.0 18.7 185.0 3650.0 \n", + "12 Dream 37.0 16.9 185.0 3000.0 \n", + "14 Dream 34.0 17.1 185.0 3400.0 \n", + "15 Dream 37.0 16.5 185.0 3400.0 \n", + "18 Dream 39.7 17.9 193.0 4250.0 \n", + "19 Dream 37.8 18.1 193.0 3750.0 \n", + "22 Dream 40.2 17.1 193.0 3400.0 \n", + "23 Dream 36.8 18.5 193.0 3500.0 \n", + "26 Dream 41.5 18.5 201.0 4000.0 \n", + "31 Dream 33.1 16.1 178.0 2900.0 \n", + "32 Dream 37.2 18.1 178.0 3900.0 \n", + "33 Dream 39.5 16.7 178.0 3250.0 \n", + "35 Dream 36.0 18.5 186.0 3100.0 \n", + "36 Dream 39.6 18.1 186.0 4450.0 \n", + "38 Dream 41.3 20.3 194.0 3550.0 \n", + "41 Dream 35.7 18.0 202.0 3550.0 \n", + "51 Dream 38.1 17.6 187.0 3425.0 \n", + "53 Dream 36.0 17.1 187.0 3700.0 \n", "\n", - " body_mass_g sex \n", - "tag_number \n", - "1172 3050.0 FEMALE \n", - "1371 3075.0 FEMALE \n", - "1417 2900.0 FEMALE \n", - "1204 3725.0 MALE \n", - "1251 3600.0 FEMALE \n", - "1422 3350.0 FEMALE \n", - "1394 3450.0 FEMALE \n", - "1163 3325.0 FEMALE \n", - "1329 3175.0 FEMALE \n", - "1406 4000.0 MALE \n", - "1196 3150.0 FEMALE \n", - "1228 3950.0 MALE \n", - "1412 3250.0 FEMALE \n", - "1142 3550.0 FEMALE \n", - "1430 3600.0 FEMALE \n", - "1333 4775.0 MALE \n", - "1414 3450.0 FEMALE \n", - "1197 3425.0 MALE \n", - "1443 4000.0 MALE \n", - "1295 4725.0 MALE \n", - "1207 3950.0 MALE \n", - "1349 3900.0 MALE \n", - "1350 4250.0 MALE \n", - "1351 3825.0 FEMALE \n", - "1116 3400.0 FEMALE \n", + " sex \n", + "0 FEMALE \n", + "1 MALE \n", + "2 MALE \n", + "4 FEMALE \n", + "5 MALE \n", + "9 MALE \n", + "10 MALE \n", + "11 MALE \n", + "12 FEMALE \n", + "14 FEMALE \n", + "15 FEMALE \n", + "18 MALE \n", + "19 MALE \n", + "22 FEMALE \n", + "23 FEMALE \n", + "26 MALE \n", + "31 FEMALE \n", + "32 MALE \n", + "33 FEMALE \n", + "35 FEMALE \n", + "36 MALE \n", + "38 MALE \n", + "41 FEMALE \n", + "51 FEMALE \n", + "53 FEMALE \n", "...\n", "\n", "[146 rows x 6 columns]" ] }, - "execution_count": 2, + "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "# set a friendlier index to uniquely identify the rows\n", - "df = df.set_index(\"tag_number\")\n", - "\n", "# filter down to the data we want to analyze\n", "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", "\n", @@ -683,9 +843,24 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 14, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "81f9aa34c7234bd88b6b7a4bc77d4b4e", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 288f0daa-a51e-45b4-86bf-d054467c4a99 is DONE. 28.9 kB processed. \n", " \n", " \n", - " tag_number\n", " species\n", " island\n", " culmen_length_mm\n", @@ -65,288 +120,319 @@ " \n", " \n", " 0\n", - " 1225\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " <NA>\n", - " <NA>\n", - " <NA>\n", - " <NA>\n", - " <NA>\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 36.6\n", + " 18.4\n", + " 184.0\n", + " 3475.0\n", + " FEMALE\n", " \n", " \n", " 1\n", - " 1278\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 42.0\n", - " 13.5\n", - " 210.0\n", - " 4150.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 39.8\n", + " 19.1\n", + " 184.0\n", + " 4650.0\n", + " MALE\n", " \n", " \n", " 2\n", - " 1275\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 46.5\n", - " 13.5\n", - " 210.0\n", - " 4550.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 40.9\n", + " 18.9\n", + " 184.0\n", + " 3900.0\n", + " MALE\n", " \n", " \n", " 3\n", - " 1233\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 43.3\n", - " 14.0\n", - " 208.0\n", - " 4575.0\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 46.5\n", + " 17.9\n", + " 192.0\n", + " 3500.0\n", " FEMALE\n", " \n", " \n", " 4\n", - " 1311\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 47.5\n", - " 14.0\n", - " 212.0\n", - " 4875.0\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 37.3\n", + " 16.8\n", + " 192.0\n", + " 3000.0\n", " FEMALE\n", " \n", " \n", " 5\n", - " 1316\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 49.1\n", - " 14.5\n", - " 212.0\n", - " 4625.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 43.2\n", + " 18.5\n", + " 192.0\n", + " 4100.0\n", + " MALE\n", " \n", " \n", " 6\n", - " 1313\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 45.5\n", - " 14.5\n", - " 212.0\n", - " 4750.0\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 46.9\n", + " 16.6\n", + " 192.0\n", + " 2700.0\n", " FEMALE\n", " \n", " \n", " 7\n", - " 1381\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 47.6\n", - " 14.5\n", - " 215.0\n", - " 5400.0\n", - " MALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 50.5\n", + " 18.4\n", + " 200.0\n", + " 3400.0\n", + " FEMALE\n", " \n", " \n", " 8\n", - " 1377\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 45.1\n", - " 14.5\n", - " 207.0\n", - " 5050.0\n", - " FEMALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 49.5\n", + " 19.0\n", + " 200.0\n", + " 3800.0\n", + " MALE\n", " \n", " \n", " 9\n", - " 1380\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 45.1\n", - " 14.5\n", - " 215.0\n", - " 5000.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 40.2\n", + " 20.1\n", + " 200.0\n", + " 3975.0\n", + " MALE\n", " \n", " \n", " 10\n", - " 1257\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 46.2\n", - " 14.5\n", - " 209.0\n", - " 4800.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 40.8\n", + " 18.9\n", + " 208.0\n", + " 4300.0\n", + " MALE\n", " \n", " \n", " 11\n", - " 1336\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 46.5\n", - " 14.5\n", - " 213.0\n", - " 4400.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 39.0\n", + " 18.7\n", + " 185.0\n", + " 3650.0\n", + " MALE\n", " \n", " \n", " 12\n", - " 1237\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 43.2\n", - " 14.5\n", - " 208.0\n", - " 4450.0\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 37.0\n", + " 16.9\n", + " 185.0\n", + " 3000.0\n", " FEMALE\n", " \n", " \n", " 13\n", - " 1302\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 48.5\n", - " 15.0\n", - " 219.0\n", - " 4850.0\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 47.0\n", + " 17.3\n", + " 185.0\n", + " 3700.0\n", " FEMALE\n", " \n", " \n", " 14\n", - " 1325\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 49.1\n", - " 15.0\n", - " 228.0\n", - " 5500.0\n", - " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 34.0\n", + " 17.1\n", + " 185.0\n", + " 3400.0\n", + " FEMALE\n", " \n", " \n", " 15\n", - " 1285\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 47.5\n", - " 15.0\n", - " 218.0\n", - " 4950.0\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 37.0\n", + " 16.5\n", + " 185.0\n", + " 3400.0\n", " FEMALE\n", " \n", " \n", " 16\n", - " 1242\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 49.6\n", - " 15.0\n", - " 216.0\n", - " 4750.0\n", - " MALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 45.7\n", + " 17.3\n", + " 193.0\n", + " 3600.0\n", + " FEMALE\n", " \n", " \n", " 17\n", - " 1246\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 47.7\n", - " 15.0\n", - " 216.0\n", - " 4750.0\n", - " FEMALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 50.6\n", + " 19.4\n", + " 193.0\n", + " 3800.0\n", + " MALE\n", " \n", " \n", " 18\n", - " 1320\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 45.5\n", - " 15.0\n", - " 220.0\n", - " 5000.0\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 39.7\n", + " 17.9\n", + " 193.0\n", + " 4250.0\n", " MALE\n", " \n", " \n", " 19\n", - " 1244\n", - " Gentoo penguin (Pygoscelis papua)\n", - " Biscoe\n", - " 46.4\n", - " 15.0\n", - " 216.0\n", - " 4700.0\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 37.8\n", + " 18.1\n", + " 193.0\n", + " 3750.0\n", + " MALE\n", + " \n", + " \n", + " 20\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 46.6\n", + " 17.8\n", + " 193.0\n", + " 3800.0\n", + " FEMALE\n", + " \n", + " \n", + " 21\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 51.3\n", + " 19.2\n", + " 193.0\n", + " 3650.0\n", + " MALE\n", + " \n", + " \n", + " 22\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 40.2\n", + " 17.1\n", + " 193.0\n", + " 3400.0\n", + " FEMALE\n", + " \n", + " \n", + " 23\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 36.8\n", + " 18.5\n", + " 193.0\n", + " 3500.0\n", " FEMALE\n", " \n", + " \n", + " 24\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " Dream\n", + " 49.6\n", + " 18.2\n", + " 193.0\n", + " 3775.0\n", + " MALE\n", + " \n", " \n", "\n", - "[347 rows x 8 columns in total]" + "

25 rows × 7 columns

\n", + "[344 rows x 7 columns in total]" ], "text/plain": [ - " tag_number species island culmen_length_mm \\\n", - "0 1225 Gentoo penguin (Pygoscelis papua) Biscoe \n", - "1 1278 Gentoo penguin (Pygoscelis papua) Biscoe 42.0 \n", - "2 1275 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", - "3 1233 Gentoo penguin (Pygoscelis papua) Biscoe 43.3 \n", - "4 1311 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", - "5 1316 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", - "6 1313 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", - "7 1381 Gentoo penguin (Pygoscelis papua) Biscoe 47.6 \n", - "8 1377 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", - "9 1380 Gentoo penguin (Pygoscelis papua) Biscoe 45.1 \n", - "10 1257 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 \n", - "11 1336 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", - "12 1237 Gentoo penguin (Pygoscelis papua) Biscoe 43.2 \n", - "13 1302 Gentoo penguin (Pygoscelis papua) Biscoe 48.5 \n", - "14 1325 Gentoo penguin (Pygoscelis papua) Biscoe 49.1 \n", - "15 1285 Gentoo penguin (Pygoscelis papua) Biscoe 47.5 \n", - "16 1242 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", - "17 1246 Gentoo penguin (Pygoscelis papua) Biscoe 47.7 \n", - "18 1320 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", - "19 1244 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 \n", - "20 1390 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", - "21 1379 Gentoo penguin (Pygoscelis papua) Biscoe 47.8 \n", - "22 1267 Gentoo penguin (Pygoscelis papua) Biscoe 50.1 \n", - "23 1389 Gentoo penguin (Pygoscelis papua) Biscoe 47.2 \n", - "24 1269 Gentoo penguin (Pygoscelis papua) Biscoe 49.6 \n", + " species island culmen_length_mm \\\n", + "0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 \n", + "1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 \n", + "2 Adelie Penguin (Pygoscelis adeliae) Dream 40.9 \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream 46.5 \n", + "4 Adelie Penguin (Pygoscelis adeliae) Dream 37.3 \n", + "5 Adelie Penguin (Pygoscelis adeliae) Dream 43.2 \n", + "6 Chinstrap penguin (Pygoscelis antarctica) Dream 46.9 \n", + "7 Chinstrap penguin (Pygoscelis antarctica) Dream 50.5 \n", + "8 Chinstrap penguin (Pygoscelis antarctica) Dream 49.5 \n", + "9 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "10 Adelie Penguin (Pygoscelis adeliae) Dream 40.8 \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream 39.0 \n", + "12 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "13 Chinstrap penguin (Pygoscelis antarctica) Dream 47.0 \n", + "14 Adelie Penguin (Pygoscelis adeliae) Dream 34.0 \n", + "15 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", + "16 Chinstrap penguin (Pygoscelis antarctica) Dream 45.7 \n", + "17 Chinstrap penguin (Pygoscelis antarctica) Dream 50.6 \n", + "18 Adelie Penguin (Pygoscelis adeliae) Dream 39.7 \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.8 \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream 46.6 \n", + "21 Chinstrap penguin (Pygoscelis antarctica) Dream 51.3 \n", + "22 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream 36.8 \n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream 49.6 \n", "\n", " culmen_depth_mm flipper_length_mm body_mass_g sex \n", - "0 \n", - "1 13.5 210.0 4150.0 FEMALE \n", - "2 13.5 210.0 4550.0 FEMALE \n", - "3 14.0 208.0 4575.0 FEMALE \n", - "4 14.0 212.0 4875.0 FEMALE \n", - "5 14.5 212.0 4625.0 FEMALE \n", - "6 14.5 212.0 4750.0 FEMALE \n", - "7 14.5 215.0 5400.0 MALE \n", - "8 14.5 207.0 5050.0 FEMALE \n", - "9 14.5 215.0 5000.0 FEMALE \n", - "10 14.5 209.0 4800.0 FEMALE \n", - "11 14.5 213.0 4400.0 FEMALE \n", - "12 14.5 208.0 4450.0 FEMALE \n", - "13 15.0 219.0 4850.0 FEMALE \n", - "14 15.0 228.0 5500.0 MALE \n", - "15 15.0 218.0 4950.0 FEMALE \n", - "16 15.0 216.0 4750.0 MALE \n", - "17 15.0 216.0 4750.0 FEMALE \n", - "18 15.0 220.0 5000.0 MALE \n", - "19 15.0 216.0 4700.0 FEMALE \n", - "20 15.0 223.0 5550.0 MALE \n", - "21 15.0 215.0 5650.0 MALE \n", - "22 15.0 225.0 5000.0 MALE \n", - "23 15.5 215.0 4975.0 FEMALE \n", - "24 16.0 225.0 5700.0 MALE \n", + "0 18.4 184.0 3475.0 FEMALE \n", + "1 19.1 184.0 4650.0 MALE \n", + "2 18.9 184.0 3900.0 MALE \n", + "3 17.9 192.0 3500.0 FEMALE \n", + "4 16.8 192.0 3000.0 FEMALE \n", + "5 18.5 192.0 4100.0 MALE \n", + "6 16.6 192.0 2700.0 FEMALE \n", + "7 18.4 200.0 3400.0 FEMALE \n", + "8 19.0 200.0 3800.0 MALE \n", + "9 20.1 200.0 3975.0 MALE \n", + "10 18.9 208.0 4300.0 MALE \n", + "11 18.7 185.0 3650.0 MALE \n", + "12 16.9 185.0 3000.0 FEMALE \n", + "13 17.3 185.0 3700.0 FEMALE \n", + "14 17.1 185.0 3400.0 FEMALE \n", + "15 16.5 185.0 3400.0 FEMALE \n", + "16 17.3 193.0 3600.0 FEMALE \n", + "17 19.4 193.0 3800.0 MALE \n", + "18 17.9 193.0 4250.0 MALE \n", + "19 18.1 193.0 3750.0 MALE \n", + "20 17.8 193.0 3800.0 FEMALE \n", + "21 19.2 193.0 3650.0 MALE \n", + "22 17.1 193.0 3400.0 FEMALE \n", + "23 18.5 193.0 3500.0 FEMALE \n", + "24 18.2 193.0 3775.0 MALE \n", "...\n", "\n", - "[347 rows x 8 columns]" + "[344 rows x 7 columns]" ] }, - "execution_count": 1, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } @@ -371,9 +457,51 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 4, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "782462924bc84b9281d6d66f7f4acbe0", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 6f0e1877-369d-4e9f-81a8-9c00ab1b57b3 is DONE. 28.9 kB processed. body_mass_g\n", " sex\n", " \n", - " \n", - " tag_number\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", " \n", " \n", " \n", - " 1172\n", + " 0\n", " Dream\n", - " 32.1\n", - " 15.5\n", - " 188.0\n", - " 3050.0\n", + " 36.6\n", + " 18.4\n", + " 184.0\n", + " 3475.0\n", " FEMALE\n", " \n", " \n", - " 1371\n", - " Biscoe\n", - " 37.7\n", - " 16.0\n", - " 183.0\n", - " 3075.0\n", - " FEMALE\n", + " 1\n", + " Dream\n", + " 39.8\n", + " 19.1\n", + " 184.0\n", + " 4650.0\n", + " MALE\n", " \n", " \n", - " 1417\n", - " Torgersen\n", - " 38.6\n", - " 17.0\n", - " 188.0\n", - " 2900.0\n", + " 2\n", + " Dream\n", + " 40.9\n", + " 18.9\n", + " 184.0\n", + " 3900.0\n", + " MALE\n", + " \n", + " \n", + " 4\n", + " Dream\n", + " 37.3\n", + " 16.8\n", + " 192.0\n", + " 3000.0\n", " FEMALE\n", " \n", " \n", - " 1204\n", + " 5\n", " Dream\n", - " 40.7\n", - " 17.0\n", - " 190.0\n", - " 3725.0\n", + " 43.2\n", + " 18.5\n", + " 192.0\n", + " 4100.0\n", " MALE\n", " \n", " \n", - " 1251\n", - " Biscoe\n", - " 37.6\n", - " 17.0\n", + " 9\n", + " Dream\n", + " 40.2\n", + " 20.1\n", + " 200.0\n", + " 3975.0\n", + " MALE\n", + " \n", + " \n", + " 10\n", + " Dream\n", + " 40.8\n", + " 18.9\n", + " 208.0\n", + " 4300.0\n", + " MALE\n", + " \n", + " \n", + " 11\n", + " Dream\n", + " 39.0\n", + " 18.7\n", " 185.0\n", - " 3600.0\n", + " 3650.0\n", + " MALE\n", + " \n", + " \n", + " 12\n", + " Dream\n", + " 37.0\n", + " 16.9\n", + " 185.0\n", + " 3000.0\n", " FEMALE\n", " \n", " \n", - " 1422\n", - " Torgersen\n", - " 35.7\n", - " 17.0\n", - " 189.0\n", - " 3350.0\n", + " 14\n", + " Dream\n", + " 34.0\n", + " 17.1\n", + " 185.0\n", + " 3400.0\n", " FEMALE\n", " \n", " \n", - " 1394\n", - " Torgersen\n", - " 40.2\n", - " 17.0\n", - " 176.0\n", - " 3450.0\n", + " 15\n", + " Dream\n", + " 37.0\n", + " 16.5\n", + " 185.0\n", + " 3400.0\n", " FEMALE\n", " \n", " \n", - " 1163\n", + " 18\n", + " Dream\n", + " 39.7\n", + " 17.9\n", + " 193.0\n", + " 4250.0\n", + " MALE\n", + " \n", + " \n", + " 19\n", + " Dream\n", + " 37.8\n", + " 18.1\n", + " 193.0\n", + " 3750.0\n", + " MALE\n", + " \n", + " \n", + " 22\n", " Dream\n", - " 36.4\n", - " 17.0\n", - " 195.0\n", - " 3325.0\n", + " 40.2\n", + " 17.1\n", + " 193.0\n", + " 3400.0\n", " FEMALE\n", " \n", " \n", - " 1329\n", - " Biscoe\n", - " 38.1\n", - " 17.0\n", - " 181.0\n", - " 3175.0\n", + " 23\n", + " Dream\n", + " 36.8\n", + " 18.5\n", + " 193.0\n", + " 3500.0\n", " FEMALE\n", " \n", " \n", - " 1406\n", - " Torgersen\n", - " 44.1\n", - " 18.0\n", - " 210.0\n", + " 26\n", + " Dream\n", + " 41.5\n", + " 18.5\n", + " 201.0\n", " 4000.0\n", " MALE\n", " \n", " \n", - " 1196\n", + " 31\n", " Dream\n", - " 36.5\n", - " 18.0\n", - " 182.0\n", - " 3150.0\n", + " 33.1\n", + " 16.1\n", + " 178.0\n", + " 2900.0\n", " FEMALE\n", " \n", " \n", - " 1228\n", - " Biscoe\n", - " 41.6\n", - " 18.0\n", - " 192.0\n", - " 3950.0\n", + " 32\n", + " Dream\n", + " 37.2\n", + " 18.1\n", + " 178.0\n", + " 3900.0\n", " MALE\n", " \n", " \n", - " 1412\n", - " Torgersen\n", - " 40.3\n", - " 18.0\n", - " 195.0\n", + " 33\n", + " Dream\n", + " 39.5\n", + " 16.7\n", + " 178.0\n", " 3250.0\n", " FEMALE\n", " \n", " \n", - " 1142\n", + " 35\n", " Dream\n", - " 35.7\n", - " 18.0\n", - " 202.0\n", - " 3550.0\n", + " 36.0\n", + " 18.5\n", + " 186.0\n", + " 3100.0\n", " FEMALE\n", " \n", " \n", - " 1430\n", - " Torgersen\n", - " 33.5\n", - " 19.0\n", - " 190.0\n", - " 3600.0\n", - " FEMALE\n", + " 36\n", + " Dream\n", + " 39.6\n", + " 18.1\n", + " 186.0\n", + " 4450.0\n", + " MALE\n", " \n", " \n", - " 1333\n", - " Biscoe\n", - " 43.2\n", - " 19.0\n", - " 197.0\n", - " 4775.0\n", + " 38\n", + " Dream\n", + " 41.3\n", + " 20.3\n", + " 194.0\n", + " 3550.0\n", " MALE\n", " \n", " \n", - " 1414\n", - " Torgersen\n", - " 38.7\n", - " 19.0\n", - " 195.0\n", - " 3450.0\n", + " 41\n", + " Dream\n", + " 35.7\n", + " 18.0\n", + " 202.0\n", + " 3550.0\n", " FEMALE\n", " \n", " \n", - " 1197\n", + " 51\n", " Dream\n", - " 41.1\n", - " 19.0\n", - " 182.0\n", + " 38.1\n", + " 17.6\n", + " 187.0\n", " 3425.0\n", - " MALE\n", - " \n", - " \n", - " 1443\n", - " Torgersen\n", - " 40.6\n", - " 19.0\n", - " 199.0\n", - " 4000.0\n", - " MALE\n", + " FEMALE\n", " \n", " \n", - " 1295\n", - " Biscoe\n", - " 41.0\n", - " 20.0\n", - " 203.0\n", - " 4725.0\n", - " MALE\n", + " 53\n", + " Dream\n", + " 36.0\n", + " 17.1\n", + " 187.0\n", + " 3700.0\n", + " FEMALE\n", " \n", " \n", "\n", + "

25 rows × 6 columns

\n", "[146 rows x 6 columns in total]" ], "text/plain": [ - " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", - "tag_number \n", - "1172 Dream 32.1 15.5 188.0 \n", - "1371 Biscoe 37.7 16.0 183.0 \n", - "1417 Torgersen 38.6 17.0 188.0 \n", - "1204 Dream 40.7 17.0 190.0 \n", - "1251 Biscoe 37.6 17.0 185.0 \n", - "1422 Torgersen 35.7 17.0 189.0 \n", - "1394 Torgersen 40.2 17.0 176.0 \n", - "1163 Dream 36.4 17.0 195.0 \n", - "1329 Biscoe 38.1 17.0 181.0 \n", - "1406 Torgersen 44.1 18.0 210.0 \n", - "1196 Dream 36.5 18.0 182.0 \n", - "1228 Biscoe 41.6 18.0 192.0 \n", - "1412 Torgersen 40.3 18.0 195.0 \n", - "1142 Dream 35.7 18.0 202.0 \n", - "1430 Torgersen 33.5 19.0 190.0 \n", - "1333 Biscoe 43.2 19.0 197.0 \n", - "1414 Torgersen 38.7 19.0 195.0 \n", - "1197 Dream 41.1 19.0 182.0 \n", - "1443 Torgersen 40.6 19.0 199.0 \n", - "1295 Biscoe 41.0 20.0 203.0 \n", - "1207 Dream 38.8 20.0 190.0 \n", - "1349 Biscoe 38.2 20.0 190.0 \n", - "1350 Biscoe 37.8 20.0 190.0 \n", - "1351 Biscoe 38.1 16.5 198.0 \n", - "1116 Dream 37.0 16.5 185.0 \n", + " island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", + "0 Dream 36.6 18.4 184.0 3475.0 \n", + "1 Dream 39.8 19.1 184.0 4650.0 \n", + "2 Dream 40.9 18.9 184.0 3900.0 \n", + "4 Dream 37.3 16.8 192.0 3000.0 \n", + "5 Dream 43.2 18.5 192.0 4100.0 \n", + "9 Dream 40.2 20.1 200.0 3975.0 \n", + "10 Dream 40.8 18.9 208.0 4300.0 \n", + "11 Dream 39.0 18.7 185.0 3650.0 \n", + "12 Dream 37.0 16.9 185.0 3000.0 \n", + "14 Dream 34.0 17.1 185.0 3400.0 \n", + "15 Dream 37.0 16.5 185.0 3400.0 \n", + "18 Dream 39.7 17.9 193.0 4250.0 \n", + "19 Dream 37.8 18.1 193.0 3750.0 \n", + "22 Dream 40.2 17.1 193.0 3400.0 \n", + "23 Dream 36.8 18.5 193.0 3500.0 \n", + "26 Dream 41.5 18.5 201.0 4000.0 \n", + "31 Dream 33.1 16.1 178.0 2900.0 \n", + "32 Dream 37.2 18.1 178.0 3900.0 \n", + "33 Dream 39.5 16.7 178.0 3250.0 \n", + "35 Dream 36.0 18.5 186.0 3100.0 \n", + "36 Dream 39.6 18.1 186.0 4450.0 \n", + "38 Dream 41.3 20.3 194.0 3550.0 \n", + "41 Dream 35.7 18.0 202.0 3550.0 \n", + "51 Dream 38.1 17.6 187.0 3425.0 \n", + "53 Dream 36.0 17.1 187.0 3700.0 \n", "\n", - " body_mass_g sex \n", - "tag_number \n", - "1172 3050.0 FEMALE \n", - "1371 3075.0 FEMALE \n", - "1417 2900.0 FEMALE \n", - "1204 3725.0 MALE \n", - "1251 3600.0 FEMALE \n", - "1422 3350.0 FEMALE \n", - "1394 3450.0 FEMALE \n", - "1163 3325.0 FEMALE \n", - "1329 3175.0 FEMALE \n", - "1406 4000.0 MALE \n", - "1196 3150.0 FEMALE \n", - "1228 3950.0 MALE \n", - "1412 3250.0 FEMALE \n", - "1142 3550.0 FEMALE \n", - "1430 3600.0 FEMALE \n", - "1333 4775.0 MALE \n", - "1414 3450.0 FEMALE \n", - "1197 3425.0 MALE \n", - "1443 4000.0 MALE \n", - "1295 4725.0 MALE \n", - "1207 3950.0 MALE \n", - "1349 3900.0 MALE \n", - "1350 4250.0 MALE \n", - "1351 3825.0 FEMALE \n", - "1116 3400.0 FEMALE \n", + " sex \n", + "0 FEMALE \n", + "1 MALE \n", + "2 MALE \n", + "4 FEMALE \n", + "5 MALE \n", + "9 MALE \n", + "10 MALE \n", + "11 MALE \n", + "12 FEMALE \n", + "14 FEMALE \n", + "15 FEMALE \n", + "18 MALE \n", + "19 MALE \n", + "22 FEMALE \n", + "23 FEMALE \n", + "26 MALE \n", + "31 FEMALE \n", + "32 MALE \n", + "33 FEMALE \n", + "35 FEMALE \n", + "36 MALE \n", + "38 MALE \n", + "41 FEMALE \n", + "51 FEMALE \n", + "53 FEMALE \n", "...\n", "\n", "[146 rows x 6 columns]" ] }, - "execution_count": 2, + "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "# set a friendlier index to uniquely identify the rows\n", - "df = df.set_index(\"tag_number\")\n", - "\n", "# filter down to the data we want to analyze\n", "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", "\n", @@ -682,7 +842,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -704,7 +864,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 6, "metadata": {}, "outputs": [], "source": [ @@ -719,9 +879,79 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 7, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "367359e2069c4d198a820d1ced057b81", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 56778fb7-779c-4e44-b2a3-04d2e174c562 is DONE. 31.9 kB processed. \n", " \n", "\n", + "

1 rows × 6 columns

\n", "[1 rows x 6 columns in total]" ], "text/plain": [ @@ -775,7 +1006,7 @@ "[1 rows x 6 columns]" ] }, - "execution_count": 5, + "execution_count": 7, "metadata": {}, "output_type": "execute_result" } @@ -787,9 +1018,79 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 8, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "025941ee47864dae956ef25bfe815da4", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 1eabd729-6c75-4087-9c87-0d95327b615c is RUNNING. \n", " predicted_body_mass_g\n", " \n", - " \n", - " tag_number\n", - " \n", - " \n", " \n", " \n", " \n", - " 1393\n", + " 292\n", " 3459.735118\n", " \n", - " \n", - " 1525\n", - " 3947.881639\n", - " \n", - " \n", - " 1524\n", - " 4304.175638\n", - " \n", - " \n", - " 1523\n", - " 3471.668379\n", - " \n", " \n", "\n", - "[4 rows x 1 columns in total]" + "

1 rows × 1 columns

\n", + "[1 rows x 1 columns in total]" ], "text/plain": [ - " predicted_body_mass_g\n", - "tag_number \n", - "1393 3459.735118\n", - "1525 3947.881639\n", - "1524 4304.175638\n", - "1523 3471.668379\n", + " predicted_body_mass_g\n", + "292 3459.735118\n", "\n", - "[4 rows x 1 columns]" + "[1 rows x 1 columns]" ] }, - "execution_count": 6, + "execution_count": 8, "metadata": {}, "output_type": "execute_result" } @@ -870,7 +1152,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 9, "metadata": {}, "outputs": [ { @@ -879,7 +1161,7 @@ "LinearRegression()" ] }, - "execution_count": 7, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -906,7 +1188,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.9" + "version": "3.10.12" }, "orig_nbformat": 4, "vscode": { diff --git a/notebooks/06 - Using ML - Large Language Models.ipynb b/notebooks/06 - Using ML - Large Language Models.ipynb index c2c0b83bef2..45a46c44af9 100644 --- a/notebooks/06 - Using ML - Large Language Models.ipynb +++ b/notebooks/06 - Using ML - Large Language Models.ipynb @@ -137,7 +137,7 @@ } ], "source": [ - "pred = model.predict(bf_df).compute()\n", + "pred = model.predict(bf_df).to_pandas()\n", "pred" ] }, diff --git a/notebooks/10 - Regionalized.ipynb b/notebooks/10 - Regionalized.ipynb new file mode 100644 index 00000000000..a7ff5db84e3 --- /dev/null +++ b/notebooks/10 - Regionalized.ipynb @@ -0,0 +1,2800 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# README\n", + "\n", + "This Notebook runs differently depending on the following environent variable:\n", + "1. BIGQUERY_LOCATION - can take values as per https://cloud.google.com/bigquery/docs/locations, e.g. `us`, `asia-east1`" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Infer location and set up data in that location if needed" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n", + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "BigQuery project: bigframes-dev\n", + "BigQuery location: us\n", + "Penguins Table: bigquery-public-data.ml_datasets.penguins\n", + "ML Model Dataset: bigframes-dev.bigframes_testing\n" + ] + } + ], + "source": [ + "# Take multi-region US as the default BQ location, where most of the BQ data lies including the BQ public datasets\n", + "BQ_LOCATION = \"us\"\n", + "PROJECT = \"bigframes-dev\"\n", + "DATASET = \"bigframes_testing\"\n", + "PENGUINS_TABLE = \"bigquery-public-data.ml_datasets.penguins\"\n", + "\n", + "\n", + "# Check for a location set in the environment and do location-specific setup if needed\n", + "\n", + "import os\n", + "import google.api_core.exceptions\n", + "from google.cloud import bigquery\n", + "import bigframes\n", + " \n", + "env_bq_location = os.getenv(\"BIGQUERY_LOCATION\")\n", + "if env_bq_location and env_bq_location != BQ_LOCATION:\n", + " BQ_LOCATION = env_bq_location.lower()\n", + "\n", + "client = bigquery.Client()\n", + "\n", + "if BQ_LOCATION != \"us\":\n", + " bq_location_normalized = BQ_LOCATION.replace('-', '_')\n", + "\n", + " # Nominate a local penguins table\n", + " penguins_table_ref = bigquery.TableReference.from_string(PENGUINS_TABLE)\n", + " penguins_local_dataset_name = f\"{DATASET}_{bq_location_normalized}\"\n", + " penguins_local_dataset_ref = bigquery.DatasetReference(project=PROJECT, dataset_id=penguins_local_dataset_name)\n", + " penguins_local_dataset = bigquery.Dataset(penguins_local_dataset_ref)\n", + " penguins_local_dataset.location = BQ_LOCATION\n", + " penguins_local_table_ref= bigquery.TableReference(penguins_local_dataset, penguins_table_ref.table_id)\n", + " penguins_local_table = str(penguins_local_table_ref)\n", + " try:\n", + " client.get_table(penguins_local_table_ref)\n", + " except google.api_core.exceptions.NotFound:\n", + " client.create_dataset(penguins_local_dataset, exists_ok=True)\n", + "\n", + " # Read the public table as an in-memory dataframe and then write to the local table\n", + " session_us = bigframes.connect()\n", + " df = session_us.read_gbq(PENGUINS_TABLE).to_pandas()\n", + " df.to_gbq(penguins_local_table)\n", + "\n", + " # Finally point the penguins table to the local table\n", + " PENGUINS_TABLE=penguins_local_table\n", + "\n", + " # Also update the dataset name used for test artifacts\n", + " DATASET = f\"{DATASET}_{bq_location_normalized}\"\n", + "\n", + "# Create the dataset to store the model if it doesn't exist \n", + "model_local_dataset = bigquery.Dataset(bigquery.DatasetReference(project=PROJECT, dataset_id=DATASET))\n", + "model_local_dataset.location = BQ_LOCATION\n", + "model_dataset = client.create_dataset(model_local_dataset, exists_ok=True)\n", + "\n", + "# Finally log the variables driving the core notebook execution\n", + "log = ('\\n'.join(f\"{name}: {str(value)}\" for name, value in {\n", + " \"BigQuery project\" : PROJECT,\n", + " \"BigQuery location\" : BQ_LOCATION,\n", + " \"Penguins Table\" : PENGUINS_TABLE,\n", + " \"ML Model Dataset\" : model_dataset.reference\n", + "}.items())) \n", + "print(log)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the BigQuery DataFrames API" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Set BigQuery DataFrames options" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import bigframes.pandas\n", + "\n", + "bigframes.pandas.options.bigquery.project = PROJECT\n", + "bigframes.pandas.options.bigquery.location = BQ_LOCATION" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize a dataframe for a BigQuery table" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "02af1103261a4e63a4c15efd26b1bc9a", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job e1a62d56-8cab-4bc1-9ad3-457f48b71d9c is RUNNING. = 4000.0]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the Remote Functions" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### BigQuery DataFrames gives you the ability to turn your custom scalar functions into a BigQuery remote function.\n", + "\n", + "It requires the GCP project to be set up appropriately and the user having sufficient privileges to use them. One can find more details on it via `help` command." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Help on function remote_function in module bigframes.pandas:\n", + "\n", + "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True)\n", + " Decorator to turn a user defined function into a BigQuery remote function.\n", + " \n", + " .. note::\n", + " Please make sure following is setup before using this API:\n", + " \n", + " 1. Have the below APIs enabled for your project:\n", + " a. BigQuery Connection API\n", + " b. Cloud Functions API\n", + " c. Cloud Run API\n", + " d. Cloud Build API\n", + " e. Artifact Registry API\n", + " f. Cloud Resource Manager API\n", + " \n", + " This can be done from the cloud console (change PROJECT_ID to yours):\n", + " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", + " Or from the gcloud CLI:\n", + " $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com\n", + " \n", + " 2. Have following IAM roles enabled for you:\n", + " a. BigQuery Data Editor (roles/bigquery.dataEditor)\n", + " b. BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", + " c. Cloud Functions Developer (roles/cloudfunctions.developer)\n", + " d. Service Account User (roles/iam.serviceAccountUser)\n", + " e. Storage Object Viewer (roles/storage.objectViewer)\n", + " f. Project IAM Admin (roles/resourcemanager.projectIamAdmin)\n", + " (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", + " \n", + " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", + " a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", + " b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", + " Alternatively, the IAM could also be setup via the gcloud CLI:\n", + " $ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"\n", + " \n", + " Args:\n", + " input_types (list(type)):\n", + " List of input data types in the user defined function.\n", + " output_type (type):\n", + " Data type of the output in the user defined function.\n", + " dataset (str, Optional):\n", + " Dataset to use to create a BigQuery function. It should be in\n", + " `.` or `` format. If this\n", + " param is not provided then session dataset id would be used.\n", + " bigquery_connection (str, Optional):\n", + " Name of the BigQuery connection. If it is pre created in the same\n", + " location as the `bigquery_client.location` then it would be used,\n", + " otherwise it would be created dynamically assuming the user has\n", + " necessary priviliges. If this param is not provided then the\n", + " bigquery connection from the session would be used.\n", + " reuse (bool, Optional):\n", + " Reuse the remote function if already exists.\n", + " `True` by default, which will result in reusing an existing remote\n", + " function (if any) that was previously created for the same udf.\n", + " Setting it to false would force creating a unique remote function.\n", + " If the required remote function does not exist then it would be\n", + " created irrespective of this param.\n", + "\n" + ] + } + ], + "source": [ + "import bigframes.pandas as pd\n", + "help(pd.remote_function)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Define a custom function, and specify the intent to turn it into a remote function.\n", + "\n", + "It requires a BigQuery connection. If the connection is not already created,\n", + "the BigQuery DataFrames package attempts to create one assuming the necessary\n", + "APIs and IAM permissions are setup in the project." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[INFO][2023-08-05 23:12:12,870][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-f9320ad496b5aeca2d7f343cbab03e2f --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmps5m0qu4z --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", + "[INFO][2023-08-05 23:13:20,660][bigframes.remote_function] Successfully created cloud function bigframes-f9320ad496b5aeca2d7f343cbab03e2f with uri (https://bigframes-f9320ad496b5aeca2d7f343cbab03e2f-7krlje3eoq-uc.a.run.app)\n", + "[INFO][2023-08-05 23:13:32,717][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", + "[INFO][2023-08-05 23:13:32,719][bigframes.remote_function] Creating BQ remote function: \n", + " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_f9320ad496b5aeca2d7f343cbab03e2f(num FLOAT64)\n", + " RETURNS STRING\n", + " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", + " OPTIONS (\n", + " endpoint = \"https://bigframes-f9320ad496b5aeca2d7f343cbab03e2f-7krlje3eoq-uc.a.run.app\"\n", + " )\n", + "[INFO][2023-08-05 23:13:33,697][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_f9320ad496b5aeca2d7f343cbab03e2f\n" + ] + } + ], + "source": [ + "@pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", + "def get_bucket(num):\n", + " if not num: return \"NA\"\n", + " boundary = 4000\n", + " return \"at_or_above_4000\" if num >= boundary else \"below_4000\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run the custom function on the BigQuery-backed dataframe" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "0dba87f5bcb74dca9efebe8f522beeff", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 5f30816f-f4d0-4063-bb9e-2905b89f717d is DONE. 2.8 kB processed. \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
body_mass_gbody_mass_bucket
03475.0below_4000
14650.0at_or_above_4000
23900.0below_4000
33500.0below_4000
43000.0below_4000
54100.0at_or_above_4000
62700.0below_4000
73400.0below_4000
83800.0below_4000
93975.0below_4000
\n", + "

10 rows × 2 columns

\n", + "[10 rows x 2 columns in total]" + ], + "text/plain": [ + " body_mass_g body_mass_bucket\n", + "0 3475.0 below_4000\n", + "1 4650.0 at_or_above_4000\n", + "2 3900.0 below_4000\n", + "3 3500.0 below_4000\n", + "4 3000.0 below_4000\n", + "5 4100.0 at_or_above_4000\n", + "6 2700.0 below_4000\n", + "7 3400.0 below_4000\n", + "8 3800.0 below_4000\n", + "9 3975.0 below_4000\n", + "\n", + "[10 rows x 2 columns]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket))\n", + "df[['body_mass_g', 'body_mass_bucket']].head(10)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using the ML API" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize a DataFrame from a BigQuery table" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "89bf7ae3fd8641bcbdc5a4614a82f48a", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job acd770bb-5ccb-463f-beec-2386132ded6b is RUNNING. \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
islandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Dream36.618.4184.03475.0FEMALE
1Dream39.819.1184.04650.0MALE
2Dream40.918.9184.03900.0MALE
4Dream37.316.8192.03000.0FEMALE
5Dream43.218.5192.04100.0MALE
9Dream40.220.1200.03975.0MALE
10Dream40.818.9208.04300.0MALE
11Dream39.018.7185.03650.0MALE
12Dream37.016.9185.03000.0FEMALE
14Dream34.017.1185.03400.0FEMALE
15Dream37.016.5185.03400.0FEMALE
18Dream39.717.9193.04250.0MALE
19Dream37.818.1193.03750.0MALE
22Dream40.217.1193.03400.0FEMALE
23Dream36.818.5193.03500.0FEMALE
26Dream41.518.5201.04000.0MALE
31Dream33.116.1178.02900.0FEMALE
32Dream37.218.1178.03900.0MALE
33Dream39.516.7178.03250.0FEMALE
35Dream36.018.5186.03100.0FEMALE
36Dream39.618.1186.04450.0MALE
38Dream41.320.3194.03550.0MALE
41Dream35.718.0202.03550.0FEMALE
51Dream38.117.6187.03425.0FEMALE
53Dream36.017.1187.03700.0FEMALE
\n", + "

25 rows × 6 columns

\n", + "[146 rows x 6 columns in total]" + ], + "text/plain": [ + " island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", + "0 Dream 36.6 18.4 184.0 3475.0 \n", + "1 Dream 39.8 19.1 184.0 4650.0 \n", + "2 Dream 40.9 18.9 184.0 3900.0 \n", + "4 Dream 37.3 16.8 192.0 3000.0 \n", + "5 Dream 43.2 18.5 192.0 4100.0 \n", + "9 Dream 40.2 20.1 200.0 3975.0 \n", + "10 Dream 40.8 18.9 208.0 4300.0 \n", + "11 Dream 39.0 18.7 185.0 3650.0 \n", + "12 Dream 37.0 16.9 185.0 3000.0 \n", + "14 Dream 34.0 17.1 185.0 3400.0 \n", + "15 Dream 37.0 16.5 185.0 3400.0 \n", + "18 Dream 39.7 17.9 193.0 4250.0 \n", + "19 Dream 37.8 18.1 193.0 3750.0 \n", + "22 Dream 40.2 17.1 193.0 3400.0 \n", + "23 Dream 36.8 18.5 193.0 3500.0 \n", + "26 Dream 41.5 18.5 201.0 4000.0 \n", + "31 Dream 33.1 16.1 178.0 2900.0 \n", + "32 Dream 37.2 18.1 178.0 3900.0 \n", + "33 Dream 39.5 16.7 178.0 3250.0 \n", + "35 Dream 36.0 18.5 186.0 3100.0 \n", + "36 Dream 39.6 18.1 186.0 4450.0 \n", + "38 Dream 41.3 20.3 194.0 3550.0 \n", + "41 Dream 35.7 18.0 202.0 3550.0 \n", + "51 Dream 38.1 17.6 187.0 3425.0 \n", + "53 Dream 36.0 17.1 187.0 3700.0 \n", + "\n", + " sex \n", + "0 FEMALE \n", + "1 MALE \n", + "2 MALE \n", + "4 FEMALE \n", + "5 MALE \n", + "9 MALE \n", + "10 MALE \n", + "11 MALE \n", + "12 FEMALE \n", + "14 FEMALE \n", + "15 FEMALE \n", + "18 MALE \n", + "19 MALE \n", + "22 FEMALE \n", + "23 FEMALE \n", + "26 MALE \n", + "31 FEMALE \n", + "32 MALE \n", + "33 FEMALE \n", + "35 FEMALE \n", + "36 MALE \n", + "38 MALE \n", + "41 FEMALE \n", + "51 FEMALE \n", + "53 FEMALE \n", + "...\n", + "\n", + "[146 rows x 6 columns]" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# filter down to the data we want to analyze\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# drop the columns we don't care about\n", + "adelie_data = adelie_data.drop(columns=[\"species\"])\n", + "\n", + "# drop rows with nulls to get our training data\n", + "training_data = adelie_data.dropna()\n", + "\n", + "# take a peek at the training data\n", + "training_data" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [], + "source": [ + "# pick feature columns and label column\n", + "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "label_columns = training_data[['body_mass_g']]\n", + "\n", + "# also get the rows that we want to make predictions for (i.e. where the feature column is null)\n", + "missing_body_mass = adelie_data[adelie_data.body_mass_g.isnull()]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Train and evaluate a linear regression model using the ML API" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "865f6bb75c5b48e4a52a3183fe3c2582", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 4b0c58e4-4752-4b96-b490-a95e3ae326c0 is DONE. 31.9 kB processed. \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.87876378553.6016340.005614181.3309110.6239510.623951
\n", + "

1 rows × 6 columns

\n", + "[1 rows x 6 columns in total]" + ], + "text/plain": [ + " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", + "0 223.878763 78553.601634 0.005614 \n", + "\n", + " median_absolute_error r2_score explained_variance \n", + "0 181.330911 0.623951 0.623951 \n", + "\n", + "[1 rows x 6 columns]" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "\n", + "# as in scikit-learn, a newly created model is just a bundle of parameters\n", + "# default parameters are fine here\n", + "model = LinearRegression()\n", + "\n", + "# this will train a temporary model in BigQuery Machine Learning\n", + "model.fit(feature_columns, label_columns)\n", + "\n", + "# check how the model performed\n", + "model.score(feature_columns, label_columns)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Make predictions using the model" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "caf8cefe06a14da2a02f31aa1e12c23a", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job fe2fe252-8433-4d20-861c-681a8dfbf2c4 is RUNNING. \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
predicted_body_mass_g
2923459.735118
\n", + "

1 rows × 1 columns

\n", + "[1 rows x 1 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g\n", + "292 3459.735118\n", + "\n", + "[1 rows x 1 columns]" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.predict(missing_body_mass)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Save the trained model to BigQuery, so we can load it later" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "LinearRegression()" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.to_gbq(f\"{DATASET}.penguins_model\", replace=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/noxfile.py b/noxfile.py index 5e54c12fb2f..d4d97ca2c4e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -16,6 +16,7 @@ from __future__ import absolute_import +from multiprocessing import Process import os import pathlib import re @@ -82,7 +83,6 @@ "docfx", "unit", "unit_noextras", - "unit_prerelease", "system", "doctest", "cover", @@ -603,24 +603,62 @@ def notebook(session): ] notebooks = [os.path.join("notebooks", nb) for nb in notebooks] + # Regionalized notebooks + notebooks_reg = { + "10 - Regionalized.ipynb": [ + "asia-southeast1", + "eu", + "europe-west4", + "southamerica-west1", + "us", + "us-central1", + ] + } + notebooks_reg = { + os.path.join("notebooks", nb): regions for nb, regions in notebooks_reg.items() + } + # For some reason nbmake exits silently with "no tests ran" message if # one of the notebook paths supplied does not exist. Let's make sure that # each path exists - for nb in notebooks: + for nb in notebooks + list(notebooks_reg): assert os.path.exists(nb), nb - # Use retries because sometimes parallel runs of the same notebook can try - # to create the same artifacts and may run into resoure conflict at the GCP - # level. - session.run( + # TODO(shobs): For some reason --retries arg masks exceptions occurred in + # notebook failures, and shows unhelpful INTERNALERROR. Investigate that + # and enable retries if we can find a way to surface the real exception + # bacause the notebook is running against real GCP and something may fail + # due to transient issues. + pytest_command = [ "py.test", - "-nauto", "--nbmake", "--nbmake-timeout=600", - "--retries=3", + ] + + # Run self-contained notebooks in single session.run + # achieve parallelization via -n + session.run( + *pytest_command, + "-nauto", *notebooks, ) + # Run regionalized notebooks in parallel session.run's, since each notebook + # takes a different region via env param. + processes = [] + for notebook, regions in notebooks_reg.items(): + for region in regions: + process = Process( + target=session.run, + args=(*pytest_command, notebook), + kwargs={"env": {"BIGQUERY_LOCATION": region}}, + ) + process.start() + processes.append(process) + + for process in processes: + process.join() + @nox.session(python="3.10") def release_dry_run(session): diff --git a/samples/snippets/quickstart.py b/samples/snippets/quickstart.py index dc8ab7ba2d0..a15ea168534 100644 --- a/samples/snippets/quickstart.py +++ b/samples/snippets/quickstart.py @@ -14,64 +14,46 @@ def run_quickstart(project_id: str): - # [START bigquery_bigframes_quickstart] - import bigframes.pandas as pd + import bigframes + + session_options = bigframes.BigQueryOptions() + session = bigframes.connect(session_options) + + your_gcp_project_id = project_id + query_or_table = "bigquery-public-data.ml_datasets.penguins" + df_session = session.read_gbq(query_or_table) + average_body_mass = df_session["body_mass_g"].mean() + print(f"average_body_mass (df_session): {average_body_mass}") - # TODO: (Optional) Setup your session with the configuration. Some of these - # settings cannot be changed once a session has started. - pd.options.bigquery.project = "your-gcp-project-id" - pd.options.bigquery.location = "us" + # [START bigquery_bigframes_quickstart] + import bigframes.pandas as bpd - # [START_EXCLUDE silent] - # Ignore the lines between "EXCLUDE" comments. They replace values so this - # sample runs in our test suite. - pd.options.bigquery.project = project_id - # [END_EXCLUDE] + # Set BigQuery DataFrames options + bpd.options.bigquery.project = your_gcp_project_id + bpd.options.bigquery.location = "us" + # Create a DataFrame from a BigQuery table query_or_table = "bigquery-public-data.ml_datasets.penguins" - df = pd.read_gbq(query_or_table) + df = bpd.read_gbq(query_or_table) # 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}") - # IMPORTANT: The `bigframes.pandas` package creates a BigQuery session for - # queries and temporary tables. A BigQuery session has a limited lifetime - # (https://cloud.google.com/bigquery/docs/sessions-intro#limitations) and - # does not support concurrent queries. For long lived applications, create - # session objects as needed, instead. - - import bigframes - - session_options = bigframes.BigQueryOptions() - session_options.project = "your-gcp-project-id" - session_options.location = "us" - - # [START_EXCLUDE silent] - # Ignore the lines between "EXCLUDE" comments. They replace values so this - # sample runs in our test suite. - session_options.project = project_id - # [END_EXCLUDE] - - session = bigframes.connect(session_options) - df_session = session.read_gbq(query_or_table) - average_body_mass = df_session["body_mass_g"].mean() - print(f"average_body_mass (df_session): {average_body_mass}") - + # Create the Linear Regression model from bigframes.ml.linear_model import LinearRegression - # filter down to the data we want to analyze + # 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 + # Drop the columns we don't care about adelie_data = adelie_data.drop(columns=["species"]) - # drop rows with nulls to get our training data + # Drop rows with nulls to get our training data training_data = adelie_data.dropna() - # pick feature columns and label column + # Pick feature columns and label column X = training_data[ [ "island", diff --git a/scripts/upload_to_google_drive.py b/scripts/upload_to_google_drive.py new file mode 100644 index 00000000000..e5791513595 --- /dev/null +++ b/scripts/upload_to_google_drive.py @@ -0,0 +1,71 @@ +# 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. + +""" +Upload latest wheel to Google Drive. + +Based on +https://github.com/googleapis/google-resumable-media-python/blob/main/google/resumable_media/requests/__init__.py + +Before running, execute the following to make sure you can use the Google Drive API: + +gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/sqlservice.login,https://www.googleapis.com/auth/drive +""" + +import pathlib + +import google.auth +import google.auth.transport.requests +import google.resumable_media._upload +import google.resumable_media.requests as resumable_requests + +repo_root = pathlib.Path(__file__).parent.parent + +# Use PATCH instead of POST to replace existing files. +google.resumable_media._upload._POST = "PATCH" + +credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/drive"]) +transport = google.auth.transport.requests.AuthorizedSession(credentials) + +wheel_id = "15fZ1DkrFDk4ibMNTzms4akpxmf2pzeAR" +wheel_path = next(iter((repo_root / "dist").glob("bigframes-*.whl"))) + +pdf_id = "1agYjxmPLrxelsaHI-lc41QHcgnQYemcX" +pdf_path = repo_root / "docs" / "_build" / "latex" / "bigframes-latest.pdf" + +uploads = ( + (wheel_id, wheel_path, "application/octet-stream"), + (pdf_id, pdf_path, "application/pdf"), +) + +upload_template = ( + "https://www.googleapis.com/upload/drive/v3/files/{file_id}?uploadType=resumable" +) +chunk_size = 1024 * 1024 # 1MB + +for file_id, file_path, content_type in uploads: + print(f"Uploading {file_path}") + transport = google.auth.transport.requests.AuthorizedSession(credentials) + upload = resumable_requests.ResumableUpload( + upload_template.format(file_id=file_id), chunk_size + ) + + with open(file_path, "rb") as stream: + response = upload.initiate( + transport, stream, metadata={}, content_type=content_type + ) + print(response) + while not upload.finished: + response = upload.transmit_next_chunk(transport) + print(response) diff --git a/setup.py b/setup.py index 311471f2627..3e36cd3ecfb 100644 --- a/setup.py +++ b/setup.py @@ -30,7 +30,7 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 3 - Alpha" dependencies = [ - "cloudpickle >= 2.2.1", + "cloudpickle >= 2.0.0", "fsspec >=2023.3.0", "gcsfs >=2023.3.0", "geopandas >=0.12.2", @@ -39,9 +39,10 @@ "google-cloud-functions >=1.10.1", "google-cloud-bigquery-connection >=1.12.0", "google-cloud-storage >=2.0.0", - "ibis-framework[bigquery] >=6.0.0", + # TODO: Relax upper bound once we have fixed `system_prerelease` tests. + "ibis-framework[bigquery] >=6.0.0,<=6.1.0", "pandas >=1.5.0", - "pydata-google-auth >=1.5.0", + "pydata-google-auth >=1.8.2", "scikit-learn >=1.2.2", "sqlalchemy >=1.4,<3.0", "ipywidgets >=7.7.1", diff --git a/testing/constraints-3.9.txt b/testing/constraints-3.9.txt index aa72e5426b7..fe3d49ef208 100644 --- a/testing/constraints-3.9.txt +++ b/testing/constraints-3.9.txt @@ -11,7 +11,7 @@ cffi==1.15.1 cfgv==3.3.1 charset-normalizer==3.1.0 click==8.1.3 -cloudpickle==2.2.1 +cloudpickle==2.0.0 colorlog==6.7.0 coverage==7.2.2 cryptography==40.0.1 @@ -81,7 +81,7 @@ pyarrow==11.0.0 pyasn1==0.4.8 pyasn1-modules==0.2.8 pycparser==2.21 -pydata-google-auth==1.5.0 +pydata-google-auth==1.8.2 Pygments==2.14.0 PyJWT==2.6.0 pyperclip==1.8.2 diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 03638698966..b6b02e4c27a 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -22,6 +22,7 @@ import google.cloud.bigquery as bigquery import google.cloud.bigquery_connection_v1 as bigquery_connection_v1 import google.cloud.exceptions +import google.cloud.functions_v2 as functions_v2 import google.cloud.storage as storage # type: ignore import ibis.backends.base import pandas as pd @@ -93,6 +94,13 @@ def bigqueryconnection_client( return session.bqconnectionclient +@pytest.fixture(scope="session") +def cloudfunctions_client( + session: bigframes.Session, +) -> functions_v2.FunctionServiceClient: + return session.cloudfunctionsclient + + @pytest.fixture(scope="session") def session() -> bigframes.Session: return bigframes.Session() @@ -665,3 +673,19 @@ def penguins_randomforest_classifier_model_name( session.bqclient.query(sql).result() finally: return model_name + + +@pytest.fixture() +def deferred_repr(): + bigframes.options.display.repr_mode = "deferred" + yield + bigframes.options.display.repr_mode = "head" + + +@pytest.fixture() +def restore_sampling_settings(): + enable_downsampling = bigframes.options.sampling.enable_downsampling + max_download_size = bigframes.options.sampling.max_download_size + yield + bigframes.options.sampling.enable_downsampling = enable_downsampling + bigframes.options.sampling.max_download_size = max_download_size diff --git a/tests/system/large/ml/test_cluster.py b/tests/system/large/ml/test_cluster.py index bb64b19076f..eae6896669b 100644 --- a/tests/system/large/ml/test_cluster.py +++ b/tests/system/large/ml/test_cluster.py @@ -12,14 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pandas +import pandas as pd +import pytest -import bigframes.ml.cluster +from bigframes.ml import cluster from tests.system.utils import assert_pandas_df_equal_ignore_ordering -def test_cluster_configure_fit_predict(session, penguins_df_default_index, dataset_id): - model = bigframes.ml.cluster.KMeans(n_clusters=3) +@pytest.mark.flaky(retries=2, delay=120) +def test_cluster_configure_fit_score_predict( + session, penguins_df_default_index, dataset_id +): + model = cluster.KMeans(n_clusters=3) df = penguins_df_default_index.dropna()[ [ @@ -37,7 +41,7 @@ def test_cluster_configure_fit_predict(session, penguins_df_default_index, datas model.fit(df) - pd_new_penguins = pandas.DataFrame.from_dict( + pd_new_penguins = pd.DataFrame.from_dict( { "test1": { "species": "Adelie Penguin (Pygoscelis adeliae)", @@ -81,13 +85,24 @@ def test_cluster_configure_fit_predict(session, penguins_df_default_index, datas pd_new_penguins.index.name = "observation" new_penguins = session.read_pandas(pd_new_penguins) - result = model.predict(new_penguins).compute() - expected = pandas.DataFrame( + + # Check score to ensure the model was fitted + score_result = model.score(new_penguins).to_pandas() + score_expected = pd.DataFrame( + {"davies_bouldin_index": [1.502182], "mean_squared_distance": [1.953408]}, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + result = model.predict(new_penguins).to_pandas() + expected = pd.DataFrame( {"CENTROID_ID": [2, 3, 1, 2]}, dtype="Int64", - index=pandas.Index( - ["test1", "test2", "test3", "test4"], dtype="string[pyarrow]" - ), + index=pd.Index(["test1", "test2", "test3", "test4"], dtype="string[pyarrow]"), ) expected.index.name = "observation" assert_pandas_df_equal_ignore_ordering(result, expected) diff --git a/tests/system/large/ml/test_compose.py b/tests/system/large/ml/test_compose.py index 52f347151ed..0c2744819d4 100644 --- a/tests/system/large/ml/test_compose.py +++ b/tests/system/large/ml/test_compose.py @@ -56,8 +56,12 @@ def test_columntransformer_standalone_fit_transform( [{"index": 1, "value": 1.0}], [{"index": 2, "value": 1.0}], ], - "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], - "scaled_flipper_length_mm": [-0.3495, -1.416, -0.9185], + "scaled_culmen_length_mm": [ + -0.811119671289163, + -0.9945520581113803, + -1.104611490204711, + ], + "scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], }, index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"), ) diff --git a/tests/system/large/ml/test_core.py b/tests/system/large/ml/test_core.py index cd829f0ae2d..ab33e5d7185 100644 --- a/tests/system/large/ml/test_core.py +++ b/tests/system/large/ml/test_core.py @@ -37,7 +37,7 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d ) # no data - report evaluation from the automatic data split - evaluate_result = model.evaluate().compute() + evaluate_result = model.evaluate().to_pandas() evaluate_expected = pandas.DataFrame( { "mean_absolute_error": [225.817334], @@ -57,13 +57,13 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d ) # evaluate on all training data - evaluate_result = model.evaluate(df).compute() + evaluate_result = model.evaluate(df).to_pandas() pandas.testing.assert_frame_equal( evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 ) # predict new labels - predictions = model.predict(new_penguins_df).compute() + predictions = model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, dtype="Float64", @@ -104,7 +104,7 @@ def test_bqml_manual_preprocessing_e2e( ) # no data - report evaluation from the automatic data split - evaluate_result = model.evaluate().compute() + evaluate_result = model.evaluate().to_pandas() evaluate_expected = pandas.DataFrame( { "mean_absolute_error": [309.477334], @@ -125,13 +125,13 @@ def test_bqml_manual_preprocessing_e2e( ) # evaluate on all training data - evaluate_result = model.evaluate(df).compute() + evaluate_result = model.evaluate(df).to_pandas() pandas.testing.assert_frame_equal( evaluate_result, evaluate_expected, check_exact=False, rtol=0.1 ) # predict new labels - predictions = model.predict(new_penguins_df).compute() + predictions = model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_body_mass_g": [3968.8, 3176.3, 3545.2]}, dtype="Float64", @@ -156,11 +156,11 @@ def test_bqml_standalone_transform(penguins_df_default_index, new_penguins_df): options={"model_type": "transform_only"}, transforms=[ "ML.STANDARD_SCALER(culmen_length_mm) OVER() AS scaled_culmen_length_mm", - "ML.ONE_HOT_ENCODER(species) OVER() AS onehotencoded_species", + "ML.ONE_HOT_ENCODER(species, 'none', 1000000, 0) OVER() AS onehotencoded_species", ], ) - transformed = model.transform(new_penguins_df).compute() + transformed = model.transform(new_penguins_df).to_pandas() expected = pandas.DataFrame( { "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], diff --git a/tests/system/large/ml/test_decomposition.py b/tests/system/large/ml/test_decomposition.py index 57507c11342..460f07b8168 100644 --- a/tests/system/large/ml/test_decomposition.py +++ b/tests/system/large/ml/test_decomposition.py @@ -12,19 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pandas +import pandas as pd -import bigframes.ml.decomposition +from bigframes.ml import decomposition -def test_decomposition_configure_fit_predict( +def test_decomposition_configure_fit_score_predict( session, penguins_df_default_index, dataset_id ): - model = bigframes.ml.decomposition.PCA(n_components=3) + model = decomposition.PCA(n_components=3) model.fit(penguins_df_default_index) - pd_new_penguins = session.read_pandas( - pandas.DataFrame( + new_penguins = session.read_pandas( + pd.DataFrame( { "tag_number": [1633, 1672, 1690], "species": [ @@ -42,17 +42,31 @@ def test_decomposition_configure_fit_predict( ).set_index("tag_number") ) - result = model.predict(pd_new_penguins).compute() - expected = pandas.DataFrame( + # Check score to ensure the model was fitted + score_result = model.score(new_penguins).to_pandas() + score_expected = pd.DataFrame( + { + "total_explained_variance_ratio": [0.812383], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + result = model.predict(new_penguins).to_pandas() + expected = pd.DataFrame( { "principal_component_1": [-1.459, 2.258, -1.685], "principal_component_2": [-1.120, -1.351, -0.874], "principal_component_3": [-0.646, 0.443, -0.704], }, dtype="Float64", - index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), ) - pandas.testing.assert_frame_equal( + pd.testing.assert_frame_equal( result.sort_index(), expected, check_exact=False, diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py index 0d5da380d1e..88c5ccd2f0d 100644 --- a/tests/system/large/ml/test_ensemble.py +++ b/tests/system/large/ml/test_ensemble.py @@ -15,10 +15,12 @@ from unittest import TestCase import pandas +import pytest import bigframes.ml.ensemble +@pytest.mark.flaky(retries=2, delay=120) def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.XGBRegressor() @@ -37,7 +39,7 @@ def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [97.368139], @@ -62,6 +64,7 @@ def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): ) +@pytest.mark.flaky(retries=2, delay=120) def test_xgbregressor_dart_booster_multiple_params( penguins_df_default_index, dataset_id ): @@ -97,7 +100,7 @@ def test_xgbregressor_dart_booster_multiple_params( model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -136,6 +139,7 @@ def test_xgbregressor_dart_booster_multiple_params( assert reloaded_model.num_parallel_tree == 2 +@pytest.mark.flaky(retries=2, delay=120) def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.XGBClassifier() @@ -153,7 +157,7 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -175,6 +179,7 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): ) +@pytest.mark.flaky(retries=2, delay=120) def test_xgbclassifier_dart_booster_multiple_params( penguins_df_default_index, dataset_id ): @@ -209,7 +214,7 @@ def test_xgbclassifier_dart_booster_multiple_params( model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -248,6 +253,7 @@ def test_xgbclassifier_dart_booster_multiple_params( assert reloaded_model.num_parallel_tree == 2 +@pytest.mark.flaky(retries=2, delay=120) def test_randomforestregressor_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.RandomForestRegressor() @@ -266,7 +272,7 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -288,6 +294,7 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset ) +@pytest.mark.flaky(retries=2, delay=120) def test_randomforestregressor_multiple_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.RandomForestRegressor( tree_method="AUTO", @@ -318,7 +325,7 @@ def test_randomforestregressor_multiple_params(penguins_df_default_index, datase model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -354,6 +361,7 @@ def test_randomforestregressor_multiple_params(penguins_df_default_index, datase assert reloaded_model.enable_global_explain is False +@pytest.mark.flaky(retries=2, delay=120) def test_randomforestclassifier_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.RandomForestClassifier() @@ -371,7 +379,7 @@ def test_randomforestclassifier_default_params(penguins_df_default_index, datase model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -393,6 +401,7 @@ def test_randomforestclassifier_default_params(penguins_df_default_index, datase ) +@pytest.mark.flaky(retries=2, delay=120) def test_randomforestclassifier_multiple_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.RandomForestClassifier( tree_method="AUTO", @@ -422,7 +431,7 @@ def test_randomforestclassifier_multiple_params(penguins_df_default_index, datas model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", diff --git a/tests/system/large/ml/test_forecasting.py b/tests/system/large/ml/test_forecasting.py index 98726eb289d..d1e2d12296d 100644 --- a/tests/system/large/ml/test_forecasting.py +++ b/tests/system/large/ml/test_forecasting.py @@ -27,7 +27,7 @@ def test_arima_plus_model_fit_score( result = model.score( new_time_series_df[["parsed_date"]], new_time_series_df[["total_visits"]] - ).compute() + ).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [154.742547], diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py index 33010a005cb..332b460fe51 100644 --- a/tests/system/large/ml/test_linear_model.py +++ b/tests/system/large/ml/test_linear_model.py @@ -35,7 +35,7 @@ def test_linear_regression_configure_fit_score(penguins_df_default_index, datase model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [225.735767], @@ -80,7 +80,7 @@ def test_linear_regression_manual_split_configure_fit_score( model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [225.735767], @@ -121,7 +121,7 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() expected = pd.DataFrame( { "precision": [0.58085], @@ -169,7 +169,7 @@ def test_logistic_regression_manual_split_configure_fit_score( model.fit(train_X, train_y) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).compute() + result = model.score(train_X, train_y).to_pandas() expected = pd.DataFrame( { "precision": [0.616753], diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 9700ba2bf65..bec1a51a991 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -13,6 +13,7 @@ # limitations under the License. import pandas as pd +import pytest from bigframes.ml import ( cluster, @@ -38,24 +39,24 @@ def test_pipeline_linear_regression_fit_score_predict( ) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - train_y = df[["body_mass_g"]] - pl.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) # Check score to ensure the model was fitted - score_result = pl.score(train_X, train_y).compute() + score_result = pl.score(X_train, y_train).to_pandas() score_expected = pd.DataFrame( { - "mean_absolute_error": [309.477334], - "mean_squared_error": [152184.227218], + "mean_absolute_error": [309.477331], + "mean_squared_error": [152184.227219], "mean_squared_log_error": [0.009524], - "median_absolute_error": [257.727777], + "median_absolute_error": [257.728263], "r2_score": [0.764356], "explained_variance": [0.764356], }, @@ -96,6 +97,61 @@ def test_pipeline_linear_regression_fit_score_predict( ) +def test_pipeline_linear_regression_series_fit_score_predict( + session, penguins_df_default_index +): + """Test a supervised model with a minimal preprocessing step""" + pl = pipeline.Pipeline( + [ + ("scale", preprocessing.StandardScaler()), + ("linreg", linear_model.LinearRegression()), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df["culmen_length_mm"] + y_train = df["body_mass_g"] + pl.fit(X_train, y_train) + + # Check score to ensure the model was fitted + score_result = pl.score(X_train, y_train).to_pandas() + score_expected = pd.DataFrame( + { + "mean_absolute_error": [528.495599], + "mean_squared_error": [421722.261808], + "mean_squared_log_error": [0.022963], + "median_absolute_error": [468.895249], + "r2_score": [0.346999], + "explained_variance": [0.346999], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + + # predict new labels + new_penguins = session.read_pandas( + pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "culmen_length_mm": [39.5, 38.5, 37.9], + } + ).set_index("tag_number") + ) + predictions = pl.predict(new_penguins["culmen_length_mm"]).to_pandas() + expected = pd.DataFrame( + {"predicted_body_mass_g": [3818.845703, 3732.022253, 3679.928123]}, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + pd.testing.assert_frame_equal( + predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 + ) + + def test_pipeline_logistic_regression_fit_score_predict( session, penguins_df_default_index ): @@ -108,18 +164,18 @@ def test_pipeline_logistic_regression_fit_score_predict( ) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - train_y = df[["sex"]] - pl.fit(train_X, train_y) + y_train = df[["sex"]] + pl.fit(X_train, y_train) # Check score to ensure the model was fitted - score_result = pl.score(train_X, train_y).compute() + score_result = pl.score(X_train, y_train).to_pandas() score_expected = pd.DataFrame( { "precision": [0.537091], @@ -166,6 +222,7 @@ def test_pipeline_logistic_regression_fit_score_predict( ) +@pytest.mark.flaky(retries=2, delay=120) def test_pipeline_xgbregressor_fit_score_predict(session, penguins_df_default_index): """Test a supervised model with a minimal preprocessing step""" pl = pipeline.Pipeline( @@ -176,26 +233,26 @@ def test_pipeline_xgbregressor_fit_score_predict(session, penguins_df_default_in ) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - train_y = df[["body_mass_g"]] - pl.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) # Check score to ensure the model was fitted - score_result = pl.score(train_X, train_y).compute() + score_result = pl.score(X_train, y_train).to_pandas() score_expected = pd.DataFrame( { - "mean_absolute_error": [203.4001727989334], - "mean_squared_error": [74898.80551717622], - "mean_squared_log_error": [0.004394266810531861], - "median_absolute_error": [152.01806640625], - "r2_score": [0.8840255831308607], - "explained_variance": [0.8858505311591299], + "mean_absolute_error": [202.298434], + "mean_squared_error": [74515.108971], + "mean_squared_log_error": [0.004365], + "median_absolute_error": [142.949219], + "r2_score": [0.88462], + "explained_variance": [0.886454], }, dtype="Float64", ) @@ -240,6 +297,7 @@ def test_pipeline_xgbregressor_fit_score_predict(session, penguins_df_default_in ) +@pytest.mark.flaky(retries=2, delay=120) def test_pipeline_random_forest_classifier_fit_score_predict( session, penguins_df_default_index ): @@ -252,26 +310,26 @@ def test_pipeline_random_forest_classifier_fit_score_predict( ) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - train_y = df[["sex"]] - pl.fit(train_X, train_y) + y_train = df[["sex"]] + pl.fit(X_train, y_train) # Check score to ensure the model was fitted - score_result = pl.score(train_X, train_y).compute() + score_result = pl.score(X_train, y_train).to_pandas() score_expected = pd.DataFrame( { - "precision": [0.587673], - "recall": [0.588781], - "accuracy": [0.88024], - "f1_score": [0.587644], - "log_loss": [0.859459], - "roc_auc": [0.971737], + "precision": [0.585505], + "recall": [0.58676], + "accuracy": [0.877246], + "f1_score": [0.585657], + "log_loss": [0.880643], + "roc_auc": [0.970697], }, dtype="Float64", ) @@ -310,7 +368,7 @@ def test_pipeline_random_forest_classifier_fit_score_predict( ) -def test_pipeline_PCA_fit_predict(session, penguins_df_default_index): +def test_pipeline_PCA_fit_score_predict(session, penguins_df_default_index): """Test a supervised model with a minimal preprocessing step""" pl = pipeline.Pipeline( [ @@ -320,14 +378,14 @@ def test_pipeline_PCA_fit_predict(session, penguins_df_default_index): ) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - pl.fit(train_X) + pl.fit(X_train) # predict new labels new_penguins = session.read_pandas( @@ -347,12 +405,27 @@ def test_pipeline_PCA_fit_predict(session, penguins_df_default_index): } ).set_index("tag_number") ) + + # Check score to ensure the model was fitted + score_result = pl.score(new_penguins).to_pandas() + score_expected = pd.DataFrame( + { + "total_explained_variance_ratio": [1.0], + }, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + predictions = pl.predict(new_penguins).to_pandas() expected = pd.DataFrame( { - "principal_component_1": [-1.115259, -1.506141, -1.471174], - "principal_component_2": [-0.074824, 0.69664, 0.406104], - "principal_component_3": [0.500012, -0.544479, 0.075849], + "principal_component_1": [-1.115259, -1.506141, -1.471173], + "principal_component_2": [-0.074825, 0.69664, 0.406103], + "principal_component_3": [0.500013, -0.544479, 0.075849], }, dtype="Float64", index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), @@ -367,7 +440,8 @@ def test_pipeline_PCA_fit_predict(session, penguins_df_default_index): ) -def test_pipeline_standard_scaler_kmeans_fit_predict( +@pytest.mark.flaky(retries=2, delay=120) +def test_pipeline_standard_scaler_kmeans_fit_score_predict( session, penguins_pandas_df_default_index ): """Test an unsupervised model with a non-BQML implementation of StandardScaler""" @@ -380,14 +454,14 @@ def test_pipeline_standard_scaler_kmeans_fit_predict( # kmeans is sensitive to the order with this configuration, so use ordered source data df = session.read_pandas(penguins_pandas_df_default_index).dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - pl.fit(train_X) + pl.fit(X_train) # predict new labels pd_new_penguins = pd.DataFrame.from_dict( @@ -452,6 +526,19 @@ def test_pipeline_standard_scaler_kmeans_fit_predict( pd_new_penguins.index.name = "observation" new_penguins = session.read_pandas(pd_new_penguins) + + # Check score to ensure the model was fitted + score_result = pl.score(new_penguins).to_pandas() + score_expected = pd.DataFrame( + {"davies_bouldin_index": [7.542981], "mean_squared_distance": [94.692409]}, + dtype="Float64", + ) + score_expected = score_expected.reindex(index=score_expected.index.astype("Int64")) + + pd.testing.assert_frame_equal( + score_result, score_expected, check_exact=False, rtol=0.1 + ) + result = pl.predict(new_penguins).to_pandas().sort_index() expected = pd.DataFrame( {"CENTROID_ID": [1, 2, 1, 2, 1, 2]}, @@ -466,7 +553,7 @@ def test_pipeline_standard_scaler_kmeans_fit_predict( def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_index): - """Test a preprocessing step that manages heterogenous data with ColumnTransformer""" + """Test a preprocessing step that manages heterogeneous data with ColumnTransformer""" pl = pipeline.Pipeline( [ ( @@ -491,9 +578,9 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind ) df = penguins_df_default_index.dropna() - train_X = df[["species", "culmen_length_mm", "flipper_length_mm"]] - train_y = df[["body_mass_g"]] - pl.fit(train_X, train_y) + X_train = df[["species", "culmen_length_mm", "flipper_length_mm"]] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) # predict new labels new_penguins = session.read_pandas( @@ -522,3 +609,129 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind pd.testing.assert_frame_equal( predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1 ) + + +def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ( + "transform", + compose.ColumnTransformer( + [ + ( + "ont_hot_encoder", + preprocessing.OneHotEncoder( + drop="most_frequent", + min_frequency=5, + max_categories=100, + ), + "species", + ), + ( + "standard_scaler", + preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ), + ), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[["species", "culmen_length_mm", "flipper_length_mm"]] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_col_transformer", replace=True + ) + + assert isinstance(pl_loaded._transform, compose.ColumnTransformer) + transformers = pl_loaded._transform.transformers_ + assert len(transformers) == 3 + + assert transformers[0][0] == "ont_hot_encoder" + assert isinstance(transformers[0][1], preprocessing.OneHotEncoder) + one_hot_encoder = transformers[0][1] + assert one_hot_encoder.drop == "most_frequent" + assert one_hot_encoder.min_frequency == 5 + assert one_hot_encoder.max_categories == 100 + assert transformers[0][2] == "species" + + assert transformers[1][0] == "standard_scaler" + assert isinstance(transformers[1][1], preprocessing.StandardScaler) + assert transformers[1][2] == "culmen_length_mm" + + assert transformers[2][0] == "standard_scaler" + assert isinstance(transformers[2][1], preprocessing.StandardScaler) + assert transformers[2][2] == "flipper_length_mm" + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False + + +def test_pipeline_standard_scaler_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ("transform", preprocessing.StandardScaler()), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_standard_scaler", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.StandardScaler) + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False + + +def test_pipeline_one_hot_encoder_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ( + "transform", + preprocessing.OneHotEncoder( + drop="most_frequent", min_frequency=5, max_categories=100 + ), + ), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "sex", + "species", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_one_hot_encoder", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.OneHotEncoder) + + one_hot_encoder = pl_loaded._transform + assert one_hot_encoder.drop == "most_frequent" + assert one_hot_encoder.min_frequency == 5 + assert one_hot_encoder.max_categories == 100 + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index 714252c9e00..8033f79c47b 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -295,7 +295,7 @@ def square(x): bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_int64_col = scalars_pandas_df["int64_col"] @@ -339,7 +339,7 @@ def add_one(x): bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(remote_add_one) bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_int64_col = scalars_pandas_df["int64_col"] @@ -392,7 +392,7 @@ def sign(num): bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(remote_sign) bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_int64_col = scalars_pandas_df["int64_col"] @@ -437,7 +437,7 @@ def circumference(radius): bf_float64_col_filtered = bf_float64_col[bf_float64_col_filter] bf_result_col = bf_float64_col_filtered.apply(remote_circumference) bf_result = ( - bf_float64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_float64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_float64_col = scalars_pandas_df["float64_col"] @@ -486,7 +486,7 @@ def find_team(num): bf_float64_col_filtered = bf_float64_col[bf_float64_col_filter] bf_result_col = bf_float64_col_filtered.apply(remote_find_team) bf_result = ( - bf_float64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_float64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_float64_col = scalars_pandas_df["float64_col"] @@ -577,7 +577,9 @@ def inner_test(): bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(remote_add_one) bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_int64_col_filtered.to_frame() + .assign(result=bf_result_col) + .to_pandas() ) pd_int64_col = scalars_pandas_df["int64_col"] @@ -675,7 +677,7 @@ def is_odd(num): bf_int64_col = scalars_df["int64_col"] bf_result_col = bf_int64_col.mask(is_odd_remote) - bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).compute() + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas() pd_int64_col = scalars_pandas_df["int64_col"] pd_result_col = pd_int64_col.mask(is_odd) @@ -718,7 +720,7 @@ def is_odd(num): # for now filter out the nulls and test the rest bf_int64_col = scalars_df["int64_col"] bf_result_col = bf_int64_col[bf_int64_col.notnull()].mask(is_odd_remote, -1) - bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).compute() + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas() pd_int64_col = scalars_pandas_df["int64_col"] pd_result_col = pd_int64_col[pd_int64_col.notnull()].mask(is_odd, -1) @@ -754,7 +756,7 @@ def test_remote_udf_lambda( bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(add_one_lambda_remote) bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() ) pd_int64_col = scalars_pandas_df["int64_col"] diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 5c3abf25dcc..8325cec50aa 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -21,11 +21,17 @@ import pandas as pd import pytest -from bigframes.ml import core, forecasting, imported, llm -import bigframes.ml.cluster -import bigframes.ml.core -import bigframes.ml.ensemble -import bigframes.ml.linear_model +import bigframes +from bigframes.ml import ( + cluster, + core, + decomposition, + ensemble, + forecasting, + imported, + linear_model, + llm, +) @pytest.fixture(scope="session") @@ -34,17 +40,15 @@ def ml_connection() -> str: @pytest.fixture(scope="session") -def penguins_bqml_linear_model( - session, penguins_linear_model_name -) -> bigframes.ml.core.BqmlModel: +def penguins_bqml_linear_model(session, penguins_linear_model_name) -> core.BqmlModel: model = session.bqclient.get_model(penguins_linear_model_name) - return bigframes.ml.core.BqmlModel(session, model) + return core.BqmlModel(session, model) @pytest.fixture(scope="function") def ephemera_penguins_bqml_linear_model( penguins_bqml_linear_model, -) -> bigframes.ml.linear_model.LinearRegression: +) -> linear_model.LinearRegression: model = penguins_bqml_linear_model return model.copy( f"{model._model.project}.{model._model.dataset_id}.{uuid.uuid4().hex}" @@ -54,18 +58,18 @@ def ephemera_penguins_bqml_linear_model( @pytest.fixture(scope="session") def penguins_linear_model( session, penguins_linear_model_name: str -) -> bigframes.ml.linear_model.LinearRegression: +) -> linear_model.LinearRegression: return cast( - bigframes.ml.linear_model.LinearRegression, + linear_model.LinearRegression, session.read_gbq_model(penguins_linear_model_name), ) @pytest.fixture(scope="function") def ephemera_penguins_linear_model( - ephemera_penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, -) -> bigframes.ml.linear_model.LinearRegression: - bf_model = bigframes.ml.linear_model.LinearRegression() + ephemera_penguins_bqml_linear_model: core.BqmlModel, +) -> linear_model.LinearRegression: + bf_model = linear_model.LinearRegression() bf_model._bqml_model = ephemera_penguins_bqml_linear_model return bf_model @@ -73,9 +77,9 @@ def ephemera_penguins_linear_model( @pytest.fixture(scope="session") def penguins_logistic_model( session, penguins_logistic_model_name -) -> bigframes.ml.linear_model.LogisticRegression: +) -> linear_model.LogisticRegression: return cast( - bigframes.ml.linear_model.LogisticRegression, + linear_model.LogisticRegression, session.read_gbq_model(penguins_logistic_model_name), ) @@ -83,9 +87,9 @@ def penguins_logistic_model( @pytest.fixture(scope="session") def penguins_xgbregressor_model( session, penguins_xgbregressor_model_name -) -> bigframes.ml.ensemble.XGBRegressor: +) -> ensemble.XGBRegressor: return cast( - bigframes.ml.ensemble.XGBRegressor, + ensemble.XGBRegressor, session.read_gbq_model(penguins_xgbregressor_model_name), ) @@ -93,9 +97,9 @@ def penguins_xgbregressor_model( @pytest.fixture(scope="session") def penguins_xgbclassifier_model( session, penguins_xgbclassifier_model_name -) -> bigframes.ml.ensemble.XGBClassifier: +) -> ensemble.XGBClassifier: return cast( - bigframes.ml.ensemble.XGBClassifier, + ensemble.XGBClassifier, session.read_gbq_model(penguins_xgbclassifier_model_name), ) @@ -103,9 +107,9 @@ def penguins_xgbclassifier_model( @pytest.fixture(scope="session") def penguins_randomforest_regressor_model( session, penguins_randomforest_regressor_model_name -) -> bigframes.ml.ensemble.RandomForestRegressor: +) -> ensemble.RandomForestRegressor: return cast( - bigframes.ml.ensemble.RandomForestRegressor, + ensemble.RandomForestRegressor, session.read_gbq_model(penguins_randomforest_regressor_model_name), ) @@ -113,9 +117,9 @@ def penguins_randomforest_regressor_model( @pytest.fixture(scope="session") def penguins_randomforest_classifier_model( session, penguins_randomforest_classifier_model_name -) -> bigframes.ml.ensemble.RandomForestClassifier: +) -> ensemble.RandomForestClassifier: return cast( - bigframes.ml.ensemble.RandomForestClassifier, + ensemble.RandomForestClassifier, session.read_gbq_model(penguins_randomforest_classifier_model_name), ) @@ -123,7 +127,7 @@ def penguins_randomforest_classifier_model( @pytest.fixture(scope="session") def penguins_kmeans_model( session: bigframes.Session, dataset_id_permanent, penguins_table_id -) -> bigframes.ml.cluster.KMeans: +) -> cluster.KMeans: """Provides a pretrained model as a test fixture that is cached across test runs. This lets us run system tests without having to wait for a model.fit(...)""" sql = f""" @@ -154,7 +158,7 @@ def penguins_kmeans_model( @pytest.fixture(scope="session") def penguins_pca_model( session: bigframes.Session, dataset_id_permanent, penguins_table_id -) -> bigframes.ml.decomposition.PCA: +) -> decomposition.PCA: # TODO(yunmengxie): Create a shared method to get different types of pretrained models. sql = f""" @@ -248,8 +252,10 @@ def ephemera_palm2_text_generator_model( @pytest.fixture(scope="session") def palm2_embedding_generator_model( session, ml_connection -) -> llm.PaLM2EmbeddingGenerator: - return llm.PaLM2EmbeddingGenerator(session=session, connection_name=ml_connection) +) -> llm.PaLM2TextEmbeddingGenerator: + return llm.PaLM2TextEmbeddingGenerator( + session=session, connection_name=ml_connection + ) @pytest.fixture(scope="session") @@ -257,7 +263,7 @@ def time_series_bqml_arima_plus_model( session, time_series_arima_plus_model_name ) -> core.BqmlModel: model = session.bqclient.get_model(time_series_arima_plus_model_name) - return bigframes.ml.core.BqmlModel(session, model) + return core.BqmlModel(session, model) @pytest.fixture(scope="session") @@ -287,8 +293,8 @@ def ephemera_imported_tensorflow_model(session) -> imported.TensorFlowModel: @pytest.fixture(scope="session") -def imported_onnx_model(session) -> imported.OnnxModel: - return imported.OnnxModel( +def imported_onnx_model(session) -> imported.ONNXModel: + return imported.ONNXModel( session=session, model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx", ) diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py index 4aefc5fa697..a003cd1ec1b 100644 --- a/tests/system/small/ml/test_cluster.py +++ b/tests/system/small/ml/test_cluster.py @@ -12,68 +12,81 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pandas +import pandas as pd -import bigframes.ml.cluster +from bigframes.ml import cluster from tests.system.utils import assert_pandas_df_equal_ignore_ordering - -def test_model_predict(session, penguins_kmeans_model: bigframes.ml.cluster.KMeans): - pd_new_penguins = pandas.DataFrame.from_dict( - { - "test1": { - "species": "Adelie Penguin (Pygoscelis adeliae)", - "island": "Dream", - "culmen_length_mm": 37.5, - "culmen_depth_mm": 18.5, - "flipper_length_mm": 199, - "body_mass_g": 4475, - "sex": "MALE", - }, - "test2": { - "species": "Chinstrap penguin (Pygoscelis antarctica)", - "island": "Dream", - "culmen_length_mm": 55.8, - "culmen_depth_mm": 19.8, - "flipper_length_mm": 207, - "body_mass_g": 4000, - "sex": "MALE", - }, - "test3": { - "species": "Adelie Penguin (Pygoscelis adeliae)", - "island": "Biscoe", - "culmen_length_mm": 39.7, - "culmen_depth_mm": 18.9, - "flipper_length_mm": 184, - "body_mass_g": 3550, - "sex": "MALE", - }, - "test4": { - "species": "Gentoo penguin (Pygoscelis papua)", - "island": "Biscoe", - "culmen_length_mm": 43.8, - "culmen_depth_mm": 13.9, - "flipper_length_mm": 208, - "body_mass_g": 4300, - "sex": "FEMALE", - }, +_PD_NEW_PENGUINS = pd.DataFrame.from_dict( + { + "test1": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Dream", + "culmen_length_mm": 37.5, + "culmen_depth_mm": 18.5, + "flipper_length_mm": 199, + "body_mass_g": 4475, + "sex": "MALE", }, - orient="index", - ) - pd_new_penguins.index.name = "observation" + "test2": { + "species": "Chinstrap penguin (Pygoscelis antarctica)", + "island": "Dream", + "culmen_length_mm": 55.8, + "culmen_depth_mm": 19.8, + "flipper_length_mm": 207, + "body_mass_g": 4000, + "sex": "MALE", + }, + "test3": { + "species": "Adelie Penguin (Pygoscelis adeliae)", + "island": "Biscoe", + "culmen_length_mm": 39.7, + "culmen_depth_mm": 18.9, + "flipper_length_mm": 184, + "body_mass_g": 3550, + "sex": "MALE", + }, + "test4": { + "species": "Gentoo penguin (Pygoscelis papua)", + "island": "Biscoe", + "culmen_length_mm": 43.8, + "culmen_depth_mm": 13.9, + "flipper_length_mm": 208, + "body_mass_g": 4300, + "sex": "FEMALE", + }, + }, + orient="index", +) - new_penguins = session.read_pandas(pd_new_penguins) - result = penguins_kmeans_model.predict(new_penguins).compute() - expected = pandas.DataFrame( + +def test_kmeans_predict(session, penguins_kmeans_model: cluster.KMeans): + new_penguins = session.read_pandas(_PD_NEW_PENGUINS) + result = penguins_kmeans_model.predict(new_penguins).to_pandas() + expected = pd.DataFrame( {"CENTROID_ID": [2, 3, 1, 2]}, dtype="Int64", - index=pandas.Index( - ["test1", "test2", "test3", "test4"], dtype="string[pyarrow]" - ), + index=pd.Index(["test1", "test2", "test3", "test4"], dtype="string[pyarrow]"), ) - expected.index.name = "observation" assert_pandas_df_equal_ignore_ordering(result, expected) +def test_kmeans_score(session, penguins_kmeans_model: cluster.KMeans): + new_penguins = session.read_pandas(_PD_NEW_PENGUINS) + result = penguins_kmeans_model.score(new_penguins).to_pandas() + expected = pd.DataFrame( + {"davies_bouldin_index": [1.523606], "mean_squared_distance": [1.965944]}, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + def test_loaded_config(penguins_kmeans_model): assert penguins_kmeans_model.n_clusters == 3 diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index d63e401dd7a..8b864d9b55c 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -27,7 +27,7 @@ def test_model_eval( penguins_bqml_linear_model, ): - result = penguins_bqml_linear_model.evaluate().compute() + result = penguins_bqml_linear_model.evaluate().to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [227.01223], @@ -52,7 +52,7 @@ def test_model_eval( def test_model_eval_with_data(penguins_bqml_linear_model, penguins_df_default_index): result = penguins_bqml_linear_model.evaluate( penguins_df_default_index.dropna() - ).compute() + ).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [225.817334], @@ -77,7 +77,7 @@ def test_model_eval_with_data(penguins_bqml_linear_model, penguins_df_default_in def test_model_predict( penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, new_penguins_df ): - predictions = penguins_bqml_linear_model.predict(new_penguins_df).compute() + predictions = penguins_bqml_linear_model.predict(new_penguins_df).to_pandas() expected = pd.DataFrame( {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, dtype="Float64", @@ -105,7 +105,7 @@ def test_model_predict_with_unnamed_index( new_penguins_df[new_penguins_df.tag_number != 1672], ) - predictions = penguins_bqml_linear_model.predict(new_penguins_df).compute() + predictions = penguins_bqml_linear_model.predict(new_penguins_df).to_pandas() expected = pd.DataFrame( {"predicted_body_mass_g": [4030.1, 3177.9]}, @@ -132,7 +132,7 @@ def test_model_generate_text( } df = bqml_palm2_text_generator_model.generate_text( llm_text_df, options=options - ).compute() + ).to_pandas() TestCase().assertSequenceEqual(df.shape, (3, 4)) TestCase().assertSequenceEqual( @@ -150,7 +150,7 @@ def test_model_generate_text( def test_model_forecast(time_series_bqml_arima_plus_model: bigframes.ml.core.BqmlModel): utc = pytz.utc - forecast = time_series_bqml_arima_plus_model.forecast().compute()[ + forecast = time_series_bqml_arima_plus_model.forecast().to_pandas()[ ["forecast_timestamp", "forecast_value"] ] expected = pd.DataFrame( diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index 7e0421129eb..01d52077505 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -12,44 +12,59 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pandas - -import bigframes.ml.decomposition - - -def test_model_predict(session, penguins_pca_model: bigframes.ml.decomposition.PCA): - new_penguins = session.read_pandas( - pandas.DataFrame( - { - "tag_number": [1633, 1672, 1690], - "species": [ - "Adelie Penguin (Pygoscelis adeliae)", - "Gentoo penguin (Pygoscelis papua)", - "Adelie Penguin (Pygoscelis adeliae)", - ], - "island": ["Dream", "Biscoe", "Torgersen"], - "culmen_length_mm": [37.8, 46.5, 41.1], - "culmen_depth_mm": [18.1, 14.8, 18.6], - "flipper_length_mm": [193.0, 217.0, 189.0], - "body_mass_g": [3750.0, 5200.0, 3325.0], - "sex": ["MALE", "FEMALE", "MALE"], - } - ).set_index("tag_number") - ) +import pandas as pd + +from bigframes.ml import decomposition + +_PD_NEW_PENGUINS = pd.DataFrame( + { + "tag_number": [1633, 1672, 1690], + "species": [ + "Adelie Penguin (Pygoscelis adeliae)", + "Gentoo penguin (Pygoscelis papua)", + "Adelie Penguin (Pygoscelis adeliae)", + ], + "island": ["Dream", "Biscoe", "Torgersen"], + "culmen_length_mm": [37.8, 46.5, 41.1], + "culmen_depth_mm": [18.1, 14.8, 18.6], + "flipper_length_mm": [193.0, 217.0, 189.0], + "body_mass_g": [3750.0, 5200.0, 3325.0], + "sex": ["MALE", "FEMALE", "MALE"], + } +).set_index("tag_number") + - predictions = penguins_pca_model.predict(new_penguins).compute() - expected = pandas.DataFrame( +def test_pca_predict(session, penguins_pca_model: decomposition.PCA): + new_penguins = session.read_pandas(_PD_NEW_PENGUINS) + + predictions = penguins_pca_model.predict(new_penguins).to_pandas() + expected = pd.DataFrame( { "principal_component_1": [-1.459, 2.258, -1.685], "principal_component_2": [-1.120, -1.351, -0.874], "principal_component_3": [-0.646, 0.443, -0.704], }, dtype="Float64", - index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), ) - pandas.testing.assert_frame_equal( + pd.testing.assert_frame_equal( predictions.sort_index(), expected, check_exact=False, rtol=0.1, ) + + +def test_pca_score(session, penguins_pca_model: decomposition.PCA): + result = penguins_pca_model.score().to_pandas() + expected = pd.DataFrame( + {"total_explained_variance_ratio": [0.812383]}, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + ) diff --git a/tests/system/small/ml/test_ensemble.py b/tests/system/small/ml/test_ensemble.py index 61c60c249f2..fde3cc431e1 100644 --- a/tests/system/small/ml/test_ensemble.py +++ b/tests/system/small/ml/test_ensemble.py @@ -36,7 +36,44 @@ def test_xgbregressor_model_score( ] ] test_y = df[["sex"]] - result = penguins_xgbregressor_model.score(test_X, test_y).compute() + result = penguins_xgbregressor_model.score(test_X, test_y).to_pandas() + expected = pandas.DataFrame( + { + "mean_absolute_error": [108.77582], + "mean_squared_error": [20943.272738], + "mean_squared_log_error": [0.00135], + "median_absolute_error": [86.313477], + "r2_score": [0.967571], + "explained_variance": [0.967609], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_xgbregressor_model_score_series( + penguins_xgbregressor_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df["sex"] + result = penguins_xgbregressor_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [108.77582], @@ -61,7 +98,7 @@ def test_xgbregressor_model_score( def test_xgbregressor_model_predict( penguins_xgbregressor_model: bigframes.ml.ensemble.XGBRegressor, new_penguins_df ): - result = penguins_xgbregressor_model.predict(new_penguins_df).compute() + result = penguins_xgbregressor_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_body_mass_g": ["4293.1538089", "3410.0271", "3357.944"]}, dtype="Float64", @@ -94,7 +131,7 @@ def test_to_gbq_saved_xgbregressor_model_scores( ] ] test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [109.016973], @@ -139,7 +176,35 @@ def test_xgbclassifier_model_score( ] ] test_y = df[["sex"]] - result = penguins_xgbclassifier_model.score(test_X, test_y).compute() + result = penguins_xgbclassifier_model.score(test_X, test_y).to_pandas() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + +def test_xgbclassifier_model_score_series( + penguins_xgbclassifier_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df["sex"] + result = penguins_xgbclassifier_model.score(test_X, test_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -155,7 +220,7 @@ def test_xgbclassifier_model_score( def test_xgbclassifier_model_predict( penguins_xgbclassifier_model: bigframes.ml.ensemble.XGBClassifier, new_penguins_df ): - result = penguins_xgbclassifier_model.predict(new_penguins_df).compute() + result = penguins_xgbclassifier_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, dtype="string[pyarrow]", @@ -188,7 +253,7 @@ def test_to_gbq_saved_xgbclassifier_model_scores( ] ] test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "precision": [1.0], @@ -235,7 +300,44 @@ def test_randomforestregressor_model_score( ] ] test_y = df[["sex"]] - result = penguins_randomforest_regressor_model.score(test_X, test_y).compute() + result = penguins_randomforest_regressor_model.score(test_X, test_y).to_pandas() + expected = pandas.DataFrame( + { + "mean_absolute_error": [317.031042], + "mean_squared_error": [159713.053504], + "mean_squared_log_error": [0.008449], + "median_absolute_error": [258.385742], + "r2_score": [0.752698], + "explained_variance": [0.756173], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigFramese + check_index_type=False, + ) + + +def test_randomforestregressor_model_score_series( + penguins_randomforest_regressor_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df["sex"] + result = penguins_randomforest_regressor_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [317.031042], @@ -261,7 +363,7 @@ def test_randomforestregressor_model_predict( penguins_randomforest_regressor_model: bigframes.ml.ensemble.RandomForestRegressor, new_penguins_df, ): - result = penguins_randomforest_regressor_model.predict(new_penguins_df).compute() + result = penguins_randomforest_regressor_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_body_mass_g": ["3897.341797", "3458.385742", "3458.385742"]}, dtype="Float64", @@ -294,7 +396,7 @@ def test_to_gbq_saved_randomforestregressor_model_scores( ] ] test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [319.239235], @@ -343,7 +445,35 @@ def test_randomforestclassifier_model_score( ] ] test_y = df[["sex"]] - result = penguins_randomforest_classifier_model.score(test_X, test_y).compute() + result = penguins_randomforest_classifier_model.score(test_X, test_y).to_pandas() + TestCase().assertSequenceEqual(result.shape, (1, 6)) + for col_name in [ + "precision", + "recall", + "accuracy", + "f1_score", + "log_loss", + "roc_auc", + ]: + assert col_name in result.columns + + +def test_randomforestclassifier_model_score_series( + penguins_randomforest_classifier_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df["sex"] + result = penguins_randomforest_classifier_model.score(test_X, test_y).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -360,7 +490,7 @@ def test_randomforestclassifier_model_predict( penguins_randomforest_classifier_model: bigframes.ml.ensemble.RandomForestClassifier, new_penguins_df, ): - result = penguins_randomforest_classifier_model.predict(new_penguins_df).compute() + result = penguins_randomforest_classifier_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, dtype="string[pyarrow]", @@ -393,7 +523,7 @@ def test_to_gbq_saved_randomforestclassifier_model_scores( ] ] test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "precision": [0.636746], diff --git a/tests/system/small/ml/test_forecasting.py b/tests/system/small/ml/test_forecasting.py index 153299f5989..cb27dd388c3 100644 --- a/tests/system/small/ml/test_forecasting.py +++ b/tests/system/small/ml/test_forecasting.py @@ -21,7 +21,7 @@ def test_model_predict(time_series_arima_plus_model): utc = pytz.utc - predictions = time_series_arima_plus_model.predict().compute() + predictions = time_series_arima_plus_model.predict().to_pandas() expected = pd.DataFrame( { "forecast_timestamp": [ @@ -42,3 +42,47 @@ def test_model_predict(time_series_arima_plus_model): rtol=0.1, check_index_type=False, ) + + +def test_model_score(time_series_arima_plus_model, new_time_series_df): + result = time_series_arima_plus_model.score( + new_time_series_df[["parsed_date"]], new_time_series_df[["total_visits"]] + ).to_pandas() + expected = pd.DataFrame( + { + "mean_absolute_error": [154.742547], + "mean_squared_error": [26844.868855], + "root_mean_squared_error": [163.844038], + "mean_absolute_percentage_error": [6.189702], + "symmetric_mean_absolute_percentage_error": [6.097155], + }, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + rtol=0.1, + check_index_type=False, + ) + + +def test_model_score_series(time_series_arima_plus_model, new_time_series_df): + result = time_series_arima_plus_model.score( + new_time_series_df["parsed_date"], new_time_series_df["total_visits"] + ).to_pandas() + expected = pd.DataFrame( + { + "mean_absolute_error": [154.742547], + "mean_squared_error": [26844.868855], + "root_mean_squared_error": [163.844038], + "mean_absolute_percentage_error": [6.189702], + "symmetric_mean_absolute_percentage_error": [6.097155], + }, + dtype="Float64", + ) + pd.testing.assert_frame_equal( + result, + expected, + rtol=0.1, + check_index_type=False, + ) diff --git a/tests/system/small/ml/test_imported.py b/tests/system/small/ml/test_imported.py index b2a9abc46e7..6274ab12450 100644 --- a/tests/system/small/ml/test_imported.py +++ b/tests/system/small/ml/test_imported.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import google.api_core.exceptions import numpy as np import pandas as pd +import pytest + +from bigframes.ml import imported def test_tensorflow_create_model(imported_tensorflow_model): @@ -23,7 +27,7 @@ def test_tensorflow_create_model(imported_tensorflow_model): def test_tensorflow_model_predict(imported_tensorflow_model, llm_text_df): df = llm_text_df.rename(columns={"prompt": "input"}) - result = imported_tensorflow_model.predict(df).compute() + result = imported_tensorflow_model.predict(df).to_pandas() # The values are non-human-readable. As they are a dense layer of Neural Network. # And since it is pretrained and imported, the model is a opaque-box. # We may want to switch to better test model and cases. @@ -44,13 +48,21 @@ def test_tensorflow_model_predict(imported_tensorflow_model, llm_text_df): ) +def test_tensorflow_model_to_gbq( + imported_tensorflow_model: imported.TensorFlowModel, dataset_id: str +): + imported_tensorflow_model.to_gbq(f"{dataset_id}.test_tf_model", replace=True) + with pytest.raises(google.api_core.exceptions.Conflict): + imported_tensorflow_model.to_gbq(f"{dataset_id}.test_tf_model") + + def test_onnx_create_model(imported_onnx_model): # Model creation doesn't return error assert imported_onnx_model is not None def test_onnx_model_predict(imported_onnx_model, onnx_iris_df): - result = imported_onnx_model.predict(onnx_iris_df).compute() + result = imported_onnx_model.predict(onnx_iris_df).to_pandas() value1 = np.array([0.9999993443489075, 0.0, 0.0]) value2 = np.array([0.0, 0.0, 0.9999993443489075]) expected = pd.DataFrame( @@ -66,3 +78,9 @@ def test_onnx_model_predict(imported_onnx_model, onnx_iris_df): check_exact=False, atol=0.1, ) + + +def test_onnx_model_to_gbq(imported_onnx_model: imported.ONNXModel, dataset_id: str): + imported_onnx_model.to_gbq(f"{dataset_id}.test_onnx_model", replace=True) + with pytest.raises(google.api_core.exceptions.Conflict): + imported_onnx_model.to_gbq(f"{dataset_id}.test_onnx_model") diff --git a/tests/system/small/ml/test_linear_model.py b/tests/system/small/ml/test_linear_model.py index d5d6957eec3..bbb7e2820c3 100644 --- a/tests/system/small/ml/test_linear_model.py +++ b/tests/system/small/ml/test_linear_model.py @@ -30,7 +30,44 @@ def test_linear_reg_model_score(penguins_linear_model, penguins_df_default_index ] ] test_y = df[["body_mass_g"]] - result = penguins_linear_model.score(test_X, test_y).compute() + result = penguins_linear_model.score(test_X, test_y).to_pandas() + expected = pandas.DataFrame( + { + "mean_absolute_error": [225.817334], + "mean_squared_error": [80540.705944], + "mean_squared_log_error": [0.004972], + "median_absolute_error": [173.080816], + "r2_score": [0.87529], + "explained_variance": [0.87529], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_linear_reg_model_score_series( + penguins_linear_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + test_y = df["body_mass_g"] + result = penguins_linear_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [225.817334], @@ -53,7 +90,7 @@ def test_linear_reg_model_score(penguins_linear_model, penguins_df_default_index def test_linear_reg_model_predict(penguins_linear_model, new_penguins_df): - predictions = penguins_linear_model.predict(new_penguins_df).compute() + predictions = penguins_linear_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, dtype="Float64", @@ -85,7 +122,7 @@ def test_to_gbq_saved_linear_reg_model_scores( ] ] test_y = df[["body_mass_g"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [227.01223], @@ -126,7 +163,44 @@ def test_logistic_model_score(penguins_logistic_model, penguins_df_default_index ] ] test_y = df[["sex"]] - result = penguins_logistic_model.score(test_X, test_y).compute() + result = penguins_logistic_model.score(test_X, test_y).to_pandas() + expected = pandas.DataFrame( + { + "precision": [0.616753], + "recall": [0.618615], + "accuracy": [0.92515], + "f1_score": [0.617681], + "log_loss": [1.498832], + "roc_auc": [0.975807], + }, + dtype="Float64", + ) + pandas.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + ) + + +def test_logistic_model_score_series( + penguins_logistic_model, penguins_df_default_index +): + df = penguins_df_default_index.dropna() + test_X = df[ + [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + ] + ] + test_y = df["sex"] + result = penguins_logistic_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "precision": [0.616753], @@ -149,7 +223,7 @@ def test_logistic_model_score(penguins_logistic_model, penguins_df_default_index def test_logsitic_model_predict(penguins_logistic_model, new_penguins_df): - predictions = penguins_logistic_model.predict(new_penguins_df).compute() + predictions = penguins_logistic_model.predict(new_penguins_df).to_pandas() expected = pandas.DataFrame( {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, dtype="string[pyarrow]", @@ -163,7 +237,7 @@ def test_logsitic_model_predict(penguins_logistic_model, new_penguins_df): ) -def test_to_gbq_saved_logsitic_model_score( +def test_logsitic_model_to_gbq_saved_score( penguins_logistic_model, dataset_id, penguins_df_default_index ): saved_model = penguins_logistic_model.to_gbq( @@ -181,7 +255,7 @@ def test_to_gbq_saved_logsitic_model_score( ] ] test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).compute() + result = saved_model.score(test_X, test_y).to_pandas() expected = pandas.DataFrame( { "precision": [0.616753], @@ -203,7 +277,7 @@ def test_to_gbq_saved_logsitic_model_score( ) -def test_to_logistic_model_gbq_replace(penguins_logistic_model, dataset_id): +def test_logistic_model_to_gbq_replace(penguins_logistic_model, dataset_id): penguins_logistic_model.to_gbq(f"{dataset_id}.test_penguins_model", replace=True) with pytest.raises(google.api_core.exceptions.Conflict): penguins_logistic_model.to_gbq(f"{dataset_id}.test_penguins_model") diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index b579e754e5a..74356c81e1b 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -25,7 +25,17 @@ def test_create_text_generator_model(palm2_text_generator_model): def test_text_generator_predict_default_params_success( palm2_text_generator_model, llm_text_df ): - df = palm2_text_generator_model.predict(llm_text_df).compute() + df = palm2_text_generator_model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +def test_text_generator_predict_series_default_params_success( + palm2_text_generator_model, llm_text_df +): + df = palm2_text_generator_model.predict(llm_text_df["prompt"]).to_pandas() TestCase().assertSequenceEqual(df.shape, (3, 1)) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] @@ -36,7 +46,7 @@ def test_text_generator_predict_arbitrary_col_label_success( palm2_text_generator_model, llm_text_df ): llm_text_df = llm_text_df.rename(columns={"prompt": "arbitrary"}) - df = palm2_text_generator_model.predict(llm_text_df).compute() + df = palm2_text_generator_model.predict(llm_text_df).to_pandas() TestCase().assertSequenceEqual(df.shape, (3, 1)) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] @@ -48,7 +58,7 @@ def test_text_generator_predict_with_params_success( ): df = palm2_text_generator_model.predict( llm_text_df, temperature=0.5, max_output_tokens=100, top_k=20, top_p=0.5 - ).compute() + ).to_pandas() TestCase().assertSequenceEqual(df.shape, (3, 1)) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] @@ -63,10 +73,22 @@ def test_create_embedding_generator_model(palm2_embedding_generator_model): def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df ): - df = palm2_embedding_generator_model.predict(llm_text_df).compute() + df = palm2_embedding_generator_model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "text_embedding" in df.columns + series = df["text_embedding"] + value = series[0] + assert isinstance(value, np.ndarray) + assert value.size == 768 + + +def test_embedding_generator_predict_series_success( + palm2_embedding_generator_model, llm_text_df +): + df = palm2_embedding_generator_model.predict(llm_text_df["prompt"]).to_pandas() TestCase().assertSequenceEqual(df.shape, (3, 1)) - assert "ml_embed_text_embedding" in df.columns - series = df["ml_embed_text_embedding"] + assert "text_embedding" in df.columns + series = df["text_embedding"] value = series[0] assert isinstance(value, np.ndarray) assert value.size == 768 diff --git a/tests/system/small/ml/test_metrics.py b/tests/system/small/ml/test_metrics.py index 5b4c486b570..b40982e2829 100644 --- a/tests/system/small/ml/test_metrics.py +++ b/tests/system/small/ml/test_metrics.py @@ -23,10 +23,20 @@ def test_r2_score_perfect_fit(session): - pd_df = pd.DataFrame({"y_true": [1, 7, 3, 2, 5], "y_pred": [1, 7, 3, 2, 5]}) + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [1, 7, 3, 2, 5], + "y_pred_arbitrary_name": [1, 7, 3, 2, 5], + } + ) df = session.read_pandas(pd_df) - assert bigframes.ml.metrics.r2_score(df[["y_true"]], df[["y_pred"]]) == 1.0 + assert ( + bigframes.ml.metrics.r2_score( + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]] + ) + == 1.0 + ) def test_r2_score_bad_fit(session): @@ -67,11 +77,28 @@ def test_r2_score_ok_fit_matches_sklearn(session): assert math.isclose(bf_result, sklearn_result) -def test_accuracy_score_perfect_fit(session): +def test_r2_score_series(session): pd_df = pd.DataFrame({"y_true": [1, 7, 3, 2, 5], "y_pred": [1, 7, 3, 2, 5]}) df = session.read_pandas(pd_df) - assert bigframes.ml.metrics.accuracy_score(df[["y_true"]], df[["y_pred"]]) == 1.0 + assert bigframes.ml.metrics.r2_score(df["y_true"], df["y_pred"]) == 1.0 + + +def test_accuracy_score_perfect_fit(session): + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [1, 7, 3, 2, 5], + "y_pred_arbitrary_name": [1, 7, 3, 2, 5], + } + ) + + df = session.read_pandas(pd_df) + assert ( + bigframes.ml.metrics.accuracy_score( + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]] + ) + == 1.0 + ) def test_accuracy_score_bad_fit(session): @@ -93,7 +120,6 @@ def test_accuracy_score_not_normailze(session): ) -@pytest.mark.skipif(sklearn_metrics is None, reason="requires sklearn") def test_accuracy_score_fit_matches_sklearn(session): pd_df = pd.DataFrame({"y_true": [1, 2, 3, 4, 5], "y_pred": [2, 3, 4, 3, 6]}) @@ -105,22 +131,42 @@ def test_accuracy_score_fit_matches_sklearn(session): assert math.isclose(bf_result, sklearn_result) +def test_accuracy_score_series(session): + pd_df = pd.DataFrame({"y_true": [1, 7, 3, 2, 5], "y_pred": [1, 7, 3, 2, 5]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.accuracy_score(df["y_true"], df["y_pred"]) == 1.0 + + def test_roc_curve_binary_classification_prediction_returns_expected(session): pd_df = pd.DataFrame( { - "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], - "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + "y_true_arbitrary_name": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score_arbitrary_name": [ + 0.1, + 0.4, + 0.35, + 0.8, + 0.65, + 0.9, + 0.5, + 0.3, + 0.6, + 0.45, + ], } ) df = session.read_pandas(pd_df) fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( - df[["y_true"]], df[["y_score"]], drop_intermediate=False + df[["y_true_arbitrary_name"]], + df[["y_score_arbitrary_name"]], + drop_intermediate=False, ) - pd_fpr = fpr.compute() - pd_tpr = tpr.compute() - pd_thresholds = thresholds.compute() + pd_fpr = fpr.to_pandas() + pd_tpr = tpr.to_pandas() + pd_thresholds = thresholds.to_pandas() pd.testing.assert_series_equal( # skip testing the first value, as it is redundant and inconsistent across sklearn versions @@ -181,9 +227,9 @@ def test_roc_curve_binary_classification_prediction_matches_sklearn(session): ) # sklearn returns float64 np arrays - np_fpr = fpr.compute().astype("float64").array - np_tpr = tpr.compute().astype("float64").array - np_thresholds = thresholds.compute().astype("float64").array + np_fpr = fpr.to_pandas().astype("float64").array + np_tpr = tpr.to_pandas().astype("float64").array + np_thresholds = thresholds.to_pandas().astype("float64").array np.testing.assert_array_equal( # skip testing the first value, as it is redundant and inconsistent across sklearn versions @@ -217,9 +263,9 @@ def test_roc_curve_binary_classification_decision_returns_expected(session): df[["y_true"]], df[["y_score"]], drop_intermediate=False ) - pd_fpr = fpr.compute() - pd_tpr = tpr.compute() - pd_thresholds = thresholds.compute() + pd_fpr = fpr.to_pandas() + pd_tpr = tpr.to_pandas() + pd_thresholds = thresholds.to_pandas() pd.testing.assert_series_equal( # skip testing the first value, as it is redundant and inconsistent across sklearn versions @@ -276,9 +322,9 @@ def test_roc_curve_binary_classification_decision_matches_sklearn(session): ) # sklearn returns float64 np arrays - np_fpr = fpr.compute().astype("float64").array - np_tpr = tpr.compute().astype("float64").array - np_thresholds = thresholds.compute().astype("float64").array + np_fpr = fpr.to_pandas().astype("float64").array + np_tpr = tpr.to_pandas().astype("float64").array + np_thresholds = thresholds.to_pandas().astype("float64").array np.testing.assert_array_equal( # skip testing the first value, as it is redundant and inconsistent across sklearn versions @@ -295,7 +341,7 @@ def test_roc_curve_binary_classification_decision_matches_sklearn(session): ) -def test_roc_auc_score_returns_expected(session): +def test_roc_curve_binary_classification_prediction_series(session): pd_df = pd.DataFrame( { "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], @@ -304,7 +350,79 @@ def test_roc_auc_score_returns_expected(session): ) df = session.read_pandas(pd_df) - score = bigframes.ml.metrics.roc_auc_score(df[["y_true"]], df[["y_score"]]) + fpr, tpr, thresholds = bigframes.ml.metrics.roc_curve( + df["y_true"], df["y_score"], drop_intermediate=False + ) + + pd_fpr = fpr.to_pandas() + pd_tpr = tpr.to_pandas() + pd_thresholds = thresholds.to_pandas() + + pd.testing.assert_series_equal( + # skip testing the first value, as it is redundant and inconsistent across sklearn versions + pd_thresholds[1:], + pd.Series( + [0.9, 0.8, 0.65, 0.6, 0.5, 0.45, 0.4, 0.35, 0.3, 0.1], + dtype="Float64", + name="thresholds", + ), + check_index=False, + ) + pd.testing.assert_series_equal( + pd_fpr, + pd.Series( + [0.0, 0.0, 0.0, 0.25, 0.25, 0.5, 0.5, 0.75, 0.75, 0.75, 1.0], + dtype="Float64", + name="fpr", + ), + check_index_type=False, + ) + pd.testing.assert_series_equal( + pd_tpr, + pd.Series( + [ + 0.0, + 0.16666667, + 0.33333333, + 0.33333333, + 0.5, + 0.5, + 0.66666667, + 0.66666667, + 0.83333333, + 1.0, + 1.0, + ], + dtype="Float64", + name="tpr", + ), + check_index_type=False, + ) + + +def test_roc_auc_score_returns_expected(session): + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score_arbitrary_name": [ + 0.1, + 0.4, + 0.35, + 0.8, + 0.65, + 0.9, + 0.5, + 0.3, + 0.6, + 0.45, + ], + } + ) + + df = session.read_pandas(pd_df) + score = bigframes.ml.metrics.roc_auc_score( + df[["y_true_arbitrary_name"]], df[["y_score_arbitrary_name"]] + ) assert score == 0.625 @@ -326,11 +444,25 @@ def test_roc_auc_score_returns_matches_sklearn(session): assert score == expected_score +def test_roc_auc_score_series(session): + pd_df = pd.DataFrame( + { + "y_true": [0, 0, 1, 1, 0, 1, 0, 1, 1, 1], + "y_score": [0.1, 0.4, 0.35, 0.8, 0.65, 0.9, 0.5, 0.3, 0.6, 0.45], + } + ) + + df = session.read_pandas(pd_df) + score = bigframes.ml.metrics.roc_auc_score(df["y_true"], df["y_score"]) + + assert score == 0.625 + + def test_auc_invalid_x_size(session): - pd_df = pd.DataFrame({"x": [0], "y": [0]}) + pd_df = pd.DataFrame({"x_arbitrary_name": [0], "y_arbitrary_name": [0]}) df = session.read_pandas(pd_df) with pytest.raises(ValueError): - bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) + bigframes.ml.metrics.auc(df[["x_arbitrary_name"]], df[["y_arbitrary_name"]]) def test_auc_nondecreasing_x(session): @@ -352,16 +484,23 @@ def test_auc_nonincreasing_x_negative(session): assert bigframes.ml.metrics.auc(df[["x"]], df[["y"]]) == -0.75 +def test_auc_series(session): + pd_df = pd.DataFrame({"x": [0, 0, 0.5, 0.5, 1], "y": [0, 0.5, 0.5, 1, 1]}) + + df = session.read_pandas(pd_df) + assert bigframes.ml.metrics.auc(df["x"], df["y"]) == 0.75 + + def test_confusion_matrix(session): pd_df = pd.DataFrame( { - "y_true": [2, 0, 2, 2, 0, 1], - "y_pred": [0, 0, 2, 2, 0, 2], + "y_true_arbitrary_name": [2, 0, 2, 2, 0, 1], + "y_pred_arbitrary_name": [0, 0, 2, 2, 0, 2], } ).astype("Int64") df = session.read_pandas(pd_df) confusion_matrix = bigframes.ml.metrics.confusion_matrix( - df[["y_true"]], df[["y_pred"]] + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]] ) expected_pd_df = pd.DataFrame( { @@ -441,7 +580,7 @@ def test_confusion_matrix_str_matches_sklearn(session): ) -def test_recall_score(session): +def test_confusion_matrix_series(session): pd_df = pd.DataFrame( { "y_true": [2, 0, 2, 2, 0, 1], @@ -449,8 +588,29 @@ def test_recall_score(session): } ).astype("Int64") df = session.read_pandas(pd_df) + confusion_matrix = bigframes.ml.metrics.confusion_matrix(df["y_true"], df["y_pred"]) + expected_pd_df = pd.DataFrame( + { + 0: [2, 0, 1], + 1: [0, 0, 0], + 2: [0, 1, 2], + } + ).astype("int64") + pd.testing.assert_frame_equal( + confusion_matrix, expected_pd_df, check_index_type=False + ) + + +def test_recall_score(session): + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [2, 0, 2, 2, 0, 1], + "y_pred_arbitrary_name": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) recall = bigframes.ml.metrics.recall_score( - df[["y_true"]], df[["y_pred"]], average=None + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]], average=None ) expected_values = [1.000000, 0.000000, 0.666667] expected_index = [0, 1, 2] @@ -497,7 +657,7 @@ def test_recall_score_str_matches_sklearn(session): pd.testing.assert_series_equal(recall, expected_recall, check_index_type=False) -def test_precision_score(session): +def test_recall_score_series(session): pd_df = pd.DataFrame( { "y_true": [2, 0, 2, 2, 0, 1], @@ -505,8 +665,24 @@ def test_precision_score(session): } ).astype("Int64") df = session.read_pandas(pd_df) + recall = bigframes.ml.metrics.recall_score(df["y_true"], df["y_pred"], average=None) + expected_values = [1.000000, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_recall = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal(recall, expected_recall, check_index_type=False) + + +def test_precision_score(session): + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [2, 0, 2, 2, 0, 1], + "y_pred_arbitrary_name": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) precision_score = bigframes.ml.metrics.precision_score( - df[["y_true"]], df[["y_pred"]], average=None + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]], average=None ) expected_values = [0.666667, 0.000000, 0.666667] expected_index = [0, 1, 2] @@ -559,7 +735,7 @@ def test_precision_score_str_matches_sklearn(session): ) -def test_f1_score(session): +def test_precision_score_series(session): pd_df = pd.DataFrame( { "y_true": [2, 0, 2, 2, 0, 1], @@ -567,8 +743,28 @@ def test_f1_score(session): } ).astype("Int64") df = session.read_pandas(pd_df) + precision_score = bigframes.ml.metrics.precision_score( + df["y_true"], df["y_pred"], average=None + ) + expected_values = [0.666667, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_precision = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal( + precision_score, expected_precision, check_index_type=False + ) + + +def test_f1_score(session): + pd_df = pd.DataFrame( + { + "y_true_arbitrary_name": [2, 0, 2, 2, 0, 1], + "y_pred_arbitrary_name": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) f1_score = bigframes.ml.metrics.f1_score( - df[["y_true"]], df[["y_pred"]], average=None + df[["y_true_arbitrary_name"]], df[["y_pred_arbitrary_name"]], average=None ) expected_values = [0.8, 0.000000, 0.666667] expected_index = [0, 1, 2] @@ -613,3 +809,19 @@ def test_f1_score_str_matches_sklearn(session): expected_index = ["ant", "bird", "cat"] expected_f1 = pd.Series(expected_values, index=expected_index) pd.testing.assert_series_equal(f1_score, expected_f1, check_index_type=False) + + +def test_f1_score_series(session): + pd_df = pd.DataFrame( + { + "y_true": [2, 0, 2, 2, 0, 1], + "y_pred": [0, 0, 2, 2, 0, 2], + } + ).astype("Int64") + df = session.read_pandas(pd_df) + f1_score = bigframes.ml.metrics.f1_score(df["y_true"], df["y_pred"], average=None) + expected_values = [0.8, 0.000000, 0.666667] + expected_index = [0, 1, 2] + expected_f1 = pd.Series(expected_values, index=expected_index) + + pd.testing.assert_series_equal(f1_score, expected_f1, check_index_type=False) diff --git a/tests/system/small/ml/test_model_selection.py b/tests/system/small/ml/test_model_selection.py index add455ab9ac..9eb36455913 100644 --- a/tests/system/small/ml/test_model_selection.py +++ b/tests/system/small/ml/test_model_selection.py @@ -15,7 +15,8 @@ import pandas as pd import pytest -import bigframes.ml.model_selection +from bigframes.ml import model_selection +import bigframes.pandas as bpd def test_train_test_split_default_correct_shape(penguins_df_default_index): @@ -27,9 +28,7 @@ def test_train_test_split_default_correct_shape(penguins_df_default_index): ] ] y = penguins_df_default_index[["body_mass_g"]] - X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( - X, y - ) + X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y) # even though the default seed is random, it should always result in this shape assert X_train.shape == (258, 3) @@ -38,6 +37,22 @@ def test_train_test_split_default_correct_shape(penguins_df_default_index): assert y_test.shape == (86, 1) +def test_train_test_split_series_default_correct_shape(penguins_df_default_index): + X = penguins_df_default_index[["species"]] + y = penguins_df_default_index["body_mass_g"] + X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y) + assert isinstance(X_train, bpd.DataFrame) + assert isinstance(X_test, bpd.DataFrame) + assert isinstance(y_train, bpd.Series) + assert isinstance(y_test, bpd.Series) + + # even though the default seed is random, it should always result in this shape + assert X_train.shape == (258, 1) + assert X_test.shape == (86, 1) + assert y_train.shape == (258,) + assert y_test.shape == (86,) + + def test_train_test_double_split_correct_shape(penguins_df_default_index): X = penguins_df_default_index[ [ @@ -47,7 +62,7 @@ def test_train_test_double_split_correct_shape(penguins_df_default_index): ] ] y = penguins_df_default_index[["body_mass_g"]] - X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( + X_train, X_test, y_train, y_test = model_selection.train_test_split( X, y, test_size=0.2, train_size=0.4 ) @@ -78,7 +93,7 @@ def test_train_test_three_dataframes_correct_shape(penguins_df_default_index): B_test, C_train, C_test, - ) = bigframes.ml.model_selection.train_test_split(A, B, C) + ) = model_selection.train_test_split(A, B, C) assert A_train.shape == (258, 2) assert A_test.shape == (86, 2) @@ -111,7 +126,7 @@ def test_train_test_split_seeded_correct_rows( ] ] y = df[["body_mass_g"]] - X_train, X_test, y_train, y_test = bigframes.ml.model_selection.train_test_split( + X_train, X_test, y_train, y_test = model_selection.train_test_split( X, y, random_state=42 ) @@ -124,31 +139,31 @@ def test_train_test_split_seeded_correct_rows( [ 144, 146, - 148, 168, 183, 186, 217, + 221, 225, - 226, 237, + 240, 244, 245, + 257, 260, 262, 263, + 264, 266, + 267, 268, - 269, - 289, 290, - 291, ], dtype="Int64", name="rowindex", ) test_index = pd.Index( - [161, 221, 240, 257, 264, 267, 278], dtype="Int64", name="rowindex" + [148, 161, 226, 269, 278, 289, 291], dtype="Int64", name="rowindex" ) all_data.index.name = "_" @@ -209,6 +224,6 @@ def test_train_test_split_value_error(penguins_df_default_index, train_size, tes ] y = penguins_df_default_index[["body_mass_g"]] with pytest.raises(ValueError): - bigframes.ml.model_selection.train_test_split( + model_selection.train_test_split( X, y, train_size=train_size, test_size=test_size ) diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index f31b93b4cc2..420a80754fb 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -20,7 +20,7 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): - # TODO(bmil): add a second test that compares output to sklearn.preprocessing.StandardScaler + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. scaler = bigframes.ml.preprocessing.StandardScaler() scaler.fit( penguins_df_default_index[ @@ -34,10 +34,9 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): ] ).to_pandas() - # If standard-scaled correctly, mean should be 0.0 and standard deviation 1.0 + # If standard-scaled correctly, mean should be 0.0 for column in result.columns: assert math.isclose(result[column].mean(), 0.0, abs_tol=1e-3) - assert math.isclose(result[column].std(), 1.0, abs_tol=1e-3) result = scaler.transform(new_penguins_df).to_pandas() @@ -48,9 +47,9 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): expected = pd.DataFrame( { - "scaled_culmen_depth_mm": [0.8349, 0.02473, 0.4805], - "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103], - "scaled_flipper_length_mm": [-0.3495, -1.4163, -0.9185], + "scaled_culmen_depth_mm": [0.836148, 0.024748, 0.48116], + "scaled_culmen_length_mm": [-0.81112, -0.994552, -1.104611], + "scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], }, dtype="Float64", index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), @@ -59,9 +58,42 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): pd.testing.assert_frame_equal(result, expected, rtol=1e-3) -def test_one_hot_encoder_encodes(penguins_df_default_index, new_penguins_df): +def test_standard_scaler_series_normalizes(penguins_df_default_index, new_penguins_df): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + scaler = bigframes.ml.preprocessing.StandardScaler() + scaler.fit(penguins_df_default_index["culmen_length_mm"]) + + result = scaler.transform(penguins_df_default_index["culmen_length_mm"]).to_pandas() + + # If standard-scaled correctly, mean should be 0.0 + for column in result.columns: + assert math.isclose(result[column].mean(), 0.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "scaled_culmen_length_mm": [ + -0.811119671289163, + -0.9945520581113803, + -1.104611490204711, + ], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_one_hot_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() - encoder.fit(penguins_df_default_index["species", "sex"]) + encoder.fit(new_penguins_df["species", "sex"]) result = encoder.transform(new_penguins_df).to_pandas() @@ -87,3 +119,91 @@ def test_one_hot_encoder_encodes(penguins_df_default_index, new_penguins_df): ) pd.testing.assert_frame_equal(result, expected) + + +def test_one_hot_encoder_series_default_params(new_penguins_df): + encoder = bigframes.ml.preprocessing.OneHotEncoder() + encoder.fit(new_penguins_df["species"]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + }, + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_one_hot_encoder_params(new_penguins_df): + encoder = bigframes.ml.preprocessing.OneHotEncoder("most_frequent", 100, 2) + encoder.fit(new_penguins_df["species", "sex"]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "onehotencoded_sex": [ + [{"index": 0, "value": 1.0}], + [{"index": 0, "value": 1.0}], + [{"index": 0, "value": 1.0}], + ], + "onehotencoded_species": [ + [{"index": 0, "value": 1.0}], + [{"index": 0, "value": 1.0}], + [{"index": 0, "value": 1.0}], + ], + }, + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_one_hot_encoder_different_data(penguins_df_default_index, new_penguins_df): + encoder = bigframes.ml.preprocessing.OneHotEncoder() + encoder.fit(penguins_df_default_index["species", "sex"]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "onehotencoded_sex": [ + [{"index": 3, "value": 1.0}], + [{"index": 2, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + }, + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +# TODO(garrettwu): add OneHotEncoder tests to compare with sklearn. diff --git a/tests/system/small/ml/test_register.py b/tests/system/small/ml/test_register.py index ad378d35084..bcf1f4a5b0e 100644 --- a/tests/system/small/ml/test_register.py +++ b/tests/system/small/ml/test_register.py @@ -28,7 +28,7 @@ def test_linear_reg_register( ) # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. assert ( - model_name + model_name[:63] # truncated in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ "vertexAiModelId" ] @@ -44,7 +44,7 @@ def test_linear_reg_register_with_params( # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. assert ( - model_name + model_name[:63] # truncated in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ "vertexAiModelId" ] @@ -62,7 +62,7 @@ def test_palm2_text_generator_register( ) # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. assert ( - model_name + model_name[:63] # truncated in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ "vertexAiModelId" ] @@ -80,7 +80,7 @@ def test_imported_tensorflow_register( ) # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. assert ( - model_name + model_name[:63] # truncated in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][ "vertexAiModelId" ] diff --git a/tests/system/small/operations/test_datetimes.py b/tests/system/small/operations/test_datetimes.py index 25d3eb69adb..7dc55b9367c 100644 --- a/tests/system/small/operations/test_datetimes.py +++ b/tests/system/small/operations/test_datetimes.py @@ -30,7 +30,7 @@ def test_day(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.day.compute() + bf_result = bf_series.dt.day.to_pandas() pd_result = scalars_pandas_df[col_name].dt.day assert_series_equal_ignoring_order( @@ -48,7 +48,7 @@ def test_date(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.date.compute() + bf_result = bf_series.dt.date.to_pandas() pd_result = scalars_pandas_df[col_name].dt.date assert_series_equal_ignoring_order( @@ -66,7 +66,7 @@ def test_dayofweek(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.dayofweek.compute() + bf_result = bf_series.dt.dayofweek.to_pandas() pd_result = scalars_pandas_df[col_name].dt.dayofweek assert_series_equal_ignoring_order(pd_result, bf_result, check_dtype=False) @@ -81,7 +81,7 @@ def test_hour(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.hour.compute() + bf_result = bf_series.dt.hour.to_pandas() pd_result = scalars_pandas_df[col_name].dt.hour assert_series_equal_ignoring_order( @@ -99,7 +99,7 @@ def test_minute(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.minute.compute() + bf_result = bf_series.dt.minute.to_pandas() pd_result = scalars_pandas_df[col_name].dt.minute assert_series_equal_ignoring_order( @@ -117,7 +117,7 @@ def test_month(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.month.compute() + bf_result = bf_series.dt.month.to_pandas() pd_result = scalars_pandas_df[col_name].dt.month assert_series_equal_ignoring_order( @@ -135,7 +135,7 @@ def test_quarter(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.quarter.compute() + bf_result = bf_series.dt.quarter.to_pandas() pd_result = scalars_pandas_df[col_name].dt.quarter assert_series_equal_ignoring_order( @@ -153,7 +153,7 @@ def test_second(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.second.compute() + bf_result = bf_series.dt.second.to_pandas() pd_result = scalars_pandas_df[col_name].dt.second assert_series_equal_ignoring_order( @@ -171,7 +171,7 @@ def test_time(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.time.compute() + bf_result = bf_series.dt.time.to_pandas() pd_result = scalars_pandas_df[col_name].dt.time assert_series_equal_ignoring_order( @@ -189,7 +189,7 @@ def test_year(scalars_dfs, col_name): pytest.skip("Pyarrow datetime objects not support in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.dt.year.compute() + bf_result = bf_series.dt.year.to_pandas() pd_result = scalars_pandas_df[col_name].dt.year assert_series_equal_ignoring_order( diff --git a/tests/system/small/operations/test_strings.py b/tests/system/small/operations/test_strings.py index 76736e3cddd..d7bf3312f09 100644 --- a/tests/system/small/operations/test_strings.py +++ b/tests/system/small/operations/test_strings.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re + import pandas as pd import pytest @@ -24,7 +26,7 @@ def test_find(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.find("W").compute() + bf_result = bf_series.str.find("W").to_pandas() pd_result = scalars_pandas_df[col_name].str.find("W") # One of type mismatches to be documented. Here, the `bf_result.dtype` is `Int64` but @@ -35,11 +37,138 @@ def test_find(scalars_dfs): ) +@pytest.mark.parametrize( + ("pat", "case", "flags", "regex"), + [ + ("hEllo", True, 0, False), + ("hEllo", False, 0, False), + ("hEllo", False, re.I, True), + (".*", True, 0, True), + (".*", True, 0, False), + ], +) +def test_str_contains(scalars_dfs, pat, case, flags, regex): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + + bf_result = bf_series.str.contains( + pat, case=case, flags=flags, regex=regex + ).to_pandas() + pd_result = scalars_pandas_df[col_name].str.contains( + pat, case=case, flags=flags, regex=regex + ) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("pat"), + [(r"(ell)(lo)"), (r"(?Ph..)"), (r"(?Pe.*o)([g-l]+)")], +) +def test_str_extract(scalars_dfs, pat): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + + bf_result = bf_series.str.extract(pat).to_pandas() + pd_result = scalars_pandas_df[col_name].str.extract(pat) + + # Pandas produces int col labels, while bq df only supports str labels at present + pd_result = pd_result.set_axis(pd_result.columns.astype(str), axis=1) + pd.testing.assert_frame_equal( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("pat", "repl", "case", "flags", "regex"), + [ + ("hEllo", "blah", True, 0, False), + ("hEllo", "blah", False, 0, False), + ("hEllo", "blah", False, re.I, True), + (".*", "blah", True, 0, True), + ("h.l", "blah", False, 0, True), + (re.compile("(?i).e.."), "blah", None, 0, True), + ], +) +def test_str_replace(scalars_dfs, pat, repl, case, flags, regex): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + + bf_result = bf_series.str.replace( + pat, repl=repl, case=case, flags=flags, regex=regex + ).to_pandas() + pd_result = scalars_pandas_df[col_name].str.replace( + pat, repl=repl, case=case, flags=flags, regex=regex + ) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("pat",), + [ + ("こん",), + ("Tag!",), + ( + ( + "Tag!", + "Hel", + ), + ), + ], +) +def test_str_startswith(scalars_dfs, pat): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + pd_series = scalars_pandas_df[col_name].astype("object") + + bf_result = bf_series.str.startswith(pat).to_pandas() + pd_result = pd_series.str.startswith(pat) + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("pat",), + [ + ("こん",), + ("Tag!",), + ( + ( + "Tag!", + "Hel", + ), + ), + ], +) +def test_str_endswith(scalars_dfs, pat): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + pd_series = scalars_pandas_df[col_name].astype("object") + + bf_result = bf_series.str.endswith(pat).to_pandas() + pd_result = pd_series.str.endswith(pat) + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + def test_len(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.len().compute() + bf_result = bf_series.str.len().to_pandas() pd_result = scalars_pandas_df[col_name].str.len() # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is `Int64` but @@ -54,7 +183,7 @@ def test_lower(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.lower().compute() + bf_result = bf_series.str.lower().to_pandas() pd_result = scalars_pandas_df[col_name].str.lower() assert_series_equal_ignoring_order( @@ -67,7 +196,7 @@ def test_reverse(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.reverse().compute() + bf_result = bf_series.str.reverse().to_pandas() pd_result = scalars_pandas_df[col_name].copy() for i in pd_result.index: cell = pd_result.loc[i] @@ -89,7 +218,7 @@ def test_slice(scalars_dfs, start, stop): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.slice(start, stop).compute() + bf_result = bf_series.str.slice(start, stop).to_pandas() pd_series = scalars_pandas_df[col_name] pd_result = pd_series.str.slice(start, stop) @@ -103,7 +232,7 @@ def test_strip(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.strip().compute() + bf_result = bf_series.str.strip().to_pandas() pd_result = scalars_pandas_df[col_name].str.strip() assert_series_equal_ignoring_order( @@ -116,7 +245,7 @@ def test_upper(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.upper().compute() + bf_result = bf_series.str.upper().to_pandas() pd_result = scalars_pandas_df[col_name].str.upper() assert_series_equal_ignoring_order( @@ -149,7 +278,7 @@ def test_isnumeric(session): df = session.read_pandas(pandas_df) pd_result = pandas_df.numeric_string_col.str.isnumeric() - bf_result = df.numeric_string_col.str.isnumeric().compute() + bf_result = df.numeric_string_col.str.isnumeric().to_pandas() assert_series_equal_ignoring_order( bf_result, @@ -163,7 +292,7 @@ def test_rstrip(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.rstrip().compute() + bf_result = bf_series.str.rstrip().to_pandas() pd_result = scalars_pandas_df[col_name].str.rstrip() assert_series_equal_ignoring_order( @@ -176,7 +305,7 @@ def test_lstrip(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.lstrip().compute() + bf_result = bf_series.str.lstrip().to_pandas() pd_result = scalars_pandas_df[col_name].str.lstrip() assert_series_equal_ignoring_order( @@ -190,7 +319,7 @@ def test_repeat(scalars_dfs, repeats): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.repeat(repeats).compute() + bf_result = bf_series.str.repeat(repeats).to_pandas() pd_result = scalars_pandas_df[col_name].str.repeat(repeats) assert_series_equal_ignoring_order( @@ -203,7 +332,7 @@ def test_capitalize(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" bf_series: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_series.str.capitalize().compute() + bf_result = bf_series.str.capitalize().to_pandas() pd_result = scalars_pandas_df[col_name].str.capitalize() assert_series_equal_ignoring_order( @@ -218,7 +347,7 @@ def test_cat_with_series(scalars_dfs): bf_filter: bigframes.series.Series = scalars_df["bool_col"] bf_left: bigframes.series.Series = scalars_df[col_name][bf_filter] bf_right: bigframes.series.Series = scalars_df[col_name] - bf_result = bf_left.str.cat(others=bf_right).compute() + bf_result = bf_left.str.cat(others=bf_right).to_pandas() pd_filter = scalars_pandas_df["bool_col"] pd_left = scalars_pandas_df[col_name][pd_filter] pd_right = scalars_pandas_df[col_name] diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 1833c674fdd..c7e17f5a2dd 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -13,6 +13,8 @@ # limitations under the License. import operator +import typing +from typing import Tuple import geopandas as gpd # type: ignore import numpy as np @@ -24,6 +26,7 @@ import bigframes import bigframes._config.display_options as display_options import bigframes.dataframe as dataframe +import bigframes.series as series from tests.system.utils import ( assert_pandas_df_equal_ignore_ordering, assert_series_equal_ignoring_order, @@ -33,7 +36,7 @@ def test_df_construct_copy(scalars_dfs): columns = ["int64_col", "string_col", "float64_col"] scalars_df, scalars_pandas_df = scalars_dfs - bf_result = dataframe.DataFrame(scalars_df, columns=columns).compute() + bf_result = dataframe.DataFrame(scalars_df, columns=columns).to_pandas() pd_result = pd.DataFrame(scalars_pandas_df, columns=columns) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -41,7 +44,7 @@ def test_df_construct_copy(scalars_dfs): def test_df_construct_pandas(scalars_dfs): columns = ["int64_too", "int64_col", "float64_col", "bool_col", "string_col"] _, scalars_pandas_df = scalars_dfs - bf_result = dataframe.DataFrame(scalars_pandas_df, columns=columns).compute() + bf_result = dataframe.DataFrame(scalars_pandas_df, columns=columns).to_pandas() pd_result = pd.DataFrame(scalars_pandas_df, columns=columns) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -56,7 +59,7 @@ def test_df_construct_pandas_set_dtype(scalars_dfs): _, scalars_pandas_df = scalars_dfs bf_result = dataframe.DataFrame( scalars_pandas_df, columns=columns, dtype="Float64" - ).compute() + ).to_pandas() pd_result = pd.DataFrame(scalars_pandas_df, columns=columns, dtype="Float64") pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -66,7 +69,7 @@ def test_df_construct_from_series(scalars_dfs): bf_result = dataframe.DataFrame( {"a": scalars_df["int64_col"], "b": scalars_df["string_col"]}, dtype="string[pyarrow]", - ).compute() + ).to_pandas() pd_result = pd.DataFrame( {"a": scalars_pandas_df["int64_col"], "b": scalars_pandas_df["string_col"]}, dtype="string[pyarrow]", @@ -78,7 +81,7 @@ def test_get_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_col" series = scalars_df[col_name] - bf_result = series.compute() + bf_result = series.to_pandas() pd_result = scalars_pandas_df[col_name] assert_series_equal_ignoring_order(bf_result, pd_result) @@ -96,7 +99,7 @@ def test_head_with_custom_column_labels(scalars_df_index, scalars_pandas_df_inde "string_col": "言語列", } bf_df = scalars_df_index.rename(columns=rename_mapping).head(3) - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = scalars_pandas_df_index.rename(columns=rename_mapping).head(3) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -107,7 +110,7 @@ def test_tail_with_custom_column_labels(scalars_df_index, scalars_pandas_df_inde "string_col": "言語列", } bf_df = scalars_df_index.rename(columns=rename_mapping).tail(3) - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = scalars_pandas_df_index.rename(columns=rename_mapping).tail(3) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -115,7 +118,7 @@ def test_tail_with_custom_column_labels(scalars_df_index, scalars_pandas_df_inde def test_get_column_by_attr(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs series = scalars_df.int64_col - bf_result = series.compute() + bf_result = series.to_pandas() pd_result = scalars_pandas_df.int64_col assert_series_equal_ignoring_order(bf_result, pd_result) @@ -124,7 +127,7 @@ def test_get_columns(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_names = ["bool_col", "float64_col", "int64_col"] df_subset = scalars_df.get(col_names) - df_pandas = df_subset.compute() + df_pandas = df_subset.to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df[col_names].columns ) @@ -140,7 +143,7 @@ def test_get_columns_default(scalars_dfs): def test_drop_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_col" - df_pandas = scalars_df.drop(columns=col_name).compute() + df_pandas = scalars_df.drop(columns=col_name).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.drop(columns=col_name).columns ) @@ -149,12 +152,22 @@ def test_drop_column(scalars_dfs): def test_drop_columns(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_names = ["int64_col", "geography_col", "time_col"] - df_pandas = scalars_df.drop(columns=col_names).compute() + df_pandas = scalars_df.drop(columns=col_names).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.drop(columns=col_names).columns ) +def test_drop_labels_axis_1(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + labels = ["int64_col", "geography_col", "time_col"] + + pd_result = scalars_pandas_df.drop(labels=labels, axis=1) + bf_result = scalars_df.drop(labels=labels, axis=1).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + def test_drop_with_custom_column_labels(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs rename_mapping = { @@ -166,17 +179,44 @@ def test_drop_with_custom_column_labels(scalars_dfs): "timestamp_col", ] bf_df = scalars_df.rename(columns=rename_mapping).drop(columns=dropped_columns) - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = scalars_pandas_df.rename(columns=rename_mapping).drop( columns=dropped_columns ) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +def test_drop_index(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = scalars_pandas_df.drop(index=[4, 1, 2]) + bf_result = scalars_df.drop(index=[4, 1, 2]).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + +def test_drop_labels_axis_0(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = scalars_pandas_df.drop(labels=[4, 1, 2], axis=0) + bf_result = scalars_df.drop(labels=[4, 1, 2], axis=0).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + +def test_drop_index_and_columns(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = scalars_pandas_df.drop(index=[4, 1, 2], columns="int64_col") + bf_result = scalars_df.drop(index=[4, 1, 2], columns="int64_col").to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + def test_rename(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name_dict = {"bool_col": "boolean_col"} - df_pandas = scalars_df.rename(columns=col_name_dict).compute() + df_pandas = scalars_df.rename(columns=col_name_dict).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns ) @@ -210,7 +250,7 @@ def test_repr_w_all_rows(scalars_dfs): def test_repr_html_w_all_rows(scalars_dfs): scalars_df, _ = scalars_dfs # get a pandas df of the expected format - df, _ = scalars_df._block.compute() + df, _ = scalars_df._block.to_pandas() pandas_df = df.set_axis(scalars_df._block.column_labels, axis=1) pandas_df.index.name = scalars_df.index.name @@ -229,7 +269,7 @@ def test_repr_html_w_all_rows(scalars_dfs): def test_df_column_name_with_space(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name_dict = {"bool_col": "bool col"} - df_pandas = scalars_df.rename(columns=col_name_dict).compute() + df_pandas = scalars_df.rename(columns=col_name_dict).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns ) @@ -238,7 +278,7 @@ def test_df_column_name_with_space(scalars_dfs): def test_df_column_name_duplicate(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name_dict = {"int64_too": "int64_col"} - df_pandas = scalars_df.rename(columns=col_name_dict).compute() + df_pandas = scalars_df.rename(columns=col_name_dict).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns ) @@ -248,7 +288,7 @@ def test_get_df_column_name_duplicate(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name_dict = {"int64_too": "int64_col"} - bf_result = scalars_df.rename(columns=col_name_dict)["int64_col"].compute() + bf_result = scalars_df.rename(columns=col_name_dict)["int64_col"].to_pandas() pd_result = scalars_pandas_df.rename(columns=col_name_dict)["int64_col"] pd.testing.assert_index_equal(bf_result.columns, pd_result.columns) @@ -257,7 +297,7 @@ def test_filter_df(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_bool_series = scalars_df["bool_col"] - bf_result = scalars_df[bf_bool_series].compute() + bf_result = scalars_df[bf_bool_series].to_pandas() pd_bool_series = scalars_pandas_df["bool_col"] pd_result = scalars_pandas_df[pd_bool_series] @@ -269,7 +309,7 @@ def test_assign_new_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs kwargs = {"new_col": 2} df = scalars_df.assign(**kwargs) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign(**kwargs) # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -284,7 +324,7 @@ def test_assign_new_column_w_loc(scalars_dfs): pd_df = scalars_pandas_df.copy() bf_df.loc[:, "new_col"] = 2 pd_df.loc[:, "new_col"] = 2 - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = pd_df # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -299,7 +339,7 @@ def test_assign_new_column_w_setitem(scalars_dfs): pd_df = scalars_pandas_df.copy() bf_df["new_col"] = 2 pd_df["new_col"] = 2 - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = pd_df # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -312,7 +352,7 @@ def test_assign_existing_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs kwargs = {"int64_col": 2} df = scalars_df.assign(**kwargs) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign(**kwargs) # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -325,7 +365,7 @@ def test_assign_series(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs column_name = "int64_col" df = scalars_df.assign(new_col=scalars_df[column_name]) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign(new_col=scalars_pandas_df[column_name]) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -335,7 +375,7 @@ def test_assign_series_overwrite(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs column_name = "int64_col" df = scalars_df.assign(**{column_name: scalars_df[column_name] + 3}) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign( **{column_name: scalars_pandas_df[column_name] + 3} ) @@ -347,7 +387,7 @@ def test_assign_sequential(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs kwargs = {"int64_col": 2, "new_col": 3, "new_col2": 4} df = scalars_df.assign(**kwargs) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign(**kwargs) # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -371,7 +411,7 @@ def test_assign_same_table_different_index_performs_self_join( ) bf_df_2 = bf_df.set_index("alternative_index") pd_df_2 = pd_df.set_index("alternative_index") - bf_result = bf_df.assign(new_col=bf_df_2[column_name] * 10).compute() + bf_result = bf_df.assign(new_col=bf_df_2[column_name] * 10).to_pandas() pd_result = pd_df.assign(new_col=pd_df_2[column_name] * 10) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -383,7 +423,7 @@ def test_assign_different_df( ): column_name = "int64_col" df = scalars_df_index.assign(new_col=scalars_df_2_index[column_name]) - bf_result = df.compute() + bf_result = df.to_pandas() # Doesn't matter to pandas if it comes from the same DF or a different DF. pd_result = scalars_pandas_df_index.assign( new_col=scalars_pandas_df_index[column_name] @@ -402,7 +442,7 @@ def test_assign_different_df_w_loc( assert "int64_col" in pd_df.columns bf_df.loc[:, "int64_col"] = bf_df2.loc[:, "int64_col"] + 1 pd_df.loc[:, "int64_col"] = pd_df.loc[:, "int64_col"] + 1 - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = pd_df # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -421,7 +461,7 @@ def test_assign_different_df_w_setitem( assert "int64_col" in pd_df.columns bf_df["int64_col"] = bf_df2["int64_col"] + 1 pd_df["int64_col"] = pd_df["int64_col"] + 1 - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_result = pd_df # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -434,7 +474,7 @@ def test_assign_callable_lambda(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs kwargs = {"new_col": lambda x: x["int64_col"] + x["int64_too"]} df = scalars_df.assign(**kwargs) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.assign(**kwargs) # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. @@ -446,7 +486,7 @@ def test_assign_callable_lambda(scalars_dfs): def test_dropna(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs df = scalars_df.dropna() - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.dropna() assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -472,7 +512,7 @@ def test_merge(scalars_dfs, merge_how): right = scalars_df[right_columns].assign(rowindex_2=scalars_df["rowindex_2"] + 2) df = left.merge(right, merge_how, on, sort=True) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df[left_columns].merge( scalars_pandas_df[right_columns].assign( @@ -506,7 +546,7 @@ def test_merge_custom_col_name(scalars_dfs, merge_how): left = left.rename(columns=rename_columns) right = scalars_df[right_columns] df = left.merge(right, merge_how, on, sort=True) - bf_result = df.compute() + bf_result = df.to_pandas() pandas_left_df = scalars_pandas_df[left_columns] pandas_left_df = pandas_left_df.rename(columns=rename_columns) @@ -536,7 +576,7 @@ def test_merge_left_on_right_on(scalars_dfs, merge_how): df = left.merge( right, merge_how, left_on="int64_too", right_on="rowindex_2", sort=True ) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df[left_columns].merge( scalars_pandas_df[right_columns], @@ -633,7 +673,7 @@ def test_empty_false(scalars_dfs): assert bf_result == pd_result -def test_empty_true(scalars_dfs): +def test_empty_true_column_filter(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = scalars_df[[]].empty @@ -642,6 +682,31 @@ def test_empty_true(scalars_dfs): assert bf_result == pd_result +def test_empty_true_row_filter(scalars_dfs: Tuple[dataframe.DataFrame, pd.DataFrame]): + scalars_df, scalars_pandas_df = scalars_dfs + bf_bool: series.Series = typing.cast(series.Series, scalars_df["bool_col"]) + pd_bool: pd.Series = scalars_pandas_df["bool_col"] + bf_false = bf_bool.notna() & (bf_bool != bf_bool) + pd_false = pd_bool.notna() & (pd_bool != pd_bool) + + bf_result = scalars_df[bf_false].empty + pd_result = scalars_pandas_df[pd_false].empty + + assert pd_result + assert bf_result == pd_result + + +def test_empty_true_memtable(session: bigframes.Session): + bf_df = dataframe.DataFrame(session=session) + pd_df = pd.DataFrame() + + bf_result = bf_df.empty + pd_result = pd_df.empty + + assert pd_result + assert bf_result == pd_result + + @pytest.mark.parametrize( ("drop",), ((True,), (False,)), @@ -650,7 +715,7 @@ def test_reset_index(scalars_df_index, scalars_pandas_df_index, drop): df = scalars_df_index.reset_index(drop=drop) assert df.index.name is None - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df_index.reset_index(drop=drop) # Pandas uses int64 instead of Int64 (nullable) dtype. @@ -666,7 +731,7 @@ def test_reset_index_then_filter( ): bf_filter = scalars_df_index["bool_col"].fillna(True) bf_df = scalars_df_index.reset_index()[bf_filter] - bf_result = bf_df.compute() + bf_result = bf_df.to_pandas() pd_filter = scalars_pandas_df_index["bool_col"].fillna(True) pd_result = scalars_pandas_df_index.reset_index()[pd_filter] @@ -693,7 +758,7 @@ def test_reset_index_with_unnamed_index( # reset_index(drop=False) creates a new column "index". assert df.columns[0] == "index" - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df_index.reset_index(drop=False) # Pandas uses int64 instead of Int64 (nullable) dtype. @@ -720,7 +785,7 @@ def test_reset_index_with_unnamed_index_and_index_column( # reset_index(drop=False) creates a new column "level_0" if the "index" column already exists. assert df.columns[0] == "level_0" - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df_index.assign( index=scalars_pandas_df_index["int64_col"] ).reset_index(drop=False) @@ -753,7 +818,7 @@ def test_reset_index_with_unnamed_index_and_index_column( def test_set_index(scalars_dfs, index_column, drop, append): scalars_df, scalars_pandas_df = scalars_dfs df = scalars_df.set_index(index_column, append=append, drop=drop) - bf_result = df.compute() + bf_result = df.to_pandas() pd_result = scalars_pandas_df.set_index(index_column, append=append, drop=drop) # Sort to disambiguate when there are duplicate index labels. @@ -767,11 +832,30 @@ def test_set_index(scalars_dfs, index_column, drop, append): pandas.testing.assert_frame_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("ascending",), + ((True,), (False,)), +) +@pytest.mark.parametrize( + ("na_position",), + (("first",), ("last",)), +) +def test_sort_index(scalars_dfs, ascending, na_position): + index_column = "int64_col" + scalars_df, scalars_pandas_df = scalars_dfs + df = scalars_df.set_index(index_column) + bf_result = df.sort_index(ascending=ascending, na_position=na_position).to_pandas() + pd_result = scalars_pandas_df.set_index(index_column).sort_index( + ascending=ascending, na_position=na_position + ) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + def test_df_abs(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs columns = ["int64_col", "int64_too", "float64_col"] - bf_result = scalars_df[columns].abs().compute() + bf_result = scalars_df[columns].abs().to_pandas() pd_result = scalars_pandas_df[columns].abs() assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -781,7 +865,7 @@ def test_df_isnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs columns = ["int64_col", "int64_too", "string_col", "bool_col"] - bf_result = scalars_df[columns].isnull().compute() + bf_result = scalars_df[columns].isnull().to_pandas() pd_result = scalars_pandas_df[columns].isnull() # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is @@ -798,7 +882,7 @@ def test_df_notnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs columns = ["int64_col", "int64_too", "string_col", "bool_col"] - bf_result = scalars_df[columns].notnull().compute() + bf_result = scalars_df[columns].notnull().to_pandas() pd_result = scalars_pandas_df[columns].notnull() # One of dtype mismatches to be documented. Here, the `bf_result.dtype` is @@ -819,6 +903,8 @@ def test_df_notnull(scalars_dfs): operator.mul, operator.truediv, operator.floordiv, + operator.eq, + operator.ne, operator.gt, operator.ge, operator.lt, @@ -830,6 +916,8 @@ def test_df_notnull(scalars_dfs): "multiply", "true_divide", "floor_divide", + "eq", + "ne", "gt", "ge", "lt", @@ -845,7 +933,7 @@ def test_scalar_binop(scalars_dfs, op, other_scalar, reverse_operands): maybe_reversed_op = (lambda x, y: op(y, x)) if reverse_operands else op - bf_result = maybe_reversed_op(scalars_df[columns], other_scalar).compute() + bf_result = maybe_reversed_op(scalars_df[columns], other_scalar).to_pandas() pd_result = maybe_reversed_op(scalars_pandas_df[columns], other_scalar) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -857,7 +945,7 @@ def test_mod(scalars_dfs, other_scalar): # This is likely a pandas bug as mod 0 is undefined in other dtypes, and most programming languages. scalars_df, scalars_pandas_df = scalars_dfs - bf_result = (scalars_df[["int64_col", "int64_too"]] % other_scalar).compute() + bf_result = (scalars_df[["int64_col", "int64_too"]] % other_scalar).to_pandas() pd_result = scalars_pandas_df[["int64_col", "int64_too"]] % other_scalar assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -867,7 +955,7 @@ def test_scalar_binop_str_exception(scalars_dfs): scalars_df, _ = scalars_dfs columns = ["string_col"] with pytest.raises(TypeError): - (scalars_df[columns] + 1).compute() + (scalars_df[columns] + 1).to_pandas() @pytest.mark.parametrize( @@ -913,7 +1001,7 @@ def test_series_binop_axis_index( df_columns = ["int64_col", "float64_col"] series_column = "int64_too" - bf_result = op(scalars_df[df_columns], scalars_df[series_column]).compute() + bf_result = op(scalars_df[df_columns], scalars_df[series_column]).to_pandas() pd_result = op(scalars_pandas_df[df_columns], scalars_pandas_df[series_column]) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -963,7 +1051,7 @@ def test_dataframe_binop_axis_index_throws_not_implemented( other_df_columns = ["int64_too"] with pytest.raises(NotImplementedError): - op(scalars_df[df_columns], scalars_df[other_df_columns]).compute() + op(scalars_df[df_columns], scalars_df[other_df_columns]).to_pandas() # Differnt table will only work for explicit index, since default index orders are arbitrary. @@ -976,7 +1064,7 @@ def test_series_binop_add_different_table( bf_result = ( scalars_df_index[df_columns] .add(scalars_df_2_index[series_column], axis="index") - .compute() + .to_pandas() ) pd_result = scalars_pandas_df_index[df_columns].add( scalars_pandas_df_index[series_column], axis="index" @@ -1001,14 +1089,12 @@ def test_series_binop_add_different_table( @all_joins def test_join_same_table(scalars_dfs, how): bf_df, pd_df = scalars_dfs - if how == "right" and pd_df.index.name != "rowindex": - pytest.skip("right join not supported without an index") - bf_df_a = bf_df[["string_col", "int64_col"]] - bf_df_b = bf_df[["float64_col"]] - bf_result = bf_df_a.join(bf_df_b, how=how).compute() - pd_df_a = pd_df[["string_col", "int64_col"]] - pd_df_b = pd_df[["float64_col"]] + bf_df_a = bf_df.set_index("int64_too")[["string_col", "int64_col"]] + bf_df_b = bf_df.set_index("int64_too")[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, how=how).to_pandas() + pd_df_a = pd_df.set_index("int64_too")[["string_col", "int64_col"]] + pd_df_b = pd_df.set_index("int64_too")[["float64_col"]] pd_result = pd_df_a.join(pd_df_b, how=how) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -1019,7 +1105,7 @@ def test_join_different_table( ): bf_df_a = scalars_df_index[["string_col", "int64_col"]] bf_df_b = scalars_df_2_index.dropna()[["float64_col"]] - bf_result = bf_df_a.join(bf_df_b, how=how).compute() + bf_result = bf_df_a.join(bf_df_b, how=how).to_pandas() pd_df_a = scalars_pandas_df_index[["string_col", "int64_col"]] pd_df_b = scalars_pandas_df_index.dropna()[["float64_col"]] pd_result = pd_df_a.join(pd_df_b, how=how) @@ -1031,7 +1117,23 @@ def test_join_duplicate_columns_raises_not_implemented(scalars_dfs): df_a = scalars_df[["string_col", "float64_col"]] df_b = scalars_df[["float64_col"]] with pytest.raises(NotImplementedError): - df_a.join(df_b, how="outer").compute() + df_a.join(df_b, how="outer").to_pandas() + + +@all_joins +def test_join_param_on(scalars_dfs, how): + bf_df, pd_df = scalars_dfs + + bf_df_a = bf_df[["string_col", "int64_col", "rowindex_2"]] + bf_df_a = bf_df_a.assign(rowindex_2=bf_df_a["rowindex_2"] + 2) + bf_df_b = bf_df[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, on="rowindex_2", how=how).to_pandas() + + pd_df_a = pd_df[["string_col", "int64_col", "rowindex_2"]] + pd_df_a = pd_df_a.assign(rowindex_2=pd_df_a["rowindex_2"] + 2) + pd_df_b = pd_df[["float64_col"]] + pd_result = pd_df_a.join(pd_df_b, on="rowindex_2", how=how) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @pytest.mark.parametrize( @@ -1050,7 +1152,7 @@ def test_dataframe_sort_values( # Test needs values to be unique bf_result = scalars_df_index.sort_values( by, ascending=ascending, na_position=na_position - ).compute() + ).to_pandas() pd_result = scalars_pandas_df_index.sort_values( by, ascending=ascending, na_position=na_position ) @@ -1061,6 +1163,22 @@ def test_dataframe_sort_values( ) +def test_dataframe_sort_values_stable(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.sort_values("int64_col", kind="stable") + .sort_values("bool_col", kind="stable") + .to_pandas() + ) + pd_result = scalars_pandas_df_index.sort_values( + "int64_col", kind="stable" + ).sort_values("bool_col", kind="stable") + + pandas.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + @pytest.mark.parametrize( ("operator", "columns"), [ @@ -1086,7 +1204,7 @@ def test_dataframe_numeric_analytic_op( # TODO: Add nullable ints (pandas 1.x has poor behavior on these) bf_series = operator(scalars_df_index[columns]) pd_series = operator(scalars_pandas_df_index[columns]) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() pd.testing.assert_frame_equal(pd_series, bf_result, check_dtype=False) @@ -1111,13 +1229,90 @@ def test_dataframe_general_analytic_op( col_names = ["int64_too", "float64_col", "int64_col", "bool_col"] bf_series = operator(scalars_df_index[col_names]) pd_series = operator(scalars_pandas_df_index[col_names]) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() pd.testing.assert_frame_equal( pd_series, bf_result, ) +def test_dataframe_agg_single_string(scalars_dfs): + numeric_cols = ["int64_col", "int64_too", "float64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[numeric_cols].agg("sum").to_pandas() + pd_result = scalars_pandas_df[numeric_cols].agg("sum") + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + +def test_dataframe_agg_multi_string(scalars_dfs): + numeric_cols = ["int64_col", "int64_too", "float64_col"] + aggregations = [ + "sum", + "mean", + "median", + "std", + "var", + "min", + "max", + "nunique", + "count", + ] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[numeric_cols].agg(aggregations).to_pandas() + pd_result = scalars_pandas_df[numeric_cols].agg(aggregations) + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + # Drop median, as it's an approximation. + bf_median = bf_result.loc["median", :] + bf_result = bf_result.drop(labels=["median"]) + pd_result = pd_result.drop(labels=["median"]) + + pd.testing.assert_frame_equal(pd_result, bf_result, check_index_type=False) + + # Double-check that median is at least plausible. + assert ( + (bf_result.loc["min", :] <= bf_median) & (bf_median <= bf_result.loc["max", :]) + ).all() + + +def test_df_describe(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + # pyarrows time columns fail in pandas + unsupported_columns = ["datetime_col", "timestamp_col", "time_col", "date_col"] + bf_result = scalars_df.describe().to_pandas() + + modified_pd_df = scalars_pandas_df.drop(columns=unsupported_columns) + pd_result = modified_pd_df.describe() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + # Drop quartiles, as they are approximate + bf_min = bf_result.loc["min", :] + bf_p25 = bf_result.loc["25%", :] + bf_p50 = bf_result.loc["50%", :] + bf_p75 = bf_result.loc["75%", :] + bf_max = bf_result.loc["max", :] + + bf_result = bf_result.drop(labels=["25%", "50%", "75%"]) + pd_result = pd_result.drop(labels=["25%", "50%", "75%"]) + + pd.testing.assert_frame_equal(pd_result, bf_result, check_index_type=False) + + # Double-check that quantiles are at least plausible. + assert ( + (bf_min <= bf_p25) + & (bf_p25 <= bf_p50) + & (bf_p50 <= bf_p50) + & (bf_p75 <= bf_max) + ).all() + + def test_ipython_key_completions_with_drop(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_names = "string_col" @@ -1206,7 +1401,7 @@ def test__dir__with_rename(scalars_dfs): ], ) def test_iloc_slice(scalars_df_index, scalars_pandas_df_index, start, stop, step): - bf_result = scalars_df_index.iloc[start:stop:step].compute() + bf_result = scalars_df_index.iloc[start:stop:step].to_pandas() pd_result = scalars_pandas_df_index.iloc[start:stop:step] # Pandas may assign non-object dtype to empty series and series index @@ -1230,7 +1425,7 @@ def test_iloc_slice_zero_step(scalars_df_index): def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.iloc[1:].iloc[1:].compute() + bf_result = scalars_df_index.iloc[1:].iloc[1:].to_pandas() pd_result = scalars_pandas_df_index.iloc[1:].iloc[1:] pd.testing.assert_frame_equal( @@ -1241,7 +1436,7 @@ def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index): @pytest.mark.parametrize( "index", - [0, 5], + [0, 5, -2], ) def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): bf_result = scalars_df_index.iloc[index] @@ -1261,7 +1456,7 @@ def test_iloc_single_integer_out_of_bound_error( def test_loc_bool_series(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.loc[scalars_df_index.bool_col].compute() + bf_result = scalars_df_index.loc[scalars_df_index.bool_col].to_pandas() pd_result = scalars_pandas_df_index.loc[scalars_pandas_df_index.bool_col] pd.testing.assert_frame_equal( @@ -1271,7 +1466,7 @@ def test_loc_bool_series(scalars_df_index, scalars_pandas_df_index): def test_loc_select_column(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.loc[:, "int64_col"].compute() + bf_result = scalars_df_index.loc[:, "int64_col"].to_pandas() pd_result = scalars_pandas_df_index.loc[:, "int64_col"] pd.testing.assert_series_equal( bf_result, @@ -1279,6 +1474,32 @@ def test_loc_select_column(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_single_index_with_duplicate(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.set_index("string_col", drop=False) + scalars_pandas_df_index = scalars_pandas_df_index.set_index( + "string_col", drop=False + ) + index = "Hello, World!" + bf_result = scalars_df_index.loc[index] + pd_result = scalars_pandas_df_index.loc[index] + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.set_index("int64_too", drop=False) + scalars_pandas_df_index = scalars_pandas_df_index.set_index("int64_too", drop=False) + index = -2345 + bf_result = scalars_df_index.loc[index] + pd_result = scalars_pandas_df_index.loc[index] + pd.testing.assert_series_equal( + bf_result.to_pandas().iloc[0, :], + pd_result, + ) + + @pytest.mark.parametrize( ("op"), [ @@ -1297,7 +1518,7 @@ def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op): col_names = ["int64_too", "float64_col", "string_col", "int64_col", "bool_col"] bf_series = op(scalars_df_index[col_names]) pd_series = op(scalars_pandas_df_index[col_names]) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() # Pandas may produce narrower numeric types, but bigframes always produces Float64 pd_series = pd_series.astype("Float64") @@ -1305,6 +1526,21 @@ def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op): pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) +def test_dataframe_aggregates_median(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col"] + bf_result = scalars_df_index[col_names].median(numeric_only=True).to_pandas() + pd_result = scalars_pandas_df_index[col_names].agg(["min", "max"]) + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + # Median is an approximation, but double-check that median is plausible. + for col in col_names: + assert (pd_result.loc["min", col] <= bf_result[col]) and ( + bf_result[col] <= pd_result.loc["max", col] + ) + + @pytest.mark.parametrize( ("op"), [ @@ -1320,7 +1556,7 @@ def test_dataframe_bool_aggregates(scalars_df_index, scalars_pandas_df_index, op ) bf_series = op(scalars_df_index) pd_series = op(scalars_pandas_df_index).astype("boolean") - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() # Pandas has object index type pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) @@ -1330,7 +1566,7 @@ def test_dataframe_prod(scalars_df_index, scalars_pandas_df_index): col_names = ["int64_too", "float64_col"] bf_series = scalars_df_index[col_names].prod() pd_series = scalars_pandas_df_index[col_names].prod() - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() # Pandas may produce narrower numeric types, but bigframes always produces Float64 pd_series = pd_series.astype("Float64") @@ -1358,7 +1594,7 @@ def test_dataframe_prod(scalars_df_index, scalars_pandas_df_index): def test_sample(scalars_dfs, frac, n, random_state): scalars_df, _ = scalars_dfs df = scalars_df.sample(frac=frac, n=n, random_state=random_state) - bf_result = df.compute() + bf_result = df.to_pandas() n = 1 if n is None else n expected_sample_size = round(frac * scalars_df.shape[0]) if frac is not None else n @@ -1384,7 +1620,7 @@ def test_sample_raises_value_error(scalars_dfs): def test_df_add_prefix(scalars_df_index, scalars_pandas_df_index, axis): if pd.__version__.startswith("1."): pytest.skip("add_prefix axis parameter not supported in pandas 1.x.") - bf_result = scalars_df_index.add_prefix("prefix_", axis).compute() + bf_result = scalars_df_index.add_prefix("prefix_", axis).to_pandas() pd_result = scalars_pandas_df_index.add_prefix("prefix_", axis) @@ -1405,7 +1641,7 @@ def test_df_add_prefix(scalars_df_index, scalars_pandas_df_index, axis): def test_df_add_suffix(scalars_df_index, scalars_pandas_df_index, axis): if pd.__version__.startswith("1."): pytest.skip("add_prefix axis parameter not supported in pandas 1.x.") - bf_result = scalars_df_index.add_suffix("_suffix", axis).compute() + bf_result = scalars_df_index.add_suffix("_suffix", axis).to_pandas() pd_result = scalars_pandas_df_index.add_suffix("_suffix", axis) @@ -1462,11 +1698,11 @@ def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index): scalars_df_index = scalars_df_index.set_index("string_col") scalars_pandas_df_index = scalars_pandas_df_index.set_index("string_col") - bf_result = scalars_df_index.loc[index_list] + bf_result = scalars_df_index.loc[index_list].to_pandas() pd_result = scalars_pandas_df_index.loc[index_list] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result, pd_result, ) @@ -1478,7 +1714,7 @@ def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.loc[index_list] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1490,7 +1726,7 @@ def test_iloc_list(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.iloc[index_list] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1501,7 +1737,7 @@ def test_iloc_empty_list(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.iloc[index_list] pd_result = scalars_pandas_df_index.iloc[index_list] - bf_result = bf_result.compute() + bf_result = bf_result.to_pandas() assert bf_result.shape == pd_result.shape # types are known to be different @@ -1510,7 +1746,7 @@ def test_rename_axis(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.rename_axis("newindexname") pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1526,7 +1762,7 @@ def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.loc[pd_string_series] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1539,7 +1775,7 @@ def test_loc_bf_index_integer_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.loc[pd_index] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1559,7 +1795,7 @@ def test_loc_bf_index_integer_index_renamed_col( pd_result = scalars_pandas_df_index.loc[pd_index] pd.testing.assert_frame_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1568,6 +1804,7 @@ def test_loc_bf_index_integer_index_renamed_col( ("subset"), [ None, + "bool_col", ["bool_col", "int64_too"], ], ) @@ -1581,7 +1818,7 @@ def test_loc_bf_index_integer_index_renamed_col( ) def test_df_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, subset): columns = ["bool_col", "int64_too", "int64_col"] - bf_series = scalars_df_index[columns].drop_duplicates(subset, keep=keep).compute() + bf_series = scalars_df_index[columns].drop_duplicates(subset, keep=keep).to_pandas() pd_series = scalars_pandas_df_index[columns].drop_duplicates(subset, keep=keep) pd.testing.assert_frame_equal( pd_series, @@ -1606,7 +1843,7 @@ def test_df_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, sub ) def test_df_duplicated(scalars_df_index, scalars_pandas_df_index, keep, subset): columns = ["bool_col", "int64_too", "int64_col"] - bf_series = scalars_df_index[columns].duplicated(subset, keep=keep).compute() + bf_series = scalars_df_index[columns].duplicated(subset, keep=keep).to_pandas() pd_series = scalars_pandas_df_index[columns].duplicated(subset, keep=keep) pd.testing.assert_series_equal(pd_series, bf_series, check_dtype=False) @@ -1625,7 +1862,7 @@ def test_df_value_counts(scalars_dfs, subset, normalize, ascending, dropna): bf_result = ( scalars_df[["string_col", "bool_col"]] .value_counts(subset, normalize=normalize, ascending=ascending, dropna=dropna) - .compute() + .to_pandas() ) pd_result = scalars_pandas_df[["string_col", "bool_col"]].value_counts( subset, normalize=normalize, ascending=ascending, dropna=dropna @@ -1670,7 +1907,7 @@ def test_df_rank_with_nulls( ascending=ascending, numeric_only=numeric_only, ) - .compute() + .to_pandas() ) pd_result = ( scalars_pandas_df_index.drop(columns=unsupported_columns) @@ -1692,3 +1929,14 @@ def test_df_rank_with_nulls( def test_df_bool_interpretation_error(scalars_df_index): with pytest.raises(ValueError): True if scalars_df_index else False + + +def test_query_job_setters(scalars_df_default_index: dataframe.DataFrame): + job_ids = set() + repr(scalars_df_default_index) + assert scalars_df_default_index.query_job is not None + job_ids.add(scalars_df_default_index.query_job.job_id) + scalars_df_default_index.to_pandas() + job_ids.add(scalars_df_default_index.query_job.job_id) + + assert len(job_ids) == 2 diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index a4318a8658f..7ad753e1bc0 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -72,15 +72,17 @@ def test_to_csv_index( gcs_folder: str, index: bool, ): + if pd.__version__.startswith("1."): + pytest.skip("date_format parameter not supported in pandas 1.x.") """Test the `to_csv` API with the `index` parameter.""" scalars_df, scalars_pandas_df = scalars_dfs index_col = None if scalars_df.index.name is not None: - path = gcs_folder + f"test_index_df_to_csv_index_{index}" + path = gcs_folder + f"test_index_df_to_csv_index_{index}*.csv" if index: index_col = scalars_df.index.name else: - path = gcs_folder + f"test_default_index_df_to_csv_index_{index}" + path = gcs_folder + f"test_default_index_df_to_csv_index_{index}*.csv" # TODO(swast): Support "date_format" parameter and make sure our # DATETIME/TIMESTAMP column export is the same format as pandas by default. @@ -90,14 +92,60 @@ def test_to_csv_index( # BigQuery-backed dataframes, so manually convert the dtypes specifically # here. dtype = scalars_df.reset_index().dtypes.to_dict() - dtype.pop("timestamp_col") dtype.pop("geography_col") + dtype.pop("rowindex") gcs_df = pd.read_csv( - path, dtype=dtype, parse_dates=["timestamp_col"], index_col=index_col + path, + dtype=dtype, + date_format={"timestamp_col": "YYYY-MM-DD HH:MM:SS Z"}, + index_col=index_col, ) convert_pandas_dtypes(gcs_df, bytes_col=True) + gcs_df.index.name = scalars_df.index.name - assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + scalars_pandas_df = scalars_pandas_df.copy() + scalars_pandas_df.index = scalars_pandas_df.index.astype("int64") + + # Ordering should be maintained for tables smaller than 1 GB. + pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df) + + +def test_to_csv_tabs( + scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame], + gcs_folder: str, +): + if pd.__version__.startswith("1."): + pytest.skip("date_format parameter not supported in pandas 1.x.") + """Test the `to_csv` API with the `sep` parameter.""" + scalars_df, scalars_pandas_df = scalars_dfs + index_col = scalars_df.index.name + path = gcs_folder + "test_to_csv_tabs*.csv" + + # TODO(swast): Support "date_format" parameter and make sure our + # DATETIME/TIMESTAMP column export is the same format as pandas by default. + scalars_df.to_csv(path, sep="\t", index=True) + + # Pandas dataframes dtypes from read_csv are not fully compatible with + # BigQuery-backed dataframes, so manually convert the dtypes specifically + # here. + dtype = scalars_df.reset_index().dtypes.to_dict() + dtype.pop("geography_col") + dtype.pop("rowindex") + gcs_df = pd.read_csv( + path, + sep="\t", + dtype=dtype, + date_format={"timestamp_col": "YYYY-MM-DD HH:MM:SS Z"}, + index_col=index_col, + ) + convert_pandas_dtypes(gcs_df, bytes_col=True) + gcs_df.index.name = scalars_df.index.name + + scalars_pandas_df = scalars_pandas_df.copy() + scalars_pandas_df.index = scalars_pandas_df.index.astype("int64") + + # Ordering should be maintained for tables smaller than 1 GB. + pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df) @pytest.mark.parametrize( @@ -190,9 +238,9 @@ def test_to_json_index_invalid_orient( ): scalars_df, scalars_pandas_df = scalars_dfs if scalars_df.index.name is not None: - path = gcs_folder + f"test_index_df_to_json_index_{index}" + path = gcs_folder + f"test_index_df_to_json_index_{index}*.jsonl" else: - path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + path = gcs_folder + f"test_default_index_df_to_json_index_{index}*.jsonl" with pytest.raises(ValueError): scalars_df.to_json(path, index=index, lines=True) @@ -208,9 +256,9 @@ def test_to_json_index_invalid_lines( ): scalars_df, scalars_pandas_df = scalars_dfs if scalars_df.index.name is not None: - path = gcs_folder + f"test_index_df_to_json_index_{index}" + path = gcs_folder + f"test_index_df_to_json_index_{index}.jsonl" else: - path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + path = gcs_folder + f"test_default_index_df_to_json_index_{index}.jsonl" with pytest.raises(NotImplementedError): scalars_df.to_json(path, index=index) @@ -227,9 +275,9 @@ def test_to_json_index_records_orient( """Test the `to_json` API with the `index` parameter.""" scalars_df, scalars_pandas_df = scalars_dfs if scalars_df.index.name is not None: - path = gcs_folder + f"test_index_df_to_json_index_{index}" + path = gcs_folder + f"test_index_df_to_json_index_{index}*.jsonl" else: - path = gcs_folder + f"test_default_index_df_to_json_index_{index}" + path = gcs_folder + f"test_default_index_df_to_json_index_{index}*.jsonl" """ Test the `to_json` API with `orient` is `records` and `lines` is True""" scalars_df.to_json(path, index=index, orient="records", lines=True) @@ -241,7 +289,13 @@ def test_to_json_index_records_orient( assert len(gcs_df.index) == len(scalars_pandas_df.index) pd.testing.assert_index_equal(gcs_df.columns, scalars_pandas_df.columns) - assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + + gcs_df.index.name = scalars_df.index.name + gcs_df.index = gcs_df.index.astype("Int64") + scalars_pandas_df.index = scalars_pandas_df.index.astype("Int64") + + # Ordering should be maintained for tables smaller than 1 GB. + pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df) @pytest.mark.parametrize( @@ -251,10 +305,12 @@ def test_to_json_index_records_orient( def test_to_parquet_index(scalars_dfs, gcs_folder, index): """Test the `to_parquet` API with the `index` parameter.""" scalars_df, scalars_pandas_df = scalars_dfs + scalars_pandas_df = scalars_pandas_df.copy() + if scalars_df.index.name is not None: - path = gcs_folder + f"test_index_df_to_parquet_{index}" + path = gcs_folder + f"test_index_df_to_parquet_{index}*.parquet" else: - path = gcs_folder + f"test_default_index_df_to_parquet_{index}" + path = gcs_folder + f"test_default_index_df_to_parquet_{index}*.parquet" # TODO(b/268693993): Type GEOGRAPHY is not currently supported for parquet. scalars_df = scalars_df.drop(columns="geography_col") @@ -265,20 +321,26 @@ def test_to_parquet_index(scalars_dfs, gcs_folder, index): # table. scalars_df.to_parquet(path, index=index) - gcs_df = pd.read_parquet(path) + gcs_df = pd.read_parquet(path.replace("*", "000000000000")) convert_pandas_dtypes(gcs_df, bytes_col=False) if index and scalars_df.index.name is not None: gcs_df = gcs_df.set_index(scalars_df.index.name) assert len(gcs_df.index) == len(scalars_pandas_df.index) pd.testing.assert_index_equal(gcs_df.columns, scalars_pandas_df.columns) - assert_pandas_df_equal_ignore_ordering(gcs_df, scalars_pandas_df) + + gcs_df.index.name = scalars_df.index.name + gcs_df.index = gcs_df.index.astype("Int64") + scalars_pandas_df.index = scalars_pandas_df.index.astype("Int64") + + # Ordering should be maintained for tables smaller than 1 GB. + pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df) def test_to_sql_query_named_index_included( session, scalars_df_index, scalars_pandas_df_index ): - sql, index_columns = scalars_df_index.to_sql_query(always_include_index=True) + sql, index_columns = scalars_df_index._to_sql_query(always_include_index=True) assert len(index_columns) == 1 index_column, is_named = index_columns[0] assert index_column == "rowindex" @@ -294,7 +356,7 @@ def test_to_sql_query_unnamed_index_excluded( session, scalars_df_default_index, scalars_pandas_df_default_index ): # The .sql property should return SQL without the unnamed indexes - sql, index_columns = scalars_df_default_index.to_sql_query( + sql, index_columns = scalars_df_default_index._to_sql_query( always_include_index=False ) assert len(index_columns) == 0 @@ -310,7 +372,7 @@ def test_to_sql_query_unnamed_index_always_include( scalars_df_default_index: bigframes.dataframe.DataFrame, scalars_pandas_df_default_index, ): - sql, index_columns = scalars_df_default_index.to_sql_query( + sql, index_columns = scalars_df_default_index._to_sql_query( always_include_index=True ) assert len(index_columns) == 1 diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 066e20bb12c..e72d75729be 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -41,11 +41,27 @@ def test_dataframe_groupby_numeric_aggregate( col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] bf_result = operator(scalars_df_index[col_names].groupby("string_col")) pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) - bf_result_computed = bf_result.compute() + bf_result_computed = bf_result.to_pandas() # Pandas std function produces float64, not matching Float64 from bigframes pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) +def test_dataframe_groupby_median(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names].groupby("string_col").median(numeric_only=True) + ) + pd_min = ( + scalars_pandas_df_index[col_names].groupby("string_col").min(numeric_only=True) + ) + pd_max = ( + scalars_pandas_df_index[col_names].groupby("string_col").max(numeric_only=True) + ) + bf_result_computed = bf_result.to_pandas() + # Median is approximate. Just check for plausibility. + assert ((pd_min <= bf_result_computed) & (bf_result_computed <= pd_max)).all().all() + + @pytest.mark.parametrize( ("operator"), [ @@ -65,7 +81,7 @@ def test_dataframe_groupby_aggregate( col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] bf_result = operator(scalars_df_index[col_names].groupby("string_col")) pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) - bf_result_computed = bf_result.compute() + bf_result_computed = bf_result.to_pandas() pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) @@ -91,7 +107,7 @@ def test_dataframe_groupby_multi_sum( .groupby(["bool_col", "int64_col"], as_index=as_index) .sum(numeric_only=True) ) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() if not as_index: # BigQuery DataFrames default indices use nullable Int64 always @@ -124,6 +140,36 @@ def test_dataframe_groupby_analytic( col_names = ["float64_col", "int64_col", "bool_col", "string_col"] bf_result = operator(scalars_df_index[col_names].groupby("string_col")) pd_result = operator(scalars_pandas_df_index[col_names].groupby("string_col")) - bf_result_computed = bf_result.compute() + bf_result_computed = bf_result.to_pandas() pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + +def test_dataframe_groupby_getitem( + scalars_df_index, + scalars_pandas_df_index, +): + col_names = ["float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names].groupby("string_col")["int64_col"].min().to_pandas() + ) + pd_result = ( + scalars_pandas_df_index[col_names].groupby("string_col")["int64_col"].min() + ) + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + +def test_dataframe_groupby_getitem_list( + scalars_df_index, + scalars_pandas_df_index, +): + col_names = ["float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names].groupby("string_col")[col_names].min().to_pandas() + ) + pd_result = ( + scalars_pandas_df_index[col_names].groupby("string_col")[col_names].min() + ) + + pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) diff --git a/tests/system/small/test_ibis.py b/tests/system/small/test_ibis.py new file mode 100644 index 00000000000..58b78e00481 --- /dev/null +++ b/tests/system/small/test_ibis.py @@ -0,0 +1,39 @@ +# 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. + +"""Tests for monkeypatched ibis code.""" + +import ibis.expr.types as ibis_types + +import bigframes +import third_party.bigframes_vendored.ibis.expr.operations as vendored_ibis_ops + + +def test_approximate_quantiles(session: bigframes.Session, scalars_table_id: str): + num_bins = 3 + ibis_client = session.ibis_client + _, dataset, table_id = scalars_table_id.split(".") + ibis_table: ibis_types.Table = ibis_client.table(table_id, database=dataset) + ibis_column: ibis_types.NumericColumn = ibis_table["int64_col"] + quantiles: ibis_types.ArrayScalar = vendored_ibis_ops.ApproximateMultiQuantile( # type: ignore + ibis_column, num_bins=num_bins + ).to_expr() + value = quantiles[1] + num_edges = quantiles.length() + + sql = ibis_client.compile(value) + num_edges_result = num_edges.to_pandas() + + assert "APPROX_QUANTILES" in sql + assert num_edges_result == num_bins + 1 diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py index 312e86d7f23..ac1f8c7220d 100644 --- a/tests/system/small/test_index.py +++ b/tests/system/small/test_index.py @@ -19,7 +19,7 @@ def test_get_index(scalars_df_index, scalars_pandas_df_index): index = scalars_df_index.index - bf_result = index.compute() + bf_result = index.to_pandas() pd_result = scalars_pandas_df_index.index assert_pandas_index_equal_ignore_index_type(bf_result, pd_result) @@ -40,7 +40,13 @@ def test_index_len(scalars_df_index, scalars_pandas_df_index): def test_index_array(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.__array__() - pd_result = scalars_pandas_df_index.__array__() + bf_result = scalars_df_index.index.__array__() + pd_result = scalars_pandas_df_index.index.__array__() numpy.array_equal(bf_result, pd_result) + + +def test_index_getitem_int(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.index[-2] + pd_result = scalars_pandas_df_index.index[-2] + assert bf_result == pd_result diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index bf2a2080de3..914be6dae4a 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -16,10 +16,11 @@ import pytest import bigframes.pandas as bpd +from tests.system.utils import assert_pandas_df_equal_ignore_ordering def test_set_multi_index(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.set_index(["bool_col", "int64_too"]).compute() + bf_result = scalars_df_index.set_index(["bool_col", "int64_too"]).to_pandas() pd_result = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -27,7 +28,7 @@ def test_set_multi_index(scalars_df_index, scalars_pandas_df_index): def test_reset_multi_index(scalars_df_index, scalars_pandas_df_index): bf_result = ( - scalars_df_index.set_index(["bool_col", "int64_too"]).reset_index().compute() + scalars_df_index.set_index(["bool_col", "int64_too"]).reset_index().to_pandas() ) pd_result = scalars_pandas_df_index.set_index( ["bool_col", "int64_too"] @@ -51,7 +52,7 @@ def test_binop_series_series_matching_multi_indices( pd_result = pd_left["int64_col"] + pd_right["int64_too"] pandas.testing.assert_series_equal( - bf_result.sort_index().compute(), pd_result.sort_index() + bf_result.sort_index().to_pandas(), pd_result.sort_index() ) @@ -67,7 +68,7 @@ def test_binop_df_series_matching_multi_indices( pd_result = pd_left[["int64_col", "int64_too"]].add(pd_right["int64_too"], axis=0) pandas.testing.assert_frame_equal( - bf_result.sort_index().compute(), pd_result.sort_index() + bf_result.sort_index().to_pandas(), pd_result.sort_index() ) @@ -80,7 +81,7 @@ def test_binop_multi_index_mono_index(scalars_df_index, scalars_pandas_df_index) bf_result = bf_left["int64_col"] + bf_right["int64_too"] pd_result = pd_left["int64_col"] + pd_right["int64_too"] - pandas.testing.assert_series_equal(bf_result.compute(), pd_result) + pandas.testing.assert_series_equal(bf_result.to_pandas(), pd_result) def test_binop_overlapping_multi_indices(scalars_df_index, scalars_pandas_df_index): @@ -93,7 +94,7 @@ def test_binop_overlapping_multi_indices(scalars_df_index, scalars_pandas_df_ind pd_result = pd_left["int64_col"] + pd_right["int64_too"] pandas.testing.assert_series_equal( - bf_result.sort_index().compute(), pd_result.sort_index() + bf_result.sort_index().to_pandas(), pd_result.sort_index() ) @@ -108,7 +109,7 @@ def test_concat_compatible_multi_indices(scalars_df_index, scalars_pandas_df_ind bf_result = bpd.concat([bf_left, bf_right]) pd_result = pandas.concat([pd_left, pd_right]) - pandas.testing.assert_frame_equal(bf_result.compute(), pd_result) + pandas.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) def test_concat_multi_indices_ignore_index(scalars_df_index, scalars_pandas_df_index): @@ -123,12 +124,12 @@ def test_concat_multi_indices_ignore_index(scalars_df_index, scalars_pandas_df_i # Pandas uses int64 instead of Int64 (nullable) dtype. pd_result.index = pd_result.index.astype(pandas.Int64Dtype()) - pandas.testing.assert_frame_equal(bf_result.compute(), pd_result) + pandas.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) def test_multi_index_loc(scalars_df_index, scalars_pandas_df_index): bf_result = ( - scalars_df_index.set_index(["int64_too", "bool_col"]).loc[[2, 0]].compute() + scalars_df_index.set_index(["int64_too", "bool_col"]).loc[[2, 0]].to_pandas() ) pd_result = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]).loc[[2, 0]] @@ -139,7 +140,7 @@ def test_multi_index_getitem_bool(scalars_df_index, scalars_pandas_df_index): bf_frame = scalars_df_index.set_index(["int64_too", "bool_col"]) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) - bf_result = bf_frame[bf_frame["int64_col"] > 0].compute() + bf_result = bf_frame[bf_frame["int64_col"] > 0].to_pandas() pd_result = pd_frame[pd_frame["int64_col"] > 0] pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -159,12 +160,30 @@ def test_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level) bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) - bf_result = bf_frame.droplevel(level).compute() + bf_result = bf_frame.droplevel(level).to_pandas() pd_result = pd_frame.droplevel(level) pandas.testing.assert_frame_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("labels", "level"), + [ + (1, 0), + ([0, 1], 0), + ([True, None], 1), + ], +) +def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, level): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame.drop(labels=labels, axis="index", level=level).to_pandas() + pd_result = pd_frame.drop(labels=labels, axis="index", level=level) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + @pytest.mark.parametrize( ("order"), [ @@ -182,7 +201,7 @@ def test_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_index, o bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) - bf_result = bf_frame.reorder_levels(order).compute() + bf_result = bf_frame.reorder_levels(order).to_pandas() pd_result = pd_frame.reorder_levels(order) pandas.testing.assert_frame_equal(bf_result, pd_result) @@ -194,7 +213,7 @@ def test_multi_index_series_groupby(scalars_df_index, scalars_pandas_df_index): bf_frame["float64_col"] .groupby([bf_frame.int64_col % 2, "bool_col"]) .mean() - .compute() + .to_pandas() ) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) pd_result = ( @@ -220,7 +239,7 @@ def test_multi_index_series_groupby_level( scalars_df_index.set_index(["int64_too", "bool_col"])["float64_col"] .groupby(level=level) .mean() - .compute() + .to_pandas() ) pd_result = ( scalars_pandas_df_index.set_index(["int64_too", "bool_col"])["float64_col"] @@ -236,7 +255,7 @@ def test_multi_index_dataframe_groupby(scalars_df_index, scalars_pandas_df_index bf_result = ( bf_frame.groupby([bf_frame.int64_col % 2, "bool_col"]) .mean(numeric_only=True) - .compute() + .to_pandas() ) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) pd_result = pd_frame.groupby([pd_frame.int64_col % 2, "bool_col"]).mean( @@ -247,27 +266,180 @@ def test_multi_index_dataframe_groupby(scalars_df_index, scalars_pandas_df_index @pytest.mark.parametrize( - ("level"), + ("level", "as_index"), [ - (1), - ([0]), - (["bool_col"]), - (["bool_col", "int64_too"]), + (1, True), + ([0], False), + (["bool_col"], True), + (["bool_col", "int64_too"], False), ], ) -def test_multi_index_dataframe_groupby_level( - scalars_df_index, scalars_pandas_df_index, level +def test_multi_index_dataframe_groupby_level_aggregate( + scalars_df_index, scalars_pandas_df_index, level, as_index ): bf_result = ( scalars_df_index.set_index(["int64_too", "bool_col"]) - .groupby(level=level) + .groupby(level=level, as_index=as_index) .mean(numeric_only=True) - .compute() + .to_pandas() ) pd_result = ( scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) - .groupby(level=level) + .groupby(level=level, as_index=as_index) .mean(numeric_only=True) ) - pandas.testing.assert_frame_equal(bf_result, pd_result) + # Pandas will have int64 index, while bigquery will have Int64 when resetting + pandas.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) + + +@pytest.mark.parametrize( + ("level", "as_index"), + [ + (1, True), + ([0], False), + ( + ["bool_col"], + True, + ), + (["bool_col", "int64_too"], False), + ], +) +def test_multi_index_dataframe_groupby_level_analytic( + scalars_df_index, scalars_pandas_df_index, level, as_index +): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"]) + .groupby(level=level, as_index=as_index, dropna=False) + .cumsum(numeric_only=True) + .to_pandas() + ) + pd_result = ( + scalars_pandas_df_index.set_index(["int64_too", "bool_col"]) + .groupby(level=level, as_index=as_index, dropna=False) + .cumsum(numeric_only=True) + ) + + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +all_joins = pytest.mark.parametrize( + ("how",), + ( + ("outer",), + ("left",), + ("right",), + ("inner",), + ), +) + + +@all_joins +# Both DFs are multi-index +def test_multi_index_dataframe_join(scalars_dfs, how): + bf_df, pd_df = scalars_dfs + + bf_df_a = bf_df.set_index((["bool_col", "rowindex_2"]))[["string_col", "int64_col"]] + bf_df_b = bf_df.assign(rowindex_2=bf_df["rowindex_2"] + 2).set_index( + (["bool_col", "rowindex_2"]) + )[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, how=how).to_pandas() + + pd_df_a = pd_df.set_index((["bool_col", "rowindex_2"]))[["string_col", "int64_col"]] + pd_df_b = pd_df.assign(rowindex_2=pd_df["rowindex_2"] + 2).set_index( + (["bool_col", "rowindex_2"]) + )[["float64_col"]] + pd_result = pd_df_a.join(pd_df_b, how=how) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@all_joins +# Only left DF is multi-index +def test_multi_index_dataframe_join_on(scalars_dfs, how): + bf_df, pd_df = scalars_dfs + + bf_df_a = bf_df.set_index((["int64_too", "bool_col"]))[ + ["string_col", "int64_col", "rowindex_2"] + ] + bf_df_a = bf_df_a.assign(rowindex_2=bf_df_a["rowindex_2"] + 2) + bf_df_b = bf_df[["float64_col"]] + bf_result = bf_df_a.join(bf_df_b, on="rowindex_2", how=how).to_pandas() + + pd_df_a = pd_df.set_index((["int64_too", "bool_col"]))[ + ["string_col", "int64_col", "rowindex_2"] + ] + pd_df_a = pd_df_a.assign(rowindex_2=pd_df_a["rowindex_2"] + 2) + pd_df_b = pd_df[["float64_col"]] + pd_result = pd_df_a.join(pd_df_b, on="rowindex_2", how=how) + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("level",), + [ + (1,), + ([0],), + (["bool_col"],), + (["bool_col", "int64_too"],), + ], +) +def test_multi_index_series_groupby_level_aggregate( + scalars_df_index, scalars_pandas_df_index, level +): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level) + .mean() + .to_pandas() + ) + pd_result = ( + scalars_pandas_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level) + .mean() + ) + + pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("level",), + [ + (1,), + ([0],), + (["bool_col"],), + (["bool_col", "int64_too"],), + ], +) +def test_multi_index_series_groupby_level_analytic( + scalars_df_index, scalars_pandas_df_index, level +): + bf_result = ( + scalars_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level, dropna=False) + .cumsum() + .to_pandas() + ) + pd_result = ( + scalars_pandas_df_index.set_index(["int64_too", "bool_col"])["float64_col"] + .groupby(level=level, dropna=False) + .cumsum() + ) + + pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + +def test_multi_index_series_rename_dict_same_type( + scalars_df_index, scalars_pandas_df_index +): + bf_result = ( + scalars_df_index.set_index(["rowindex_2", "int64_too"])["string_col"] + .rename({1: 100, 2: 200}) + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index(["rowindex_2", "int64_too"])[ + "string_col" + ].rename({1: 100, 2: 200}) + + pandas.testing.assert_series_equal( + bf_result, pd_result, check_dtype=False, check_index_type=False + ) diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index a612ad946fa..98bafc6392c 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -21,7 +21,7 @@ def test_concat_dataframe(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = bpd.concat(11 * [scalars_df]) - bf_result = bf_result.compute() + bf_result = bf_result.to_pandas() pd_result = pd.concat(11 * [scalars_pandas_df]) pd.testing.assert_frame_equal(bf_result, pd_result) @@ -32,7 +32,7 @@ def test_concat_series(scalars_dfs): bf_result = bpd.concat( [scalars_df.int64_col, scalars_df.int64_too, scalars_df.int64_col] ) - bf_result = bf_result.compute() + bf_result = bf_result.to_pandas() pd_result = pd.concat( [ scalars_pandas_df.int64_col, @@ -45,10 +45,10 @@ def test_concat_series(scalars_dfs): @pytest.mark.parametrize( - ("how",), + ("how"), [ - ("inner",), - ("outer",), + ("inner"), + ("outer"), ], ) def test_concat_dataframe_mismatched_columns(scalars_dfs, how): @@ -56,9 +56,52 @@ def test_concat_dataframe_mismatched_columns(scalars_dfs, how): cols2 = ["int64_col", "string_col", "int64_too"] scalars_df, scalars_pandas_df = scalars_dfs bf_result = bpd.concat([scalars_df[cols1], scalars_df[cols2]], join=how) - bf_result = bf_result.compute() + bf_result = bf_result.to_pandas() pd_result = pd.concat( - [scalars_pandas_df[cols1], scalars_pandas_df[cols2]], join=how + [scalars_pandas_df[cols1], scalars_pandas_df[cols2]], + join=how, ) pd.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("how",), + [ + ("inner",), + ("outer",), + ], +) +def test_concat_axis_1(scalars_dfs, how): + if pd.__version__.startswith("1."): + pytest.skip("pandas has different behavior in 1.x") + scalars_df, scalars_pandas_df = scalars_dfs + cols1 = ["int64_col", "float64_col", "rowindex_2"] + cols2 = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + part1 = scalars_df[cols1] + part1.index.name = "newindexname" + # Offset the rows somewhat so that outer join can have an effect. + part2 = ( + scalars_df[cols2] + .assign(rowindex_2=scalars_df["rowindex_2"] + 2) + .sort_values(["string_col"], kind="stable") + ) + part3 = scalars_df["int64_too"].cumsum().iloc[2:] + + bf_result = bpd.concat([part1, part2, part3], join=how, axis=1) + + # Copy since modifying index + pd_part1 = scalars_pandas_df.copy()[cols1] + pd_part1.index.name = "newindexname" + # Offset the rows somewhat so that outer join can have an effect. + pd_part2 = ( + scalars_pandas_df[cols2] + .assign(rowindex_2=scalars_pandas_df["rowindex_2"] + 2) + .sort_values(["string_col"], kind="stable") + ) + pd_part3 = scalars_pandas_df["int64_too"].cumsum().iloc[2:] + + pd_result = pd.concat([pd_part1, pd_part2, pd_part3], join=how, axis=1) + + pd.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index 261e0d5b14e..96697dbcabe 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime +from unittest import mock + import google.api_core.exceptions +import google.auth +import google.auth.exceptions import pytest import bigframes.pandas as bpd @@ -79,6 +84,10 @@ def test_read_gbq_start_sets_session_location( # There should still be the previous location set in the bigquery options assert bpd.options.bigquery.location == tokyo_location + # Reset the location to be able to query another location + bpd.options.bigquery.location = None + assert not bpd.options.bigquery.location + # Starting over the user journey with read_gbq* should work for a table # in another location, in this case US df = read_method(query) @@ -143,3 +152,188 @@ def test_read_gbq_after_session_start_must_comply_with_default_location( # read_gbq* from a table in the default location should work df = read_method(query) assert df is not None + + +@pytest.mark.parametrize( + ("read_method", "query_prefix"), + [ + (bpd.read_gbq, None), + (bpd.read_gbq, "SELECT COUNT(1) FROM "), + (bpd.read_gbq_table, None), + (bpd.read_gbq_query, "SELECT COUNT(1) FROM "), + ], + ids=[ + "read_gbq-on-table-name", + "read_gbq-on-sql", + "read_gbq_table-on-table-name", + "read_gbq_query-on-sql", + ], +) +def test_read_gbq_must_comply_with_set_location_US( + test_data_tables, + test_data_tables_tokyo, + dataset_id_permanent_tokyo, + read_method, + query_prefix, +): + # Form query as a table name or a SQL depending on the test scenario + query_tokyo = test_data_tables_tokyo["scalars"] + query = test_data_tables["scalars"] + if query_prefix: + query_tokyo = f"{query_prefix} {query_tokyo}" + query = f"{query_prefix} {query}" + + # Initially there is no location set in the bigquery options + assert not bpd.options.bigquery.location + + # Explicitly set location + bpd.options.bigquery.location = "US" + assert bpd.options.bigquery.location == "US" + + # Starting user journey with read_gbq* from another location should fail + with pytest.raises( + google.api_core.exceptions.NotFound, + match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + ): + read_method(query_tokyo) + + # Starting user journey with read_gbq* should work for a table in the same + # location, in this case tokyo + df = read_method(query) + assert df is not None + + +@pytest.mark.parametrize( + ("read_method", "query_prefix"), + [ + (bpd.read_gbq, None), + (bpd.read_gbq, "SELECT COUNT(1) FROM "), + (bpd.read_gbq_table, None), + (bpd.read_gbq_query, "SELECT COUNT(1) FROM "), + ], + ids=[ + "read_gbq-on-table-name", + "read_gbq-on-sql", + "read_gbq_table-on-table-name", + "read_gbq_query-on-sql", + ], +) +def test_read_gbq_must_comply_with_set_location_non_US( + tokyo_location, + test_data_tables, + test_data_tables_tokyo, + dataset_id_permanent, + read_method, + query_prefix, +): + # Form query as a table name or a SQL depending on the test scenario + query_tokyo = test_data_tables_tokyo["scalars"] + query = test_data_tables["scalars"] + if query_prefix: + query_tokyo = f"{query_prefix} {query_tokyo}" + query = f"{query_prefix} {query}" + + # Initially there is no location set in the bigquery options + assert not bpd.options.bigquery.location + + # Explicitly set location + bpd.options.bigquery.location = tokyo_location + assert bpd.options.bigquery.location == tokyo_location + + # Starting user journey with read_gbq* from another location should fail + with pytest.raises( + google.api_core.exceptions.NotFound, + match=f"404 Not found: Dataset {dataset_id_permanent} was not found in location {tokyo_location}", + ): + read_method(query) + + # Starting user journey with read_gbq* should work for a table in the same + # location, in this case tokyo + df = read_method(query_tokyo) + assert df is not None + + +def test_reset_session_after_bq_session_ended(): + # Use a simple test query to verify that default session works to interact + # with BQ + test_query = "SELECT 1" + + # Confirm that there is a session id in the default session + session = bpd.get_global_session() + assert session._session_id + + # Confirm that session works as usual + df = bpd.read_gbq(test_query) + assert df is not None + + # Abort the session to simulate the auto-expiration + # https://cloud.google.com/bigquery/docs/sessions-terminating#auto-terminate_a_session + abort_session_query = "CALL BQ.ABORT_SESSION()" + query_job = session.bqclient.query(abort_session_query) + query_job.result() # blocks until finished + + # Confirm that session is unusable to run any jobs + with pytest.raises( + google.api_core.exceptions.BadRequest, + match=f"Session {session._session_id} has expired and is no longer available.", + ): + query_job = session.bqclient.query(test_query) + query_job.result() # blocks until finished + + # Confirm that as a result bigframes.pandas interface is unusable + with pytest.raises( + google.api_core.exceptions.BadRequest, + match=f"Session {session._session_id} has expired and is no longer available.", + ): + bpd.read_gbq(test_query) + + # Now try to reset session and verify that it works + bpd.reset_session() + assert bpd._global_session is None + + # Now verify that use is able to start over + df = bpd.read_gbq(test_query) + assert df is not None + + +def test_reset_session_after_credentials_need_reauthentication(monkeypatch): + # Use a simple test query to verify that default session works to interact + # with BQ + test_query = "SELECT 1" + + # Confirm that default session has BQ client with valid credentials + session = bpd.get_global_session() + assert session.bqclient._credentials.valid + + # Confirm that default session works as usual + df = bpd.read_gbq(test_query) + assert df is not None + + with monkeypatch.context() as m: + # Simulate expired credentials to trigger the credential refresh flow + m.setattr(session.bqclient._credentials, "expiry", datetime.datetime.utcnow()) + assert not session.bqclient._credentials.valid + + # Simulate an exception during the credential refresh flow + m.setattr( + session.bqclient._credentials, + "refresh", + mock.Mock(side_effect=google.auth.exceptions.RefreshError()), + ) + + # Confirm that session is unusable to run any jobs + with pytest.raises(google.auth.exceptions.RefreshError): + query_job = session.bqclient.query(test_query) + query_job.result() # blocks until finished + + # Confirm that as a result bigframes.pandas interface is unusable + with pytest.raises(google.auth.exceptions.RefreshError): + bpd.read_gbq(test_query) + + # Now verify that resetting the session works + bpd.reset_session() + assert bpd._global_session is None + + # Now verify that use is able to start over + df = bpd.read_gbq(test_query) + assert df is not None diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index 6eeee3a3a43..00380c26390 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -12,6 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import tempfile + +import pandas as pd + import bigframes as bf import bigframes.formatting_helpers as formatting_helpers @@ -20,11 +24,12 @@ def test_progress_bar_dataframe( penguins_df_default_index: bf.dataframe.DataFrame, capsys ): bf.options.display.progress_bar = "notebook" - penguins_df_default_index.compute() + penguins_df_default_index.to_pandas() html_check = "HTML(value=" open_job_check = "Open Job" lines = capsys.readouterr().out.split("\n") - lines = filter(None, lines) + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 assert penguins_df_default_index.query_job is not None for line in lines: assert html_check in line and open_job_check in line @@ -33,11 +38,12 @@ def test_progress_bar_dataframe( def test_progress_bar_series(penguins_df_default_index: bf.dataframe.DataFrame, capsys): bf.options.display.progress_bar = "notebook" series = penguins_df_default_index["body_mass_g"].head(10) - series.compute() + series.to_pandas() html_check = "HTML(value=" open_job_check = "Open Job" lines = capsys.readouterr().out.split("\n") - lines = filter(None, lines) + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 assert series.query_job is not None for line in lines: assert html_check in line and open_job_check in line @@ -49,18 +55,63 @@ def test_progress_bar_scalar(penguins_df_default_index: bf.dataframe.DataFrame, html_check = "HTML(value=" open_job_check = "Open Job" lines = capsys.readouterr().out.split("\n") - lines = filter(None, lines) + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 for line in lines: assert html_check in line and open_job_check in line -def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame): +def test_progress_bar_read_gbq(session: bf.Session, penguins_table_id: str, capsys): + bf.options.display.progress_bar = "notebook" + session.read_gbq(penguins_table_id) + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 + for line in lines: + assert html_check in line and open_job_check in line + + +def test_progress_bar_extract_jobs( + penguins_df_default_index: bf.dataframe.DataFrame, gcs_folder, capsys +): + bf.options.display.progress_bar = "notebook" + path = gcs_folder + "test_read_csv_progress_bar*.csv" + penguins_df_default_index.to_csv(path) + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 + for line in lines: + assert html_check in line and open_job_check in line + + +def test_progress_bar_load_jobs( + session: bf.Session, penguins_pandas_df_default_index: pd.DataFrame, capsys +): + bf.options.display.progress_bar = "notebook" + with tempfile.TemporaryDirectory() as dir: + path = dir + "/test_read_csv_progress_bar*.csv" + penguins_pandas_df_default_index.to_csv(path, index=False) + session.read_csv(path) + html_check = "HTML(value=" + open_job_check = "Open Job" + lines = capsys.readouterr().out.split("\n") + lines = [line for line in lines if len(line) > 0] + assert len(lines) > 0 + for line in lines: + assert html_check in line and open_job_check in line + + +def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): bf.options.display.progress_bar = "notebook" penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = ( False ) - penguins_df_default_index.compute() - query_job_repr = formatting_helpers.repr_query_job( + penguins_df_default_index.to_pandas() + query_job_repr = formatting_helpers.repr_query_job_html( penguins_df_default_index.query_job ).value string_checks = [ @@ -72,3 +123,33 @@ def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame): ] for string in string_checks: assert string in query_job_repr + + +def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame): + penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = ( + False + ) + penguins_df_default_index.to_pandas() + query_job_repr = formatting_helpers.repr_query_job( + penguins_df_default_index.query_job + ) + string_checks = [ + "Job", + "Destination Table", + "Slot Time", + "Bytes Processed", + "Cache hit", + ] + for string in string_checks: + assert string in query_job_repr + + +def test_query_job_dry_run( + penguins_df_default_index: bf.dataframe.DataFrame, capsys, deferred_repr +): + repr(penguins_df_default_index) + repr(penguins_df_default_index["body_mass_g"]) + lines = capsys.readouterr().out.split("\n") + lines = filter(None, lines) + for line in lines: + assert "Computation deferred. Computation will process" in line diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 4c75b5d3d27..fe4b1c5a97a 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -16,7 +16,7 @@ import pytest import bigframes -from bigframes.remote_function import remote_function +from bigframes.remote_function import read_gbq_function, remote_function from tests.system.utils import assert_pandas_df_equal_ignore_ordering @@ -28,6 +28,38 @@ def bq_cf_connection() -> str: return "bigframes-rf-conn" +@pytest.fixture(scope="module") +def bq_cf_connection_location() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us --project_id=bigframes-dev bigframes-rf-conn + """ + return "us.bigframes-rf-conn" + + +@pytest.fixture(scope="module") +def bq_cf_connection_location_mistached() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us-east1 --project_id=bigframes-dev bigframes-rf-conn + """ + return "us-east1.bigframes-rf-conn" + + +@pytest.fixture(scope="module") +def bq_cf_connection_location_project() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us --project_id=bigframes-dev bigframes-rf-conn + """ + return "bigframes-dev.us.bigframes-rf-conn" + + +@pytest.fixture(scope="module") +def bq_cf_connection_location_project_mistached() -> str: + """Pre-created BQ connection to invoke cloud function for bigframes-dev + $ bq show --connection --location=us-east1 --project_id=bigframes-metrics bigframes-rf-conn + """ + return "bigframes-metrics.us-east1.bigframes-rf-conn" + + @pytest.fixture(scope="module") def session_with_bq_connection(bq_cf_connection) -> bigframes.Session: return bigframes.Session( @@ -35,10 +67,42 @@ def session_with_bq_connection(bq_cf_connection) -> bigframes.Session: ) +@pytest.fixture(scope="module") +def session_with_bq_connection_location_specified( + bq_cf_connection_location, +) -> bigframes.Session: + return bigframes.Session( + bigframes.BigQueryOptions(remote_udf_connection=bq_cf_connection_location) + ) + + +@pytest.fixture(scope="module") +def session_with_bq_connection_location_mistached( + bq_cf_connection_location_mistached, +) -> bigframes.Session: + return bigframes.Session( + bigframes.BigQueryOptions( + remote_udf_connection=bq_cf_connection_location_mistached + ) + ) + + +@pytest.fixture(scope="module") +def session_with_bq_connection_location_project_specified( + bq_cf_connection_location_project, +) -> bigframes.Session: + return bigframes.Session( + bigframes.BigQueryOptions( + remote_udf_connection=bq_cf_connection_location_project + ) + ) + + @pytest.mark.flaky(retries=2, delay=120) def test_remote_function_direct_no_session_param( bigquery_client, bigqueryconnection_client, + cloudfunctions_client, scalars_dfs, dataset_id_permanent, bq_cf_connection, @@ -48,6 +112,7 @@ def test_remote_function_direct_no_session_param( int, bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection, # See e2e tests for tests that actually deploy the Cloud Function. @@ -56,13 +121,112 @@ def test_remote_function_direct_no_session_param( def square(x): return x * x + assert square.bigframes_remote_function + assert square.bigframes_cloud_function + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_no_session_param_location_specified( + bigquery_client, + bigqueryconnection_client, + cloudfunctions_client, + scalars_dfs, + dataset_id_permanent, + bq_cf_connection_location, +): + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection_location, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_no_session_param_location_mistached( + bigquery_client, + bigqueryconnection_client, + cloudfunctions_client, + scalars_dfs, + dataset_id_permanent, + bq_cf_connection_location_mistached, +): + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection_location_mistached, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + scalars_df, scalars_pandas_df = scalars_dfs bf_int64_col = scalars_df["int64_col"] bf_int64_col_filter = bf_int64_col.notnull() bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) pd_int64_col = scalars_pandas_df["int64_col"] pd_int64_col_filter = pd_int64_col.notnull() @@ -78,6 +242,78 @@ def square(x): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_no_session_param_location_project_specified( + bigquery_client, + bigqueryconnection_client, + cloudfunctions_client, + scalars_dfs, + dataset_id_permanent, + bq_cf_connection_location_project, +): + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection_location_project, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_direct_no_session_param_project_mismatched( + bigquery_client, + bigqueryconnection_client, + cloudfunctions_client, + dataset_id_permanent, + bq_cf_connection_location_project_mistached, +): + with pytest.raises(ValueError): + + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection_location_project_mistached, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x + + @pytest.mark.flaky(retries=2, delay=120) def test_remote_function_direct_session_param(session_with_bq_connection, scalars_dfs): @remote_function( @@ -94,7 +330,9 @@ def square(x): bf_int64_col_filter = bf_int64_col.notnull() bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) pd_int64_col = scalars_pandas_df["int64_col"] pd_int64_col_filter = pd_int64_col.notnull() @@ -129,7 +367,9 @@ def square(x): bf_int64_col_filter = bf_int64_col.notnull() bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) pd_int64_col = scalars_pandas_df["int64_col"] pd_int64_col_filter = pd_int64_col.notnull() @@ -166,7 +406,9 @@ def square(x): bf_int64_col_filter = bf_int64_col.notnull() bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) pd_int64_col = scalars_pandas_df["int64_col"] pd_int64_col_filter = pd_int64_col.notnull() @@ -210,7 +452,9 @@ def square(x): bf_int64_col_filter = bf_int64_col.notnull() bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = bf_int64_col_filtered.to_frame().assign(result=bf_result_col).compute() + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) pd_int64_col = scalars_pandas_df["int64_col"] pd_int64_col_filter = pd_int64_col.notnull() @@ -238,7 +482,7 @@ def add_one(x): bf_int64_df = scalars_df[int64_cols] bf_int64_df_filtered = bf_int64_df.dropna() - bf_result = bf_int64_df_filtered.applymap(remote_add_one).compute() + bf_result = bf_int64_df_filtered.applymap(remote_add_one).to_pandas() pd_int64_df = scalars_pandas_df[int64_cols] pd_int64_df_filtered = pd_int64_df.dropna() @@ -264,7 +508,7 @@ def add_one(x): int64_cols = ["int64_col", "int64_too"] bf_int64_df = scalars_df[int64_cols] - bf_result = bf_int64_df.applymap(remote_add_one, na_action="ignore").compute() + bf_result = bf_int64_df.applymap(remote_add_one, na_action="ignore").to_pandas() pd_int64_df = scalars_pandas_df[int64_cols] pd_result = pd_int64_df.applymap(add_one, na_action="ignore") @@ -276,3 +520,56 @@ def add_one(x): pd_result[col] = pd_result[col].astype(pd_int64_df[col].dtype) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_read_gbq_function_like_original( + bigquery_client, + bigqueryconnection_client, + cloudfunctions_client, + scalars_df_index, + dataset_id_permanent, + bq_cf_connection, +): + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + dataset=dataset_id_permanent, + cloud_functions_client=cloudfunctions_client, + bigquery_connection=bq_cf_connection, + reuse=True, + ) + def square1(x): + return x * x + + square2 = read_gbq_function( + function_name=square1.bigframes_remote_function, + bigquery_client=bigquery_client, + ) + + # The newly-created function (square1) should have a remote function AND a + # cloud function associated with it, while the read-back version (square2) + # should only have a remote function. + assert square1.bigframes_remote_function + assert square1.bigframes_cloud_function + + assert square2.bigframes_remote_function + assert not hasattr(square2, "bigframes_cloud_function") + + # They should point to the same function. + assert square1.bigframes_remote_function == square2.bigframes_remote_function + + # The result of applying them should be the same. + int64_col = scalars_df_index["int64_col"] + int64_col_filter = int64_col.notnull() + int64_col_filtered = int64_col[int64_col_filter] + + s1_result_col = int64_col_filtered.apply(square1) + s1_result = int64_col_filtered.to_frame().assign(result=s1_result_col) + + s2_result_col = int64_col_filtered.apply(square2) + s2_result = int64_col_filtered.to_frame().assign(result=s2_result_col) + + assert_pandas_df_equal_ignore_ordering(s1_result.to_pandas(), s2_result.to_pandas()) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 751521bd754..1c1e7b035bd 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -13,6 +13,7 @@ # limitations under the License. import math +import re import tempfile import geopandas as gpd # type: ignore @@ -21,6 +22,7 @@ import pyarrow as pa # type: ignore import pytest +import bigframes.pandas import bigframes.series as series from tests.system.utils import ( assert_pandas_df_equal_ignore_ordering, @@ -32,7 +34,7 @@ def test_series_construct_copy(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = series.Series( scalars_df["int64_col"], name="test_series", dtype="Float64" - ).compute() + ).to_pandas() pd_result = pd.Series( scalars_pandas_df["int64_col"], name="test_series", dtype="Float64" ) @@ -48,11 +50,11 @@ def test_series_construct_pandas(scalars_dfs): scalars_pandas_df["int64_col"], name="test_series", dtype="Float64" ) assert bf_result.shape == pd_result.shape - pd.testing.assert_series_equal(bf_result.compute(), pd_result) + pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result) def test_series_construct_from_list(): - bf_result = series.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64").compute() + bf_result = series.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64").to_pandas() pd_result = pd.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64") # BigQuery DataFrame default indices use nullable Int64 always @@ -61,6 +63,22 @@ def test_series_construct_from_list(): pd.testing.assert_series_equal(bf_result, pd_result) +def test_series_construct_from_list_escaped_strings(): + """Check that special characters are supported.""" + strings = [ + "string\nwith\nnewline", + "string\twith\ttabs", + "string\\with\\backslashes", + ] + bf_result = series.Series(strings, name="test_series", dtype="string[pyarrow]") + pd_result = pd.Series(strings, name="test_series", dtype="string[pyarrow]") + + # BigQuery DataFrame default indices use nullable Int64 always + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result) + + @pytest.mark.parametrize( ["col_name", "expected_dtype"], [ @@ -83,7 +101,7 @@ def test_series_construct_from_list(): def test_get_column(scalars_dfs, col_name, expected_dtype): scalars_df, scalars_pandas_df = scalars_dfs series = scalars_df[col_name] - series_pandas = series.compute() + series_pandas = series.to_pandas() assert series_pandas.dtype == expected_dtype assert series_pandas.shape[0] == scalars_pandas_df.shape[0] @@ -94,6 +112,54 @@ def test_series_get_column_default(scalars_dfs): assert result == "default_val" +def test_series_get_with_default_index(scalars_dfs): + col_name = "float64_col" + key = 2 + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].get(key) + pd_result = scalars_pandas_df[col_name].get(key) + assert bf_result.to_pandas().iloc[0] == pd_result + + +@pytest.mark.parametrize( + ("index_col", "key"), + ( + ("int64_too", 2), + ("string_col", "Hello, World!"), + ("int64_too", slice(2, 6)), + ), +) +def test_series___getitem__(scalars_dfs, index_col, key): + col_name = "float64_col" + scalars_df, scalars_pandas_df = scalars_dfs + scalars_df = scalars_df.set_index(index_col, drop=False) + scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False) + bf_result = scalars_df[col_name][key] + pd_result = scalars_pandas_df[col_name][key] + pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result) + + +def test_series___getitem___with_int_key(scalars_dfs): + col_name = "int64_too" + index_col = "string_col" + key = 2 + scalars_df, scalars_pandas_df = scalars_dfs + scalars_df = scalars_df.set_index(index_col, drop=False) + scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False) + bf_result = scalars_df[col_name][key] + pd_result = scalars_pandas_df[col_name][key] + assert bf_result == pd_result + + +def test_series___getitem___with_default_index(scalars_dfs): + col_name = "float64_col" + key = 2 + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name][key] + pd_result = scalars_pandas_df[col_name][key] + assert bf_result.to_pandas().iloc[0] == pd_result + + @pytest.mark.parametrize( ("col_name",), ( @@ -103,7 +169,7 @@ def test_series_get_column_default(scalars_dfs): ) def test_abs(scalars_dfs, col_name): scalars_df, scalars_pandas_df = scalars_dfs - bf_result = scalars_df[col_name].abs().compute() + bf_result = scalars_df[col_name].abs().to_pandas() pd_result = scalars_pandas_df[col_name].abs() assert_series_equal_ignoring_order(pd_result, bf_result) @@ -112,7 +178,7 @@ def test_abs(scalars_dfs, col_name): def test_fillna(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" - bf_result = scalars_df[col_name].fillna("Missing").compute() + bf_result = scalars_df[col_name].fillna("Missing").to_pandas() pd_result = scalars_pandas_df[col_name].fillna("Missing") assert_series_equal_ignoring_order( pd_result, @@ -120,6 +186,25 @@ def test_fillna(scalars_dfs): ) +def test_series_agg_single_string(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_col"].agg("sum") + pd_result = scalars_pandas_df["int64_col"].agg("sum") + assert math.isclose(pd_result, bf_result) + + +def test_series_agg_multi_string(scalars_dfs): + aggregations = ["sum", "mean", "std", "var", "min", "max", "nunique", "count"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_col"].agg(aggregations).to_pandas() + pd_result = scalars_pandas_df["int64_col"].agg(aggregations) + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + @pytest.mark.parametrize( ("col_name",), ( @@ -176,6 +261,36 @@ def test_kurt(scalars_dfs, col_name): assert math.isclose(pd_result, bf_result) +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_col",), + ), +) +def test_skew(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[col_name].skew() + pd_result = scalars_pandas_df[col_name].skew() + assert math.isclose(pd_result, bf_result) + + +def test_skew_undefined(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_col"].iloc[:2].skew() + pd_result = scalars_pandas_df["int64_col"].iloc[:2].skew() + # both should be pd.NA + assert pd_result is bf_result + + +def test_kurt_undefined(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_col"].iloc[:3].kurt() + pd_result = scalars_pandas_df["int64_col"].iloc[:3].kurt() + # both should be pd.NA + assert pd_result is bf_result + + @pytest.mark.parametrize( ("col_name",), ( @@ -198,7 +313,7 @@ def test_var(scalars_dfs, col_name): ), ) def test_mode_stat(scalars_df_index, scalars_pandas_df_index, col_name): - bf_result = scalars_df_index[col_name].mode().compute() + bf_result = scalars_df_index[col_name].mode().to_pandas() pd_result = scalars_pandas_df_index[col_name].mode() ## Mode implicitly resets index, and bigframes default indices use nullable Int64 @@ -244,7 +359,7 @@ def test_series_int_int_operators_scalar( maybe_reversed_op = (lambda x, y: operator(y, x)) if reverse_operands else operator - bf_result = maybe_reversed_op(scalars_df["int64_col"], other_scalar).compute() + bf_result = maybe_reversed_op(scalars_df["int64_col"], other_scalar).to_pandas() pd_result = maybe_reversed_op(scalars_pandas_df["int64_col"], other_scalar) assert_series_equal_ignoring_order(pd_result, bf_result) @@ -270,7 +385,7 @@ def test_series_bool_bool_operators_scalar( maybe_reversed_op = (lambda x, y: operator(y, x)) if reverse_operands else operator - bf_result = maybe_reversed_op(scalars_df["bool_col"], other_scalar).compute() + bf_result = maybe_reversed_op(scalars_df["bool_col"], other_scalar).to_pandas() pd_result = maybe_reversed_op(scalars_pandas_df["bool_col"], other_scalar) assert_series_equal_ignoring_order(pd_result.astype(pd.BooleanDtype()), bf_result) @@ -309,12 +424,101 @@ def test_series_bool_bool_operators_scalar( ) def test_series_int_int_operators_series(scalars_dfs, operator): scalars_df, scalars_pandas_df = scalars_dfs - bf_result = operator(scalars_df["int64_col"], scalars_df["int64_too"]).compute() + bf_result = operator(scalars_df["int64_col"], scalars_df["int64_too"]).to_pandas() pd_result = operator(scalars_pandas_df["int64_col"], scalars_pandas_df["int64_too"]) - assert_series_equal_ignoring_order(pd_result, bf_result) +@pytest.mark.parametrize( + ("col_x",), + [ + ("int64_col",), + ("int64_too",), + ("float64_col",), + ], +) +@pytest.mark.parametrize( + ("col_y",), + [ + ("int64_col",), + ("int64_too",), + ("float64_col",), + ], +) +@pytest.mark.parametrize( + ("method",), + [ + ("mod",), + ("rmod",), + ], +) +def test_mods(scalars_dfs, col_x, col_y, method): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = getattr(scalars_df[col_x], method)(scalars_df[col_y]).to_pandas() + pd_result = getattr(scalars_pandas_df[col_x], method)(scalars_pandas_df[col_y]) + pd.testing.assert_series_equal(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_x",), + [ + ("int64_col",), + ("float64_col",), + ], +) +@pytest.mark.parametrize( + ("col_y",), + [ + ("int64_col",), + ("float64_col",), + ], +) +@pytest.mark.parametrize( + ("method",), + [ + ("divmod",), + ("rdivmod",), + ], +) +def test_divmods_series(scalars_dfs, col_x, col_y, method): + scalars_df, scalars_pandas_df = scalars_dfs + bf_div_result, bf_mod_result = getattr(scalars_df[col_x], method)(scalars_df[col_y]) + pd_div_result, pd_mod_result = getattr(scalars_pandas_df[col_x], method)( + scalars_pandas_df[col_y] + ) + pd.testing.assert_series_equal(pd_div_result, bf_div_result.to_pandas()) + pd.testing.assert_series_equal(pd_mod_result, bf_mod_result.to_pandas()) + + +@pytest.mark.parametrize( + ("col_x",), + [ + ("int64_col",), + ("float64_col",), + ], +) +@pytest.mark.parametrize( + ("other",), + [ + (-1000,), + (678,), + ], +) +@pytest.mark.parametrize( + ("method",), + [ + ("divmod",), + ("rdivmod",), + ], +) +def test_divmods_scalars(scalars_dfs, col_x, other, method): + scalars_df, scalars_pandas_df = scalars_dfs + bf_div_result, bf_mod_result = getattr(scalars_df[col_x], method)(other) + pd_div_result, pd_mod_result = getattr(scalars_pandas_df[col_x], method)(other) + pd.testing.assert_series_equal(pd_div_result, bf_div_result.to_pandas()) + pd.testing.assert_series_equal(pd_mod_result, bf_mod_result.to_pandas()) + + @pytest.mark.parametrize( ("other",), [ @@ -324,7 +528,7 @@ def test_series_int_int_operators_series(scalars_dfs, operator): ) def test_series_add_scalar(scalars_dfs, other): scalars_df, scalars_pandas_df = scalars_dfs - bf_result = (scalars_df["float64_col"] + other).compute() + bf_result = (scalars_df["float64_col"] + other).to_pandas() pd_result = scalars_pandas_df["float64_col"] + other assert_series_equal_ignoring_order(pd_result, bf_result) @@ -340,7 +544,7 @@ def test_series_add_scalar(scalars_dfs, other): ) def test_series_add_bigframes_series(scalars_dfs, left_col, right_col): scalars_df, scalars_pandas_df = scalars_dfs - bf_result = (scalars_df[left_col] + scalars_df[right_col]).compute() + bf_result = (scalars_df[left_col] + scalars_df[right_col]).to_pandas() pd_result = scalars_pandas_df[left_col] + scalars_pandas_df[right_col] assert_series_equal_ignoring_order(pd_result, bf_result) @@ -360,7 +564,7 @@ def test_series_add_bigframes_series_nested( scalars_df, scalars_pandas_df = scalars_dfs bf_result = ( (scalars_df[left_col] + scalars_df[right_col]) + scalars_df[righter_col] - ).compute() + ).to_pandas() pd_result = ( scalars_pandas_df[left_col] + scalars_pandas_df[right_col] ) + scalars_pandas_df[righter_col] @@ -375,12 +579,12 @@ def test_series_add_different_table_default_index( bf_result = ( scalars_df_default_index["float64_col"] + scalars_df_2_default_index["float64_col"] - ).compute() + ).to_pandas() pd_result = ( # Default index may not have a well defined order, but it should at - # least be consistent across compute() calls. - scalars_df_default_index["float64_col"].compute() - + scalars_df_2_default_index["float64_col"].compute() + # least be consistent across to_pandas() calls. + scalars_df_default_index["float64_col"].to_pandas() + + scalars_df_2_default_index["float64_col"].to_pandas() ) # TODO(swast): Can remove sort_index() when there's default ordering. pd.testing.assert_series_equal(bf_result.sort_index(), pd_result.sort_index()) @@ -394,7 +598,7 @@ def test_series_add_different_table_with_index( # When index values are unique, we can emulate with values from the same # DataFrame. pd_result = scalars_pandas_df["float64_col"] + scalars_pandas_df["int64_col"] - pd.testing.assert_series_equal(bf_result.compute(), pd_result) + pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result) def test_reset_index_drop(scalars_df_index, scalars_pandas_df_index): @@ -413,7 +617,7 @@ def test_reset_index_drop(scalars_df_index, scalars_pandas_df_index): # BigQuery DataFrames default indices use nullable Int64 always pd_result.index = pd_result.index.astype("Int64") - pd.testing.assert_series_equal(bf_result.compute(), pd_result) + pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result) @pytest.mark.parametrize( @@ -440,7 +644,7 @@ def test_reset_index_no_drop(scalars_df_index, scalars_pandas_df_index, name): # BigQuery DataFrames default indices use nullable Int64 always pd_result.index = pd_result.index.astype("Int64") - pd.testing.assert_frame_equal(bf_result.compute(), pd_result) + pd.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) def test_series_add_pandas_series_not_implemented(scalars_dfs): @@ -451,7 +655,7 @@ def test_series_add_pandas_series_not_implemented(scalars_dfs): + pd.Series( [1, 1, 1, 1], ) - ).compute() + ).to_pandas() def test_copy(scalars_df_index, scalars_pandas_df_index): @@ -467,14 +671,14 @@ def test_copy(scalars_df_index, scalars_pandas_df_index): pd_copy.loc[0] = 5.6 pd_series.loc[0] = 3.4 - assert bf_copy.compute().loc[0] != bf_series.compute().loc[0] - pd.testing.assert_series_equal(bf_copy.compute(), pd_copy) + assert bf_copy.to_pandas().loc[0] != bf_series.to_pandas().loc[0] + pd.testing.assert_series_equal(bf_copy.to_pandas(), pd_copy) def test_isnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "float64_col" - bf_series = scalars_df[col_name].isnull().compute() + bf_series = scalars_df[col_name].isnull().to_pandas() pd_series = scalars_pandas_df[col_name].isnull() # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but @@ -485,7 +689,7 @@ def test_isnull(scalars_dfs): def test_notnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" - bf_series = scalars_df[col_name].notnull().compute() + bf_series = scalars_df[col_name].notnull().to_pandas() pd_series = scalars_pandas_df[col_name].notnull() # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but @@ -496,7 +700,7 @@ def test_notnull(scalars_dfs): def test_round(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "float64_col" - bf_result = scalars_df[col_name].round().compute() + bf_result = scalars_df[col_name].round().to_pandas() pd_result = scalars_pandas_df[col_name].round() assert_series_equal_ignoring_order(pd_result, bf_result) @@ -505,7 +709,7 @@ def test_round(scalars_dfs): def test_eq_scalar(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_result = scalars_df[col_name].eq(0).compute() + bf_result = scalars_df[col_name].eq(0).to_pandas() pd_result = scalars_pandas_df[col_name].eq(0) assert_series_equal_ignoring_order(pd_result, bf_result) @@ -514,7 +718,7 @@ def test_eq_scalar(scalars_dfs): def test_eq_wider_type_scalar(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_result = scalars_df[col_name].eq(1.0).compute() + bf_result = scalars_df[col_name].eq(1.0).to_pandas() pd_result = scalars_pandas_df[col_name].eq(1.0) assert_series_equal_ignoring_order(pd_result, bf_result) @@ -523,7 +727,7 @@ def test_eq_wider_type_scalar(scalars_dfs): def test_ne_scalar(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_result = (scalars_df[col_name] != 0).compute() + bf_result = (scalars_df[col_name] != 0).to_pandas() pd_result = scalars_pandas_df[col_name] != 0 assert_series_equal_ignoring_order(pd_result, bf_result) @@ -532,7 +736,7 @@ def test_ne_scalar(scalars_dfs): def test_eq_int_scalar(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_result = (scalars_df[col_name] == 0).compute() + bf_result = (scalars_df[col_name] == 0).to_pandas() pd_result = scalars_pandas_df[col_name] == 0 assert_series_equal_ignoring_order(pd_result, bf_result) @@ -549,7 +753,7 @@ def test_eq_int_scalar(scalars_dfs): def test_eq_same_type_series(scalars_dfs, col_name): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" - bf_result = (scalars_df[col_name] == scalars_df[col_name]).compute() + bf_result = (scalars_df[col_name] == scalars_df[col_name]).to_pandas() pd_result = scalars_pandas_df[col_name] == scalars_pandas_df[col_name] # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but @@ -564,18 +768,18 @@ def test_loc_setitem_cell(scalars_df_index, scalars_pandas_df_index): pd_series = scalars_pandas_df_index["string_col"].copy() bf_series.loc[2] = "This value isn't in the test data." pd_series.loc[2] = "This value isn't in the test data." - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() pd_result = pd_series pd.testing.assert_series_equal(bf_result, pd_result) # Per Copy-on-Write semantics, other references to the original DataFrame # should remain unchanged. - pd.testing.assert_series_equal(bf_original.compute(), pd_original) + pd.testing.assert_series_equal(bf_original.to_pandas(), pd_original) def test_ne_obj_series(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" - bf_result = (scalars_df[col_name] != scalars_df[col_name]).compute() + bf_result = (scalars_df[col_name] != scalars_df[col_name]).to_pandas() pd_result = scalars_pandas_df[col_name] != scalars_pandas_df[col_name] # One of dtype mismatches to be documented. Here, the `bf_series.dtype` is `BooleanDtype` but @@ -586,7 +790,7 @@ def test_ne_obj_series(scalars_dfs): def test_indexing_using_unselected_series(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" - bf_result = scalars_df[col_name][scalars_df["int64_too"].eq(0)].compute() + bf_result = scalars_df[col_name][scalars_df["int64_too"].eq(0)].to_pandas() pd_result = scalars_pandas_df[col_name][scalars_pandas_df["int64_too"].eq(0)] assert_series_equal_ignoring_order( @@ -600,7 +804,7 @@ def test_indexing_using_selected_series(scalars_dfs): col_name = "string_col" bf_result = scalars_df[col_name][ scalars_df["string_col"].eq("Hello, World!") - ].compute() + ].to_pandas() pd_result = scalars_pandas_df[col_name][ scalars_pandas_df["string_col"].eq("Hello, World!") ] @@ -618,7 +822,7 @@ def test_nested_filter(scalars_dfs): bool_col = scalars_df["bool_col"] == bool( True ) # Convert from nullable bool to nonnullable bool usable as indexer - bf_result = string_col[int64_too == 0][~bool_col].compute() + bf_result = string_col[int64_too == 0][~bool_col].to_pandas() pd_string_col = scalars_pandas_df["string_col"] pd_int64_too = scalars_pandas_df["int64_too"] @@ -633,12 +837,35 @@ def test_nested_filter(scalars_dfs): ) +def test_binop_repeated_application_does_row_identity_joins(scalars_dfs): + """Make sure row identity joins kick in so that we don't do way more joins than expected.""" + scalars_df, scalars_pandas_df = scalars_dfs + bf_series = scalars_df["int64_col"] + pd_series = scalars_pandas_df["int64_col"] + + num_joins = 10 + for _ in range(num_joins): + bf_series = bf_series + bf_series + pd_series = pd_series + pd_series + + bf_result = bf_series.to_pandas() + pd_result = pd_series + assert_series_equal_ignoring_order( + bf_result, + pd_result, + ) + + bf_sql, _ = bf_series.to_frame()._to_sql_query(always_include_index=True) + selects = re.findall("SELECT", bf_sql.upper()) + assert 0 < len(selects) < (num_joins // 2) + + def test_binop_opposite_filters(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs int64_col1 = scalars_df["int64_col"] int64_col2 = scalars_df["int64_col"] bool_col = scalars_df["bool_col"] - bf_result = (int64_col1[bool_col] + int64_col2[bool_col.__invert__()]).compute() + bf_result = (int64_col1[bool_col] + int64_col2[bool_col.__invert__()]).to_pandas() pd_int64_col1 = scalars_pandas_df["int64_col"] pd_int64_col2 = scalars_pandas_df["int64_col"] @@ -656,7 +883,7 @@ def test_binop_left_filtered(scalars_dfs): int64_col = scalars_df["int64_col"] float64_col = scalars_df["float64_col"] bool_col = scalars_df["bool_col"] - bf_result = (int64_col[bool_col] + float64_col).compute() + bf_result = (int64_col[bool_col] + float64_col).to_pandas() pd_int64_col = scalars_pandas_df["int64_col"] pd_float64_col = scalars_pandas_df["float64_col"] @@ -674,7 +901,7 @@ def test_binop_right_filtered(scalars_dfs): int64_col = scalars_df["int64_col"] float64_col = scalars_df["float64_col"] bool_col = scalars_df["bool_col"] - bf_result = (float64_col + int64_col[bool_col]).compute() + bf_result = (float64_col + int64_col[bool_col]).to_pandas() pd_int64_col = scalars_pandas_df["int64_col"] pd_float64_col = scalars_pandas_df["float64_col"] @@ -695,6 +922,16 @@ def test_mean(scalars_dfs): assert math.isclose(pd_result, bf_result) +def test_median(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + bf_result = scalars_df[col_name].median() + pd_max = scalars_pandas_df[col_name].max() + pd_min = scalars_pandas_df[col_name].min() + # Median is approximate, so just check for plausibility. + assert pd_min < bf_result < pd_max + + def test_repr(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs if scalars_pandas_df.index.name != "rowindex": @@ -722,6 +959,19 @@ def test_product(scalars_dfs): assert math.isclose(pd_result, bf_result) +def test_cumprod(scalars_dfs): + if pd.__version__.startswith("1."): + pytest.skip("Series.cumprod NA mask are different in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "float64_col" + bf_result = scalars_df[col_name].cumprod() + pd_result = scalars_pandas_df[col_name].cumprod() + pd.testing.assert_series_equal( + pd_result, + bf_result.to_pandas(), + ) + + def test_count(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_col" @@ -762,7 +1012,7 @@ def test_groupby_sum(scalars_dfs): scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).sum() ) # TODO(swast): Update groupby to use index based on group by key(s). - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() assert_series_equal_ignoring_order( pd_series, bf_result, @@ -780,7 +1030,7 @@ def test_groupby_std(scalars_dfs): .std() .astype(pd.Float64Dtype()) ) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() assert_series_equal_ignoring_order( pd_series, bf_result, @@ -795,7 +1045,7 @@ def test_groupby_var(scalars_dfs): pd_series = ( scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).var() ) - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() assert_series_equal_ignoring_order( pd_series, bf_result, @@ -815,7 +1065,7 @@ def test_groupby_level_sum(scalars_dfs): # TODO(swast): Update groupby to use index based on group by key(s). pd.testing.assert_series_equal( pd_series.sort_index(), - bf_series.compute().sort_index(), + bf_series.to_pandas().sort_index(), ) @@ -831,7 +1081,7 @@ def test_groupby_level_list_sum(scalars_dfs): # TODO(swast): Update groupby to use index based on group by key(s). pd.testing.assert_series_equal( pd_series.sort_index(), - bf_series.compute().sort_index(), + bf_series.to_pandas().sort_index(), ) @@ -847,13 +1097,36 @@ def test_groupby_mean(scalars_dfs): .mean() ) # TODO(swast): Update groupby to use index based on group by key(s). - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() assert_series_equal_ignoring_order( pd_series, bf_result, ) +def test_groupby_median(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_too" + bf_series = ( + scalars_df[col_name].groupby(scalars_df["string_col"], dropna=False).median() + ) + pd_max = ( + scalars_pandas_df[col_name] + .groupby(scalars_pandas_df["string_col"], dropna=False) + .max() + ) + pd_min = ( + scalars_pandas_df[col_name] + .groupby(scalars_pandas_df["string_col"], dropna=False) + .min() + ) + # TODO(swast): Update groupby to use index based on group by key(s). + bf_result = bf_series.to_pandas() + + # Median is approximate, so just check that it's plausible. + assert ((pd_min <= bf_result) & (bf_result <= pd_max)).all() + + def test_groupby_prod(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" @@ -862,7 +1135,7 @@ def test_groupby_prod(scalars_dfs): scalars_pandas_df[col_name].groupby(scalars_pandas_df["int64_col"]).prod() ) # TODO(swast): Update groupby to use index based on group by key(s). - bf_result = bf_series.compute() + bf_result = bf_series.to_pandas() assert_series_equal_ignoring_order( pd_series, bf_result, @@ -897,7 +1170,7 @@ def test_groupby_window_ops(scalars_df_index, scalars_pandas_df_index, operator) group_key = "int64_too" # has some duplicates values, good for grouping bf_series = ( operator(scalars_df_index[col_name].groupby(scalars_df_index[group_key])) - ).compute() + ).to_pandas() pd_series = operator( scalars_pandas_df_index[col_name].groupby(scalars_pandas_df_index[group_key]) ).astype(pd.Int64Dtype()) @@ -909,7 +1182,7 @@ def test_groupby_window_ops(scalars_df_index, scalars_pandas_df_index, operator) def test_drop_label(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" - bf_series = scalars_df_index[col_name].drop(1).compute() + bf_series = scalars_df_index[col_name].drop(1).to_pandas() pd_series = scalars_pandas_df_index[col_name].drop(1) pd.testing.assert_series_equal( pd_series, @@ -919,7 +1192,7 @@ def test_drop_label(scalars_df_index, scalars_pandas_df_index): def test_drop_label_list(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" - bf_series = scalars_df_index[col_name].drop([1, 3]).compute() + bf_series = scalars_df_index[col_name].drop([1, 3]).to_pandas() pd_series = scalars_pandas_df_index[col_name].drop([1, 3]) pd.testing.assert_series_equal( pd_series, @@ -943,7 +1216,7 @@ def test_drop_label_list(scalars_df_index, scalars_pandas_df_index): ], ) def test_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, col_name): - bf_series = scalars_df_index[col_name].drop_duplicates(keep=keep).compute() + bf_series = scalars_df_index[col_name].drop_duplicates(keep=keep).to_pandas() pd_series = scalars_pandas_df_index[col_name].drop_duplicates(keep=keep) pd.testing.assert_series_equal( pd_series, @@ -980,7 +1253,7 @@ def test_unique(scalars_df_index, scalars_pandas_df_index, col_name): ], ) def test_duplicated(scalars_df_index, scalars_pandas_df_index, keep, col_name): - bf_series = scalars_df_index[col_name].duplicated(keep=keep).compute() + bf_series = scalars_df_index[col_name].duplicated(keep=keep).to_pandas() pd_series = scalars_pandas_df_index[col_name].duplicated(keep=keep) pd.testing.assert_series_equal(pd_series, bf_series, check_dtype=False) @@ -1021,7 +1294,7 @@ def test_empty_false(scalars_dfs): assert pd_result == bf_result -def test_empty_true(scalars_dfs): +def test_empty_true_row_filter(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = scalars_df["string_col"][ @@ -1031,9 +1304,21 @@ def test_empty_true(scalars_dfs): scalars_pandas_df["string_col"] == "won't find this" ].empty + assert pd_result assert pd_result == bf_result +def test_empty_true_memtable(session: bigframes.Session): + bf_series: series.Series = series.Series(session=session) + pd_series: pd.Series = pd.Series() + + bf_result = bf_series.empty + pd_result = pd_series.empty + + assert pd_result + assert bf_result == pd_result + + def test_dtype(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs @@ -1058,7 +1343,7 @@ def test_head(scalars_dfs): if scalars_df.index.name is None: pytest.skip("Require explicit index for offset ops.") - bf_result = scalars_df["string_col"].head(2).compute() + bf_result = scalars_df["string_col"].head(2).to_pandas() pd_result = scalars_pandas_df["string_col"].head(2) assert_series_equal_ignoring_order( @@ -1073,7 +1358,7 @@ def test_tail(scalars_dfs): if scalars_df.index.name is None: pytest.skip("Require explicit index for offset ops.") - bf_result = scalars_df["string_col"].tail(2).compute() + bf_result = scalars_df["string_col"].tail(2).to_pandas() pd_result = scalars_pandas_df["string_col"].tail(2) assert_series_equal_ignoring_order( @@ -1088,7 +1373,7 @@ def test_head_then_scalar_operation(scalars_dfs): if scalars_df.index.name is None: pytest.skip("Require explicit index for offset ops.") - bf_result = (scalars_df["float64_col"].head(1) + 4).compute() + bf_result = (scalars_df["float64_col"].head(1) + 4).to_pandas() pd_result = scalars_pandas_df["float64_col"].head(1) + 4 pd.testing.assert_series_equal( @@ -1105,7 +1390,7 @@ def test_head_then_series_operation(scalars_dfs): bf_result = ( scalars_df["float64_col"].head(4) + scalars_df["float64_col"].head(2) - ).compute() + ).to_pandas() pd_result = scalars_pandas_df["float64_col"].head(4) + scalars_pandas_df[ "float64_col" ].head(2) @@ -1118,7 +1403,7 @@ def test_head_then_series_operation(scalars_dfs): def test_shift(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" - bf_result = scalars_df_index[col_name].shift().compute() + bf_result = scalars_df_index[col_name].shift().to_pandas() # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA pd_result = scalars_pandas_df_index[col_name].shift().astype(pd.Int64Dtype()) @@ -1133,7 +1418,7 @@ def test_cumsum_int(scalars_df_index, scalars_pandas_df_index): pytest.skip("Series.cumsum NA mask are different in pandas 1.x.") col_name = "int64_col" - bf_result = scalars_df_index[col_name].cumsum().compute() + bf_result = scalars_df_index[col_name].cumsum().to_pandas() # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA pd_result = scalars_pandas_df_index[col_name].cumsum().astype(pd.Int64Dtype()) @@ -1143,6 +1428,27 @@ def test_cumsum_int(scalars_df_index, scalars_pandas_df_index): ) +def test_cumsum_int_ordered(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("1."): + pytest.skip("Series.cumsum NA mask are different in pandas 1.x.") + + col_name = "int64_col" + bf_result = ( + scalars_df_index.sort_values(by="rowindex_2")[col_name].cumsum().to_pandas() + ) + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = ( + scalars_pandas_df_index.sort_values(by="rowindex_2")[col_name] + .cumsum() + .astype(pd.Int64Dtype()) + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + @pytest.mark.parametrize( ("na_option",), [ @@ -1167,7 +1473,7 @@ def test_cumsum_int(scalars_df_index, scalars_pandas_df_index): def test_rank_with_nulls(scalars_df_index, scalars_pandas_df_index, na_option, method): col_name = "bool_col" bf_result = ( - scalars_df_index[col_name].rank(na_option=na_option, method=method).compute() + scalars_df_index[col_name].rank(na_option=na_option, method=method).to_pandas() ) pd_result = ( scalars_pandas_df_index[col_name] @@ -1191,7 +1497,7 @@ def test_rank_with_nulls(scalars_df_index, scalars_pandas_df_index, na_option, m ) def test_nlargest(scalars_df_index, scalars_pandas_df_index, keep): col_name = "bool_col" - bf_result = scalars_df_index[col_name].nlargest(4, keep=keep).compute() + bf_result = scalars_df_index[col_name].nlargest(4, keep=keep).to_pandas() pd_result = scalars_pandas_df_index[col_name].nlargest(4, keep=keep) pd.testing.assert_series_equal( @@ -1200,6 +1506,29 @@ def test_nlargest(scalars_df_index, scalars_pandas_df_index, keep): ) +@pytest.mark.parametrize( + ("periods",), + [ + (1,), + (2,), + (-1,), + ], +) +def test_diff(scalars_df_index, scalars_pandas_df_index, periods): + bf_result = scalars_df_index["int64_col"].diff(periods=periods).to_pandas() + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = ( + scalars_pandas_df_index["int64_col"] + .diff(periods=periods) + .astype(pd.Int64Dtype()) + ) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + @pytest.mark.parametrize( ("keep",), [ @@ -1210,7 +1539,7 @@ def test_nlargest(scalars_df_index, scalars_pandas_df_index, keep): ) def test_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): col_name = "bool_col" - bf_result = scalars_df_index[col_name].nsmallest(2, keep=keep).compute() + bf_result = scalars_df_index[col_name].nsmallest(2, keep=keep).to_pandas() pd_result = scalars_pandas_df_index[col_name].nsmallest(2, keep=keep) pd.testing.assert_series_equal( @@ -1221,7 +1550,7 @@ def test_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): def test_rank_ints(scalars_df_index, scalars_pandas_df_index): col_name = "int64_too" - bf_result = scalars_df_index[col_name].rank().compute() + bf_result = scalars_df_index[col_name].rank().to_pandas() pd_result = scalars_pandas_df_index[col_name].rank().astype(pd.Float64Dtype()) pd.testing.assert_series_equal( @@ -1230,9 +1559,33 @@ def test_rank_ints(scalars_df_index, scalars_pandas_df_index): ) +def test_cast_float_to_int(scalars_df_index, scalars_pandas_df_index): + col_name = "float64_col" + bf_result = scalars_df_index[col_name].astype(pd.Int64Dtype()).to_pandas() + # cumsum does not behave well on nullable floats in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index[col_name].astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_cast_float_to_bool(scalars_df_index, scalars_pandas_df_index): + col_name = "float64_col" + bf_result = scalars_df_index[col_name].astype(pd.BooleanDtype()).to_pandas() + # cumsum does not behave well on nullable floats in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index[col_name].astype(pd.BooleanDtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + def test_cumsum_nested(scalars_df_index, scalars_pandas_df_index): col_name = "float64_col" - bf_result = scalars_df_index[col_name].cumsum().cumsum().cumsum().compute() + bf_result = scalars_df_index[col_name].cumsum().cumsum().cumsum().to_pandas() # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA pd_result = ( scalars_pandas_df_index[col_name] @@ -1252,7 +1605,7 @@ def test_cumsum_int_filtered(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" bf_col = scalars_df_index[col_name] - bf_result = bf_col[bf_col > -2].cumsum().compute() + bf_result = bf_col[bf_col > -2].cumsum().to_pandas() pd_col = scalars_pandas_df_index[col_name] # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA @@ -1266,7 +1619,7 @@ def test_cumsum_int_filtered(scalars_df_index, scalars_pandas_df_index): def test_cumsum_float(scalars_df_index, scalars_pandas_df_index): col_name = "float64_col" - bf_result = scalars_df_index[col_name].cumsum().compute() + bf_result = scalars_df_index[col_name].cumsum().to_pandas() # cumsum does not behave well on nullable floats in pandas, produces object type and never ignores NA pd_result = scalars_pandas_df_index[col_name].cumsum().astype(pd.Float64Dtype()) @@ -1278,7 +1631,7 @@ def test_cumsum_float(scalars_df_index, scalars_pandas_df_index): def test_cummin_int(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" - bf_result = scalars_df_index[col_name].cummin().compute() + bf_result = scalars_df_index[col_name].cummin().to_pandas() pd_result = scalars_pandas_df_index[col_name].cummin() pd.testing.assert_series_equal( @@ -1289,7 +1642,7 @@ def test_cummin_int(scalars_df_index, scalars_pandas_df_index): def test_cummax_int(scalars_df_index, scalars_pandas_df_index): col_name = "int64_col" - bf_result = scalars_df_index[col_name].cummax().compute() + bf_result = scalars_df_index[col_name].cummax().to_pandas() pd_result = scalars_pandas_df_index[col_name].cummax() pd.testing.assert_series_equal( @@ -1302,7 +1655,7 @@ def test_value_counts(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_result = scalars_df[col_name].value_counts().compute() + bf_result = scalars_df[col_name].value_counts().to_pandas() pd_result = scalars_pandas_df[col_name].value_counts() # Older pandas version may not have these values, bigframes tries to emulate 2.0+ @@ -1315,9 +1668,29 @@ def test_value_counts(scalars_dfs): ) +def test_value_counts_w_cut(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "int64_col" + + bf_cut = bigframes.pandas.cut(scalars_df[col_name], 3, labels=False) + pd_cut = pd.cut(scalars_pandas_df[col_name], 3, labels=False) + + bf_result = bf_cut.value_counts().to_pandas() + pd_result = pd_cut.value_counts() + # Older pandas version may not have these values, bigframes tries to emulate 2.0+ + pd_result.name = "count" + pd_result.index.name = col_name + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.Int64Dtype()), + ) + + def test_iloc_nested(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index["string_col"].iloc[1:].iloc[1:].compute() + bf_result = scalars_df_index["string_col"].iloc[1:].iloc[1:].to_pandas() pd_result = scalars_pandas_df_index["string_col"].iloc[1:].iloc[1:] pd.testing.assert_series_equal( @@ -1347,7 +1720,7 @@ def test_iloc_nested(scalars_df_index, scalars_pandas_df_index): ], ) def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, step): - bf_result = scalars_df_index["string_col"].iloc[start:stop:step].compute() + bf_result = scalars_df_index["string_col"].iloc[start:stop:step].to_pandas() pd_result = scalars_pandas_df_index["string_col"].iloc[start:stop:step] # Pandas may assign non-object dtype to empty series and series index @@ -1362,7 +1735,7 @@ def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, ste def test_series_add_prefix(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index["int64_too"].add_prefix("prefix_").compute() + bf_result = scalars_df_index["int64_too"].add_prefix("prefix_").to_pandas() pd_result = scalars_pandas_df_index["int64_too"].add_prefix("prefix_") @@ -1375,7 +1748,7 @@ def test_series_add_prefix(scalars_df_index, scalars_pandas_df_index): def test_series_add_suffix(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index["int64_too"].add_suffix("_suffix").compute() + bf_result = scalars_df_index["int64_too"].add_suffix("_suffix").to_pandas() pd_result = scalars_pandas_df_index["int64_too"].add_suffix("_suffix") @@ -1391,7 +1764,7 @@ def test_where_with_series(scalars_df_index, scalars_pandas_df_index): bf_result = ( scalars_df_index["int64_col"] .where(scalars_df_index["bool_col"], scalars_df_index["int64_too"]) - .compute() + .to_pandas() ) pd_result = scalars_pandas_df_index["int64_col"].where( scalars_pandas_df_index["bool_col"], scalars_pandas_df_index["int64_too"] @@ -1411,7 +1784,7 @@ def test_where_with_different_indices(scalars_df_index, scalars_pandas_df_index) scalars_df_index["bool_col"].iloc[2:], scalars_df_index["int64_too"].iloc[:5], ) - .compute() + .to_pandas() ) pd_result = ( scalars_pandas_df_index["int64_col"] @@ -1430,7 +1803,7 @@ def test_where_with_different_indices(scalars_df_index, scalars_pandas_df_index) def test_where_with_default(scalars_df_index, scalars_pandas_df_index): bf_result = ( - scalars_df_index["int64_col"].where(scalars_df_index["bool_col"]).compute() + scalars_df_index["int64_col"].where(scalars_df_index["bool_col"]).to_pandas() ) pd_result = scalars_pandas_df_index["int64_col"].where( scalars_pandas_df_index["bool_col"] @@ -1446,7 +1819,7 @@ def test_clip(scalars_df_index, scalars_pandas_df_index): col_bf = scalars_df_index["int64_col"] lower_bf = scalars_df_index["int64_too"] - 1 upper_bf = scalars_df_index["int64_too"] + 1 - bf_result = col_bf.clip(lower_bf, upper_bf).compute() + bf_result = col_bf.clip(lower_bf, upper_bf).to_pandas() col_pd = scalars_pandas_df_index["int64_col"] lower_pd = scalars_pandas_df_index["int64_too"] - 1 @@ -1463,7 +1836,7 @@ def test_clip_filtered_two_sided(scalars_df_index, scalars_pandas_df_index): col_bf = scalars_df_index["int64_col"].iloc[::2] lower_bf = scalars_df_index["int64_too"].iloc[2:] - 1 upper_bf = scalars_df_index["int64_too"].iloc[:5] + 1 - bf_result = col_bf.clip(lower_bf, upper_bf).compute() + bf_result = col_bf.clip(lower_bf, upper_bf).to_pandas() col_pd = scalars_pandas_df_index["int64_col"].iloc[::2] lower_pd = scalars_pandas_df_index["int64_too"].iloc[2:] - 1 @@ -1479,7 +1852,7 @@ def test_clip_filtered_two_sided(scalars_df_index, scalars_pandas_df_index): def test_clip_filtered_one_sided(scalars_df_index, scalars_pandas_df_index): col_bf = scalars_df_index["int64_col"].iloc[::2] lower_bf = scalars_df_index["int64_too"].iloc[2:] - 1 - bf_result = col_bf.clip(lower_bf, None).compute() + bf_result = col_bf.clip(lower_bf, None).to_pandas() col_pd = scalars_pandas_df_index["int64_col"].iloc[::2] lower_pd = scalars_pandas_df_index["int64_too"].iloc[2:] - 1 @@ -1510,7 +1883,9 @@ def test_dot(scalars_dfs): ], ) def test_between(scalars_df_index, scalars_pandas_df_index, left, right, inclusive): - bf_result = scalars_df_index["int64_col"].between(left, right, inclusive).compute() + bf_result = ( + scalars_df_index["int64_col"].between(left, right, inclusive).to_pandas() + ) pd_result = scalars_pandas_df_index["int64_col"].between(left, right, inclusive) pd.testing.assert_series_equal( @@ -1522,7 +1897,7 @@ def test_between(scalars_df_index, scalars_pandas_df_index, left, right, inclusi def test_to_frame(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs - bf_result = scalars_df["int64_col"].to_frame().compute() + bf_result = scalars_df["int64_col"].to_frame().to_pandas() pd_result = scalars_pandas_df["int64_col"].to_frame() assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) @@ -1651,7 +2026,7 @@ def test_sort_values(scalars_df_index, scalars_pandas_df_index, ascending, na_po bf_result = ( scalars_df_index["int64_col"] .sort_values(ascending=ascending, na_position=na_position) - .compute() + .to_pandas() ) pd_result = scalars_pandas_df_index["int64_col"].sort_values( ascending=ascending, na_position=na_position @@ -1671,7 +2046,9 @@ def test_sort_values(scalars_df_index, scalars_pandas_df_index, ascending, na_po ], ) def test_sort_index(scalars_df_index, scalars_pandas_df_index, ascending): - bf_result = scalars_df_index["int64_too"].sort_index(ascending=ascending).compute() + bf_result = ( + scalars_df_index["int64_too"].sort_index(ascending=ascending).to_pandas() + ) pd_result = scalars_pandas_df_index["int64_too"].sort_index(ascending=ascending) pd.testing.assert_series_equal( @@ -1685,7 +2062,7 @@ def test_mask_default_value(scalars_dfs): bf_col = scalars_df["int64_col"] bf_col_masked = bf_col.mask(bf_col % 2 == 1) - bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).compute() + bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).to_pandas() pd_col = scalars_pandas_df["int64_col"] pd_col_masked = pd_col.mask(pd_col % 2 == 1) @@ -1699,7 +2076,7 @@ def test_mask_custom_value(scalars_dfs): bf_col = scalars_df["int64_col"] bf_col_masked = bf_col.mask(bf_col % 2 == 1, -1) - bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).compute() + bf_result = bf_col.to_frame().assign(int64_col_masked=bf_col_masked).to_pandas() pd_col = scalars_pandas_df["int64_col"] pd_col_masked = pd_col.mask(pd_col % 2 == 1, -1) @@ -1739,14 +2116,14 @@ def test_mask_custom_value(scalars_dfs): ], ) def test_astype(scalars_df_index, scalars_pandas_df_index, column, to_type): - bf_result = scalars_df_index[column].astype(to_type).compute() + bf_result = scalars_df_index[column].astype(to_type).to_pandas() pd_result = scalars_pandas_df_index[column].astype(to_type) pd.testing.assert_series_equal(bf_result, pd_result) @pytest.mark.parametrize( "index", - [0, 5], + [0, 5, -2], ) def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): bf_result = scalars_df_index.string_col.iloc[index] @@ -1763,7 +2140,7 @@ def test_iloc_single_integer_out_of_bound_error( def test_loc_bool_series_explicit_index(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.string_col.loc[scalars_df_index.bool_col].compute() + bf_result = scalars_df_index.string_col.loc[scalars_df_index.bool_col].to_pandas() pd_result = scalars_pandas_df_index.string_col.loc[scalars_pandas_df_index.bool_col] pd.testing.assert_series_equal( @@ -1777,7 +2154,7 @@ def test_loc_bool_series_default_index( ): bf_result = scalars_df_default_index.string_col.loc[ scalars_df_default_index.bool_col - ].compute() + ].to_pandas() pd_result = scalars_pandas_df_default_index.string_col.loc[ scalars_pandas_df_default_index.bool_col ] @@ -1815,7 +2192,19 @@ def test_rename(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.string_col.rename("newname") pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), + pd_result, + ) + + +def test_rename_dict_same_type(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.rename({1: 100, 2: 200}) + pd_result = scalars_pandas_df_index.string_col.rename({1: 100, 2: 200}) + + pd_result.index = pd_result.index.astype("Int64") + + pd.testing.assert_series_equal( + bf_result.to_pandas(), pd_result, ) @@ -1825,7 +2214,7 @@ def test_rename_axis(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.string_col.rename_axis("newindexname") pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1842,7 +2231,7 @@ def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.string_col.loc[index_list] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1854,7 +2243,7 @@ def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.bool_col.loc[index_list] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1866,7 +2255,7 @@ def test_iloc_list(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.string_col.iloc[index_list] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1880,7 +2269,7 @@ def test_iloc_list_nameless(scalars_df_index, scalars_pandas_df_index): pd_result = pd_series.iloc[index_list] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1895,7 +2284,7 @@ def test_loc_list_nameless(scalars_df_index, scalars_pandas_df_index): pd_result = pd_series.loc[index_list] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1911,7 +2300,7 @@ def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.date_col.loc[pd_string_series] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) @@ -1924,11 +2313,46 @@ def test_loc_bf_index_integer_index(scalars_df_index, scalars_pandas_df_index): pd_result = scalars_pandas_df_index.date_col.loc[pd_index] pd.testing.assert_series_equal( - bf_result.compute(), + bf_result.to_pandas(), pd_result, ) +def test_loc_single_index_with_duplicate(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.set_index("string_col", drop=False) + scalars_pandas_df_index = scalars_pandas_df_index.set_index( + "string_col", drop=False + ) + index = "Hello, World!" + bf_result = scalars_df_index.date_col.loc[index] + pd_result = scalars_pandas_df_index.date_col.loc[index] + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.set_index("int64_too", drop=False) + scalars_pandas_df_index = scalars_pandas_df_index.set_index("int64_too", drop=False) + index = -2345 + bf_result = scalars_df_index.date_col.loc[index] + pd_result = scalars_pandas_df_index.date_col.loc[index] + assert bf_result.to_pandas().iloc[0] == pd_result + + def test_series_bool_interpretation_error(scalars_df_index): with pytest.raises(ValueError): True if scalars_df_index["string_col"] else False + + +def test_query_job_setters(scalars_dfs): + job_ids = set() + df, _ = scalars_dfs + series = df["int64_col"] + assert series.query_job is not None + repr(series) + job_ids.add(series.query_job.job_id) + series.to_pandas() + job_ids.add(series.query_job.job_id) + assert len(job_ids) == 2 diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 31ab4dee121..599b8aabbc9 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -27,6 +27,8 @@ import bigframes.dtypes import bigframes.ml.linear_model +FIRST_FILE = "000000000000" + def test_read_gbq_tokyo( session_tokyo: bigframes.Session, @@ -35,7 +37,7 @@ def test_read_gbq_tokyo( tokyo_location: str, ): df = session_tokyo.read_gbq(scalars_table_tokyo, index_col=["rowindex"]) - result = df.sort_index().compute() + result = df.sort_index().to_pandas() expected = scalars_pandas_df_index _, query_job = df._block.expr.start_query() @@ -86,7 +88,11 @@ def test_read_gbq_w_col_order( @pytest.mark.parametrize( ("query_or_table", "index_col"), [ - pytest.param("{scalars_table_id}", ["bool_col", "int64_col"], id="multiindex"), + pytest.param( + "{scalars_table_id}", + ["bool_col", "int64_col"], + id="unique_multiindex_table", + ), pytest.param( """SELECT t.float64_col * 2 AS my_floats, @@ -97,11 +103,43 @@ def test_read_gbq_w_col_order( ["my_strings"], id="string_index", ), + pytest.param( + "SELECT GENERATE_UUID() AS uuid, 0 AS my_value FROM UNNEST(GENERATE_ARRAY(1, 20))", + ["uuid"], + id="unique_uuid_index_query", + ), pytest.param( "{scalars_table_id}", ["bool_col"], id="non_unique_index", ), + pytest.param( + "{scalars_table_id}", + ["float64_col"], + id="non_unique_float_index", + ), + pytest.param( + "{scalars_table_id}", + [ + "timestamp_col", + "float64_col", + "datetime_col", + "int64_too", + ], + id="multi_part_index_direct", + ), + pytest.param( + "SELECT * FROM {scalars_table_id}", + [ + "timestamp_col", + "float64_col", + "string_col", + "bool_col", + "int64_col", + "int64_too", + ], + id="multi_part_index_w_query", + ), ], ) def test_read_gbq_w_index_col( @@ -116,6 +154,11 @@ def test_read_gbq_w_index_col( ) assert list(df.index.names) == index_col + # Verify that we get the expected number of results. + bf_shape = df.shape + result = df.to_pandas() + assert bf_shape == result.shape + @pytest.mark.parametrize( ("query_or_table", "max_results"), @@ -152,7 +195,7 @@ def test_read_gbq_w_max_results( query_or_table.format(scalars_table_id=scalars_table_id), max_results=max_results, ) - bf_result = df.compute() + bf_result = df.to_pandas() assert bf_result.shape[0] == max_results @@ -181,7 +224,7 @@ def test_read_pandas(session, scalars_dfs): df = session.read_pandas(scalars_pandas_df) assert df._block._expr._ordering is not None - result = df.compute() + result = df.to_pandas() expected = scalars_pandas_df pd.testing.assert_frame_equal(result, expected) @@ -189,7 +232,7 @@ def test_read_pandas(session, scalars_dfs): def test_read_pandas_multi_index(session, scalars_pandas_df_multi_index): df = session.read_pandas(scalars_pandas_df_multi_index) - result = df.compute() + result = df.to_pandas() pd.testing.assert_frame_equal(result, scalars_pandas_df_multi_index) @@ -199,7 +242,8 @@ def test_read_pandas_rowid_exists_adds_suffix(session, scalars_pandas_df_default ) df = session.read_pandas(scalars_pandas_df_default_index) - assert df._block._expr._ordering.ordering_id == "rowid_2" + total_order_col = df._block._expr._ordering.total_order_col + assert total_order_col and total_order_col.column_id == "rowid_2" def test_read_pandas_tokyo( @@ -208,7 +252,7 @@ def test_read_pandas_tokyo( tokyo_location: str, ): df = session_tokyo.read_pandas(scalars_pandas_df_index) - result = df.compute() + result = df.to_pandas() expected = scalars_pandas_df_index _, query_job = df._block.expr.start_query() @@ -220,14 +264,15 @@ def test_read_pandas_tokyo( def test_read_csv_gcs_default_engine(session, scalars_dfs, gcs_folder): scalars_df, _ = scalars_dfs if scalars_df.index.name is not None: - path = gcs_folder + "test_read_csv_gcs_default_engine_w_index.csv" + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index*.csv" else: - path = gcs_folder + "test_read_csv_gcs_default_engine_wo_index.csv" + path = gcs_folder + "test_read_csv_gcs_default_engine_wo_index*.csv" + read_path = path.replace("*", FIRST_FILE) scalars_df.to_csv(path, index=False) dtype = scalars_df.dtypes.to_dict() dtype.pop("geography_col") df = session.read_csv( - path, + read_path, # Convert default pandas dtypes to match BigQuery DataFrames dtypes. dtype=dtype, ) @@ -247,9 +292,9 @@ def test_read_csv_gcs_default_engine(session, scalars_dfs, gcs_folder): def test_read_csv_gcs_bq_engine(session, scalars_dfs, gcs_folder): scalars_df, _ = scalars_dfs if scalars_df.index.name is not None: - path = gcs_folder + "test_read_csv_gcs_bq_engine_w_index.csv" + path = gcs_folder + "test_read_csv_gcs_bq_engine_w_index*.csv" else: - path = gcs_folder + "test_read_csv_gcs_bq_engine_wo_index.csv" + path = gcs_folder + "test_read_csv_gcs_bq_engine_wo_index*.csv" scalars_df.to_csv(path, index=False) df = session.read_csv(path, engine="bigquery") @@ -419,22 +464,25 @@ def test_read_csv_default_engine_throws_not_implemented_error( match, ): path = ( - gcs_folder + "test_read_csv_gcs_default_engine_throws_not_implemented_error.csv" + gcs_folder + + "test_read_csv_gcs_default_engine_throws_not_implemented_error*.csv" ) + read_path = path.replace("*", FIRST_FILE) scalars_df_index.to_csv(path) with pytest.raises(NotImplementedError, match=match): - session.read_csv(path, **kwargs) + session.read_csv(read_path, **kwargs) def test_read_csv_gcs_default_engine_w_header(session, scalars_df_index, gcs_folder): - path = gcs_folder + "test_read_csv_gcs_default_engine_w_header.csv" + path = gcs_folder + "test_read_csv_gcs_default_engine_w_header*.csv" + read_path = path.replace("*", FIRST_FILE) scalars_df_index.to_csv(path) # Skips header=N rows, normally considers the N+1th row as the header, but overridden by # passing the `names` argument. In this case, pandas will skip the N+1th row too, take # the column names from `names`, and begin reading data from the N+2th row. df = session.read_csv( - path, + read_path, header=2, names=scalars_df_index.columns.to_list(), ) @@ -443,7 +491,7 @@ def test_read_csv_gcs_default_engine_w_header(session, scalars_df_index, gcs_fol def test_read_csv_gcs_bq_engine_w_header(session, scalars_df_index, gcs_folder): - path = gcs_folder + "test_read_csv_gcs_bq_engine_w_header.csv" + path = gcs_folder + "test_read_csv_gcs_bq_engine_w_header*.csv" scalars_df_index.to_csv(path, index=False) # Skip the header and the first 2 data rows. Without provided schema, the column names @@ -487,10 +535,11 @@ def test_read_csv_local_bq_engine_w_header(session, scalars_pandas_df_index): def test_read_csv_gcs_default_engine_w_index_col_name( session, scalars_df_default_index, gcs_folder ): - path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_name.csv" + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_name*.csv" + read_path = path.replace("*", FIRST_FILE) scalars_df_default_index.to_csv(path) - df = session.read_csv(path, index_col="rowindex") + df = session.read_csv(read_path, index_col="rowindex") scalars_df_default_index = scalars_df_default_index.set_index( "rowindex" ).sort_index() @@ -501,11 +550,12 @@ def test_read_csv_gcs_default_engine_w_index_col_name( def test_read_csv_gcs_default_engine_w_index_col_index( session, scalars_df_default_index, gcs_folder ): - path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_index.csv" + path = gcs_folder + "test_read_csv_gcs_default_engine_w_index_col_index*.csv" + read_path = path.replace("*", FIRST_FILE) scalars_df_default_index.to_csv(path) index_col = scalars_df_default_index.columns.to_list().index("rowindex") - df = session.read_csv(path, index_col=index_col) + df = session.read_csv(read_path, index_col=index_col) scalars_df_default_index = scalars_df_default_index.set_index( "rowindex" ).sort_index() @@ -559,11 +609,12 @@ def test_read_csv_local_default_engine_w_index_col_index( ) def test_read_csv_gcs_w_usecols(session, scalars_df_index, gcs_folder, engine): path = gcs_folder + "test_read_csv_gcs_w_usecols" - path = path + "_default_engine.csv" if engine is None else path + "_bq_engine.csv" + path = path + "_default_engine*.csv" if engine is None else path + "_bq_engine*.csv" + read_path = path.replace("*", FIRST_FILE) if engine is None else path scalars_df_index.to_csv(path) # df should only have 1 column which is bool_col. - df = session.read_csv(path, usecols=["bool_col"], engine=engine) + df = session.read_csv(read_path, usecols=["bool_col"], engine=engine) assert len(df.columns) == 1 diff --git a/tests/system/small/test_window.py b/tests/system/small/test_window.py index 34b56fcd730..e2f0fe999b1 100644 --- a/tests/system/small/test_window.py +++ b/tests/system/small/test_window.py @@ -19,42 +19,31 @@ @pytest.mark.parametrize( ("windowing"), [ - (lambda x: x.expanding()), - (lambda x: x.rolling(3, min_periods=3)), - (lambda x: x.groupby(x % 2).rolling(3, min_periods=3)), - (lambda x: x.groupby(x % 3).expanding(min_periods=2)), - ], - ids=[ - "expanding", - "rolling", - "rollinggroupby", - "expandinggroupby", + pytest.param(lambda x: x.expanding(), id="expanding"), + pytest.param(lambda x: x.rolling(3, min_periods=3), id="rolling"), + pytest.param( + lambda x: x.groupby(x % 2).rolling(3, min_periods=3), id="rollinggroupby" + ), + pytest.param( + lambda x: x.groupby(x % 3).expanding(min_periods=2), id="expandinggroupby" + ), ], ) @pytest.mark.parametrize( ("agg_op"), [ - (lambda x: x.sum()), - (lambda x: x.min()), - (lambda x: x.max()), - (lambda x: x.mean()), - (lambda x: x.count()), - (lambda x: x.std()), - (lambda x: x.var()), - ], - ids=[ - "sum", - "min", - "max", - "mean", - "count", - "std", - "var", + pytest.param(lambda x: x.sum(), id="sum"), + pytest.param(lambda x: x.min(), id="min"), + pytest.param(lambda x: x.max(), id="max"), + pytest.param(lambda x: x.mean(), id="mean"), + pytest.param(lambda x: x.count(), id="count"), + pytest.param(lambda x: x.std(), id="std"), + pytest.param(lambda x: x.var(), id="var"), ], ) def test_window_agg_ops(scalars_df_index, scalars_pandas_df_index, windowing, agg_op): col_name = "int64_too" - bf_series = agg_op(windowing(scalars_df_index[col_name])).compute() + bf_series = agg_op(windowing(scalars_df_index[col_name])).to_pandas() pd_series = agg_op(windowing(scalars_pandas_df_index[col_name])) # Pandas always converts to float64, even for min/max/count, which is not desired diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 92cb52a6815..6f5c96da49d 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -45,8 +45,10 @@ def test_standard_scaler_produces_correct_sql(): def test_one_hot_encoder_produces_correct_sql(): - sql = ml_sql.ml_one_hot_encoder("col_a", "encoded_col_a") - assert sql == "ML.ONE_HOT_ENCODER(col_a) OVER() AS encoded_col_a" + sql = ml_sql.ml_one_hot_encoder("col_a", "none", 1000000, 0, "encoded_col_a") + assert ( + sql == "ML.ONE_HOT_ENCODER(col_a, 'none', 1000000, 0) OVER() AS encoded_col_a" + ) def test_create_model_produces_correct_sql(): @@ -57,7 +59,7 @@ def test_create_model_produces_correct_sql(): ) assert ( sql - == """CREATE MODEL `my_dataset.my_model` + == """CREATE TEMP MODEL `my_dataset.my_model` my_options_sql AS my_source_sql""" ) @@ -65,14 +67,14 @@ def test_create_model_produces_correct_sql(): def test_create_model_transform_produces_correct_sql(): sql = ml_sql.create_model( - model_name="my_dataset.my_model", + model_name="my_model", source_sql="my_source_sql", options_sql="my_options_sql", transform_sql="my_transform_sql", ) assert ( sql - == """CREATE MODEL `my_dataset.my_model` + == """CREATE TEMP MODEL `my_model` my_transform_sql my_options_sql AS my_source_sql""" @@ -81,18 +83,30 @@ def test_create_model_transform_produces_correct_sql(): def test_create_remote_model_produces_correct_sql(): sql = ml_sql.create_remote_model( - model_name="my_dataset.my_model", + model_name="my_model", connection_name="my_project.us.my_connection", options_sql="my_options_sql", ) assert ( sql - == """CREATE MODEL `my_dataset.my_model` + == """CREATE TEMP MODEL `my_model` REMOTE WITH CONNECTION `my_project.us.my_connection` my_options_sql""" ) +def test_create_imported_model_produces_correct_sql(): + sql = ml_sql.create_imported_model( + model_name="my_model", + options_sql="my_options_sql", + ) + assert ( + sql + == """CREATE TEMP MODEL `my_model` +my_options_sql""" + ) + + def test_alter_model_correct_sql(): sql = ml_sql.alter_model( model_name="my_dataset.my_model", diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 25286189637..123dae79395 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -12,80 +12,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import mock - import ibis -from ibis.expr.types import Column, Table +from ibis.expr.types import Table from bigframes import core +ORDERING = core.ExpressionOrdering( + [ + core.OrderingColumnReference("int64_col"), + core.OrderingColumnReference("string_col"), + ], + total_ordering_columns=frozenset(["int64_col", "string_col"]), +) + def test_constructor_from_ibis_table_adds_all_columns( session, scalars_ibis_table: Table ): - actual = core.ArrayValue(session=session, table=scalars_ibis_table) + columns = tuple(scalars_ibis_table[key] for key in scalars_ibis_table.columns) + actual = core.ArrayValue( + session=session, table=scalars_ibis_table, columns=columns, ordering=ORDERING + ) assert actual._table is scalars_ibis_table assert len(actual._columns) == len(scalars_ibis_table.columns) -def test_builder_doesnt_change_original(session): - mock_table = mock.create_autospec(Table) - mock_column = mock.create_autospec(Column) - original = core.ArrayValue(session=session, table=mock_table, columns=[mock_column]) - assert original._table is mock_table - assert len(original._columns) == 1 - assert original._columns[0] is mock_column - - # Create a new expression from a builder. - builder = original.builder() - new_table = mock.create_autospec(Table) - assert new_table is not mock_table - builder.table = new_table - new_column = mock.create_autospec(Column) - assert new_column is not mock_column - builder.columns.append(new_column) - actual = builder.build() - - # Expected values are present. - assert actual._table is new_table - assert len(actual._columns) == 2 - assert actual._columns[0] is mock_column - assert actual._columns[1] is new_column - # Don't modify the original. - assert original._table is mock_table - assert len(original._columns) == 1 - assert original._columns[0] is mock_column - - -def test_projection_doesnt_change_original(session): - mock_table = mock.create_autospec(Table) - mock_column = mock.create_autospec(Column) - original = core.ArrayValue(session=session, table=mock_table, columns=[mock_column]) - assert original._table is mock_table - assert len(original._columns) == 1 - assert original._columns[0] is mock_column - - # Create a new expression from a projection. - new_column_1 = mock.create_autospec(Column) - new_column_2 = mock.create_autospec(Column) - assert new_column_1 is not mock_column - assert new_column_2 is not mock_column - actual = original.projection([new_column_1, mock_column, new_column_2]) - - # Expected values are present. - assert actual._table is mock_table - assert len(actual._columns) == 3 - assert actual._columns[0] is new_column_1 - assert actual._columns[1] is mock_column - assert actual._columns[2] is new_column_2 - # Don't modify the original. - assert original._table is mock_table - assert len(original._columns) == 1 - assert original._columns[0] is mock_column - - def test_to_ibis_expr_with_projection(session, scalars_ibis_table: Table): - expr = core.ArrayValue(session=session, table=scalars_ibis_table).projection( + columns = tuple(scalars_ibis_table[key] for key in scalars_ibis_table.columns) + expr = core.ArrayValue( + session=session, table=scalars_ibis_table, columns=columns, ordering=ORDERING + ).projection( [ scalars_ibis_table["int64_col"], ibis.literal(123456789).name("literals"), diff --git a/tests/unit/test_dtypes.py b/tests/unit/test_dtypes.py index ec49315f0fe..dafed08980c 100644 --- a/tests/unit/test_dtypes.py +++ b/tests/unit/test_dtypes.py @@ -175,10 +175,9 @@ def test_literal_to_ibis_scalar_throws_on_incompatible_literal(): def test_remote_function_io_types_are_supported_bigframes_types(): - from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type + from bigframes.remote_function import ( + _supported_io_ibis_types as rf_supported_io_ibis_types, + ) - from bigframes.remote_function import _supported_io_types as rf_supported_io_types - - for python_type in rf_supported_io_types: - ibis_type = python_type_to_bigquery_type(python_type) + for ibis_type in rf_supported_io_ibis_types: assert ibis_type in bigframes.dtypes.IBIS_TO_BIGFRAMES diff --git a/tests/unit/test_formatting_helper.py b/tests/unit/test_formatting_helper.py new file mode 100644 index 00000000000..ea29869e824 --- /dev/null +++ b/tests/unit/test_formatting_helper.py @@ -0,0 +1,17 @@ +import pytest + +import bigframes.formatting_helpers as formatter + + +@pytest.mark.parametrize( + "test_input, expected", [(None, "N/A"), ("string", "N/A"), (100000, "100.0 kB")] +) +def test_get_formatted_bytes(test_input, expected): + assert formatter.get_formatted_bytes(test_input) == expected + + +@pytest.mark.parametrize( + "test_input, expected", [(None, None), ("string", "string"), (100000, "a minute")] +) +def test_get_formatted_time(test_input, expected): + assert formatter.get_formatted_time(test_input) == expected diff --git a/tests/unit/test_formatting_helpers.py b/tests/unit/test_formatting_helpers.py new file mode 100644 index 00000000000..9db9b372e21 --- /dev/null +++ b/tests/unit/test_formatting_helpers.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. + +import unittest.mock as mock + +import google.api_core.exceptions as api_core_exceptions +import google.cloud.bigquery as bigquery +import pytest + +import bigframes.constants as constants +import bigframes.formatting_helpers as formatting_helpers + + +def test_wait_for_query_job_error_includes_feedback_link(): + mock_query_job = mock.create_autospec(bigquery.QueryJob) + mock_query_job.result.side_effect = api_core_exceptions.BadRequest( + "Test message 123." + ) + + with pytest.raises(api_core_exceptions.BadRequest) as cap_exc: + formatting_helpers.wait_for_query_job(mock_query_job) + + cap_exc.match("Test message 123.") + cap_exc.match(constants.FEEDBACK_LINK) + + +def test_wait_for_job_error_includes_feedback_link(): + mock_job = mock.create_autospec(bigquery.LoadJob) + mock_job.result.side_effect = api_core_exceptions.BadRequest("Test message 123.") + + with pytest.raises(api_core_exceptions.BadRequest) as cap_exc: + formatting_helpers.wait_for_job(mock_job) + + cap_exc.match("Test message 123.") + cap_exc.match(constants.FEEDBACK_LINK) diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index 3e8e16ee6ab..a178a454388 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -15,10 +15,12 @@ import inspect import re import sys +import unittest.mock as mock +import pandas as pd import pytest -import bigframes.pandas +import bigframes.pandas as bpd import bigframes.session leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE) @@ -79,3 +81,31 @@ def test_method_matches_session(method_name: str): 1: ] assert pandas_signature.return_annotation == session_signature.return_annotation + + +def test_cut_raises_with_labels(): + with pytest.raises(NotImplementedError, match="Only labels=False"): + mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True) + bigframes.pandas.cut(mock_series, 4, labels=["a", "b", "c", "d"]) + + +@pytest.mark.parametrize( + ("bins",), + ( + (0,), + (-1,), + ), +) +def test_cut_raises_with_invalid_bins(bins: int): + with pytest.raises(ValueError, match="`bins` should be a positive integer."): + mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True) + bigframes.pandas.cut(mock_series, bins, labels=False) + + +def test_pandas_attribute(): + assert bpd.NA is pd.NA + assert bpd.BooleanDtype is pd.BooleanDtype + assert bpd.Float64Dtype is pd.Float64Dtype + assert bpd.Int64Dtype is pd.Int64Dtype + assert bpd.StringDtype is pd.StringDtype + assert bpd.ArrowDtype is pd.ArrowDtype diff --git a/third_party/bigframes_vendored/ibis/LICENSE.txt b/third_party/bigframes_vendored/ibis/LICENSE.txt new file mode 100644 index 00000000000..d6456956733 --- /dev/null +++ b/third_party/bigframes_vendored/ibis/LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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/third_party/bigframes_vendored/ibis/README.md b/third_party/bigframes_vendored/ibis/README.md new file mode 100644 index 00000000000..8a00750e920 --- /dev/null +++ b/third_party/bigframes_vendored/ibis/README.md @@ -0,0 +1,196 @@ +# Ibis + +[![Documentation Status](https://img.shields.io/badge/docs-docs.ibis--project.org-blue.svg)](http://ibis-project.org) +[![Anaconda-Server Badge](https://anaconda.org/conda-forge/ibis-framework/badges/version.svg)](https://anaconda.org/conda-forge/ibis-framework) +[![PyPI](https://img.shields.io/pypi/v/ibis-framework.svg)](https://pypi.org/project/ibis-framework) +[![Build status](https://github.com/ibis-project/ibis/actions/workflows/ibis-main.yml/badge.svg)](https://github.com/ibis-project/ibis/actions/workflows/ibis-main.yml?query=branch%3Amaster) +[![Build status](https://github.com/ibis-project/ibis/actions/workflows/ibis-backends.yml/badge.svg)](https://github.com/ibis-project/ibis/actions/workflows/ibis-backends.yml?query=branch%3Amaster) +[![Codecov branch](https://img.shields.io/codecov/c/github/ibis-project/ibis/master.svg)](https://codecov.io/gh/ibis-project/ibis) + +## What is Ibis? + +Ibis is a Python library that provides a lightweight, universal interface for data wrangling. It helps Python users explore and transform data of any size, stored anywhere. + +Ibis has three primary components: + +1. **A dataframe API for Python**. + Python users can write Ibis code to manipulate tabular data. +2. **Interfaces to 15+ query engines.** + Wherever data is stored, people can use Ibis as their API of choice to communicate with any of those query engines. +3. **Deferred execution**. + Ibis uses deferred execution, so execution of code is pushed to the query engine. + Users can execute at the speed of their backend, not their local computer. + +## Why Use Ibis? + +Ibis aims to be a future-proof solution to interacting with data using Python and can accomplish this goal through its main features: + +- **Familiar API**: Ibis’s API design borrows from popular APIs like pandas and dplyr that most users already know and like to use. +- **Consistent syntax**: Ibis aims to be a universal Python API for tabular data of any size, big or small. +- **Deferred execution**: Ibis pushes code execution to the query engine and only moves required data into memory when necessary. + Analytics workflows are faster and more efficient +- **Interactive mode**: Ibis provides an interactive mode in which users can quickly diagnose problems, explore data, and mock up workflows and pipelines locally. +- **10+ supported backends**: Ibis supports multiple query engines and DataFrame APIs. + Use one interface to transform with your data wherever it lives: from DataFrames in pandas to Parquet files through DuckDB to tables in BigQuery. +- **Minimize rewrites**: Teams can often keep their Ibis code the same regardless of backend changes, like increasing or decreasing computing power, changing the number or size of their databases, or switching backends entirely. +- **Flexibility when you need it**: When Ibis doesn't support something, it provides a way to jump directly into SQL. + +## Common Use Cases + +- **Speed up prototype to production.** + Scale code written and tested locally to a distributed system or cloud SQL engine with minimal rewrites. +- **Boost performance of existing Python or pandas code.** + For example a general rule of thumb for pandas is "Have 5 to 10 times as much RAM as the size of your dataset". + When a dataset exceeds this rule using in-memory frameworks like pandas can be slow. + Instead, using Ibis will significantly speed up your workflows because of its deferred execution. + Ibis also empowers you to switch to a faster database engine, without changing much of your code. +- **Get rid of long, error-prone, `f`-strings.** + Ibis provides one syntax for multiple query engines and dataframe APIs that lets you avoid learning new flavors of SQL or other framework-specific code. + Learn the syntax once and use that syntax anywhere. + +## Backends + +Ibis acts as a universal frontend to the following systems: + +- [Apache Arrow DataFusion](https://ibis-project.org/backends/datafusion/) (experimental) +- [Apache Druid](https://ibis-project.org/backends/druid/) (experimental) +- [Apache Impala](https://ibis-project.org/backends/impala/) +- [Apache PySpark](https://ibis-project.org/backends/pyspark/) +- [BigQuery](https://ibis-project.org/backends/bigquery/) +- [ClickHouse](https://ibis-project.org/backends/clickhouse/) +- [Dask](https://ibis-project.org/backends/dask/) +- [DuckDB](https://ibis-project.org/backends/duckdb/) +- [HeavyAI](https://github.com/heavyai/ibis-heavyai) +- [MySQL](https://ibis-project.org/backends/mysql/) +- [Oracle](https://ibis-project.org/backends/oracle/) (experimental) +- [Pandas](https://ibis-project.org/backends/pandas/) +- [Polars](https://ibis-project.org/backends/polars/) (experimental) +- [PostgreSQL](https://ibis-project.org/backends/postgresql/) +- [SQL Server](https://ibis-project.org/backends/mssql/) +- [SQLite](https://ibis-project.org/backends/sqlite/) +- [Snowflake](https://ibis-project.org/backends/snowflake) (experimental) +- [Trino](https://ibis-project.org/backends/trino/) (experimental) + +The list of supported backends is continuously growing. Anyone can get involved +in adding new ones! Learn more about contributing to ibis in our contributing +documentation at https://github.com/ibis-project/ibis/blob/master/docs/CONTRIBUTING.md + +## Installation + +Install Ibis from PyPI with: + +```bash +pip install 'ibis-framework[duckdb]' +``` + +Or from conda-forge with: + +```bash +conda install ibis-framework -c conda-forge +``` + +(It’s a common mistake to `pip install ibis`. If you try to use Ibis and get errors early on try uninstalling `ibis` and installing `ibis-framework`) + +To discover ibis, we suggest starting with the DuckDB backend (which is included by default in the conda-forge package). The DuckDB backend is performant and fully featured. + +To use ibis with other backends, include the backend name in brackets for PyPI: + +```bash +pip install 'ibis-framework[postgres]' +``` + +Or use `ibis-$BACKEND` where `$BACKEND` is the specific backend you want to use when installing from conda-forge: + +```bash +conda install ibis-postgres -c conda-forge +``` + +## Getting Started with Ibis + +We provide a number of tutorial and example notebooks in the +[ibis-examples](https://github.com/ibis-project/ibis-examples). The easiest way +to try these out is through the online interactive notebook environment +provided here: +[![Binder](https://static.mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/ibis-project/ibis-examples/main) + +You can also get started analyzing any dataset, anywhere with just a few lines +of Ibis code. Here’s an example of how to use Ibis with a SQLite database. + +Download the SQLite database from the `ibis-tutorial-data` GCS (Google Cloud +Storage) bucket, then connect to it using ibis. + +```bash +curl -LsS -o geography.db 'https://storage.googleapis.com/ibis-tutorial-data/geography.db' +``` + +Connect to the database and show the available tables + +```python +>>> import ibis +>>> from ibis import _ +>>> ibis.options.interactive = True +>>> con = ibis.sqlite.connect("geography.db") +>>> con.tables +Tables +------ +- countries +- gdp +- independence +``` + +Choose the `countries` table and preview its first few rows + +```python +>>> countries = con.tables.countries +>>> countries.head() +┏━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━┓ +┃ iso_alpha2 ┃ iso_alpha3 ┃ iso_numeric ┃ fips ┃ name ┃ capital ┃ area_km2 ┃ population ┃ continent ┃ +┡━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━┩ +│ string │ string │ int32 │ string │ string │ string │ float64 │ int32 │ string │ +├────────────┼────────────┼─────────────┼────────┼──────────────────────┼──────────────────┼──────────┼────────────┼───────────┤ +│ AD │ AND │ 20 │ AN │ Andorra │ Andorra la Vella │ 468.0 │ 84000 │ EU │ +│ AE │ ARE │ 784 │ AE │ United Arab Emirates │ Abu Dhabi │ 82880.0 │ 4975593 │ AS │ +│ AF │ AFG │ 4 │ AF │ Afghanistan │ Kabul │ 647500.0 │ 29121286 │ AS │ +│ AG │ ATG │ 28 │ AC │ Antigua and Barbuda │ St. Johns │ 443.0 │ 86754 │ NA │ +│ AI │ AIA │ 660 │ AV │ Anguilla │ The Valley │ 102.0 │ 13254 │ NA │ +└────────────┴────────────┴─────────────┴────────┴──────────────────────┴──────────────────┴──────────┴────────────┴───────────┘ +``` + +Show the 5 least populous countries in Asia + +```python + +>>> ( +... countries.filter(_.continent == "AS") +... .select("name", "population") +... .order_by(_.population) +... .limit(5) +... ) +┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓ +┃ name ┃ population ┃ +┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩ +│ string │ int32 │ +├────────────────────────────────┼────────────┤ +│ Cocos [Keeling] Islands │ 628 │ +│ British Indian Ocean Territory │ 4000 │ +│ Brunei │ 395027 │ +│ Maldives │ 395650 │ +│ Macao │ 449198 │ +└────────────────────────────────┴────────────┘ +``` + +## Community and Contributing + +Ibis is an open source project and welcomes contributions from anyone in the community. + +- Read [the contributing guide](https://github.com/ibis-project/ibis/blob/master/docs/CONTRIBUTING.md). +- We care about keeping the community welcoming for all. Check out [the code of conduct](https://github.com/ibis-project/ibis/blob/master/docs/CODE_OF_CONDUCT.md). +- The Ibis project is open sourced under the [Apache License](https://github.com/ibis-project/ibis/blob/master/LICENSE.txt). + +Join our community here: + +- Twitter: https://twitter.com/IbisData +- Gitter: https://gitter.im/ibis-dev/Lobby +- StackOverflow: https://stackoverflow.com/questions/tagged/ibis + +For more information visit https://ibis-project.org/. diff --git a/third_party/bigframes_vendored/ibis/__init__.py b/third_party/bigframes_vendored/ibis/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/ibis/backends/__init__.py b/third_party/bigframes_vendored/ibis/backends/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/ibis/backends/bigquery/__init__.py b/third_party/bigframes_vendored/ibis/backends/bigquery/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py new file mode 100644 index 00000000000..d209284ab7d --- /dev/null +++ b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py @@ -0,0 +1,19 @@ +# Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/registry.py +"""Module to convert from Ibis expression to SQL string.""" + +from ibis.backends.bigquery.registry import OPERATION_REGISTRY + +import third_party.bigframes_vendored.ibis.expr.operations as vendored_ibis_ops + + +def _approx_quantiles(translator, op: vendored_ibis_ops.ApproximateMultiQuantile): + arg = translator.translate(op.arg) + num_bins = translator.translate(op.num_bins) + return f"APPROX_QUANTILES({arg}, {num_bins})" + + +patched_ops = { + vendored_ibis_ops.ApproximateMultiQuantile: _approx_quantiles, +} + +OPERATION_REGISTRY.update(patched_ops) diff --git a/third_party/bigframes_vendored/ibis/expr/__init__.py b/third_party/bigframes_vendored/ibis/expr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/ibis/expr/operations/__init__.py b/third_party/bigframes_vendored/ibis/expr/operations/__init__.py new file mode 100644 index 00000000000..f3ab753a3b1 --- /dev/null +++ b/third_party/bigframes_vendored/ibis/expr/operations/__init__.py @@ -0,0 +1,4 @@ +# Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/__init__.py +from __future__ import annotations + +from third_party.bigframes_vendored.ibis.expr.operations.reductions import * # noqa: F403 diff --git a/third_party/bigframes_vendored/ibis/expr/operations/reductions.py b/third_party/bigframes_vendored/ibis/expr/operations/reductions.py new file mode 100644 index 00000000000..5e6ad9ecf2b --- /dev/null +++ b/third_party/bigframes_vendored/ibis/expr/operations/reductions.py @@ -0,0 +1,23 @@ +# Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/reductions.py + +from __future__ import annotations + +import ibis.expr.datatypes as dt +from ibis.expr.operations.reductions import Filterable, Reduction +import ibis.expr.rules as rlz + + +class ApproximateMultiQuantile(Filterable, Reduction): + """Calculate (approximately) evenly-spaced quantiles. + + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/approximate_aggregate_functions#approx_quantiles + """ + + arg = rlz.any + num_bins = rlz.value(dt.int64) + output_dtype = dt.Array(dt.float64) + + +__all__ = [ + "ApproximateMultiQuantile", +] diff --git a/third_party/bigframes_vendored/pandas/core/config_init.py b/third_party/bigframes_vendored/pandas/core/config_init.py index 10f5546f108..198654015e0 100644 --- a/third_party/bigframes_vendored/pandas/core/config_init.py +++ b/third_party/bigframes_vendored/pandas/core/config_init.py @@ -16,10 +16,44 @@ Encapsulates configuration for displaying objects. Attributes: - max_columns (int): - If max_columns is exceeded, switch to truncate view. - max_rows (int): - If max_rows is exceeded, switch to truncate view. - progress_bar Optional(str): - Determines if progress bars are shown during job runs. Use auto | notebook | terminal. Set to 'None' to remove progress bars + max_columns (int, default 20): + If `max_columns` is exceeded, switch to truncate view. + max_rows (int, default 25): + If `max_rows` is exceeded, switch to truncate view. + progress_bar (Optional(str), default "auto"): + Determines if progress bars are shown during job runs. + Valid values are `auto`, `notebook`, and `terminal`. Set + to `None` to remove progress bars. + repr_mode (Literal[`head`, `deferred`]): + `head`: + Execute, download, and display results (limited to head) from + dataframe and series objects during repr. + `deferred`: + Prevent executions from repr statements in dataframe and series objects. + Instead estimated bytes processed will be shown. Dataframe and Series + objects can still be computed with methods that explicitly execute and + download results. +""" + +sampling_options_doc = """ +Encapsulates configuration for data sampling. + +Attributes: + max_download_size (int, default 500): + Download size threshold in MB. If value set to None, the download size + won't be checked. + enable_downsampling (bool, default False): + Whether to enable downsampling, If max_download_size is exceeded when + downloading data (e.g., to_pandas()), the data will be downsampled + if enable_downsampling is True, otherwise, an error will be raised. + sampling_method (str, default "uniform"): + Downsampling algorithms to be chosen from, the choices are: + "head": This algorithm returns a portion of the data from + the beginning. It is fast and requires minimal computations + to perform the downsampling.; "uniform": This algorithm returns + uniform random samples of the data. + random_state (int, default None): + The seed for the uniform downsampling algorithm. If provided, + the uniform method may take longer to execute and require more + computation. """ diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 93b7c7683cf..653b65c8346 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -13,6 +13,8 @@ from typing import Iterable, Literal, Mapping, Optional, Sequence, Union +import numpy + from third_party.bigframes_vendored.pandas.core.generic import NDFrame # ----------------------------------------------------------------------- @@ -33,8 +35,61 @@ def shape(self) -> tuple[int, int]: """Return a tuple representing the dimensionality of the DataFrame.""" raise NotImplementedError("abstract method") + @property + def axes(self) -> list: + """ + Return a list representing the axes of the DataFrame. + + It has the row axis labels and column axis labels as the only members. + They are returned in that order. + + Examples + .. code-block:: + + df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + df.axes + [RangeIndex(start=0, stop=2, step=1), Index(['col1', 'col2'], + dtype='object')] + """ + return [self.index, self.columns] + + @property + def values(self) -> numpy.ndarray: + """Return the values of DataFrame in the form of a NumPy array. + + Args: + dytype (default None): + The dtype to pass to `numpy.asarray()`. + copy (bool, default False): + Whether to ensure that the returned value is not a view + on another array. + na_value (default None): + The value to use for missing values. + """ + raise NotImplementedError("abstract method") + # ---------------------------------------------------------------------- # IO methods (to / from other formats) + def to_numpy( + self, dtype=None, copy=False, na_value=None, **kwargs + ) -> numpy.ndarray: + """ + Convert the DataFrame to a NumPy array. + + Args: + dtype (None): + The dtype to pass to `numpy.asarray()`. + copy (bool, default None): + Whether to ensure that the returned value is not a view + on another array. + na_value (Any, default None): + The value to use for missing values. The default value + depends on dtype and the dtypes of the DataFrame columns. + + Returns: + numpy.ndarray: The converted NumPy array. + """ + raise NotImplementedError("abstract method") def to_gbq( self, @@ -42,15 +97,16 @@ def to_gbq( *, if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", index: bool = True, + ordering_id: Optional[str] = None, ) -> None: - """Write a DataFrame to a Google BigQuery table. + """Write a DataFrame to a BigQuery table. Args: - destination_table : str + destination_table (str): Name of table to be written, in the form ``dataset.tablename`` or ``project.dataset.tablename``. - if_exists : str, default 'fail' + if_exists (str, default 'fail'): Behavior when the destination table exists. Value can be one of: ``'fail'`` @@ -60,7 +116,12 @@ def to_gbq( ``'append'`` If table exists, insert data. Create if does not exist. - index: whether write row names (index) or not. + index (bool. default True): + whether write row names (index) or not. + + ordering_id (Optional[str], default None): + If set, write the ordering of the DataFrame as a column in the + result table with this name. """ raise NotImplementedError("abstract method") @@ -70,19 +131,19 @@ def to_parquet( *, index: bool = True, ) -> None: - """Write a DataFrame to the binary parquet format. + """Write a DataFrame to the binary Parquet format. This function writes the dataframe as a `parquet file - `_ to Google Cloud Storage. + `_ to Cloud Storage. Args: - path: - Destination URI(s) of GCS files(s) to store the extracted dataframe + path (str): + Destination URI(s) of Cloud Storage files(s) to store the extracted dataframe in format of ``gs:///``. If the data size is more than 1GB, you must use a wildcard to export the data into multiple files and the size of the files varies. - index : bool, default None + index (bool, default True): If ``True``, include the dataframe's index(es) in the file output. If ``False``, they will not be written to the file. @@ -101,6 +162,12 @@ def assign(self, **kwargs) -> DataFrame: Returns a new object with all original columns in addition to new ones. Existing columns that are re-assigned will be overwritten. + .. note:: + Assigning multiple columns within the same ``assign`` is possible. + Later items in '\*\*kwargs' may refer to newly created or modified + columns in 'df'; items are computed and assigned into 'df' in + order. + Args: kwargs: A dictionary of ``{str: values}``. The column names are @@ -108,14 +175,8 @@ def assign(self, **kwargs) -> DataFrame: are simply assigned to the column. Returns: - A new DataFrame with the new columns in addition to + DataFrame: A new DataFrame with the new columns in addition to all the existing columns. - - Notes: - Assigning multiple columns within the same ``assign`` is possible. - Later items in '\*\*kwargs' may refer to newly created or modified - columns in 'df'; items are computed and assigned into 'df' in - order. """ raise NotImplementedError("abstract method") @@ -123,24 +184,31 @@ def assign(self, **kwargs) -> DataFrame: # Reindexing and alignment def drop( - self, - *, - columns: Union[str, Iterable[str]], + self, labels=None, *, axis=0, index=None, columns=None, level=None ) -> DataFrame | None: """Drop specified labels from columns. Remove columns by directly specifying column names. Args: + labels: + Index or column labels to drop. + axis: + Whether to drop labels from the index (0 or 'index') or + columns (1 or 'columns'). + index: + Alternative to specifying axis (``labels, axis=0`` + is equivalent to ``index=labels``). columns: - Single label or list-like. - + Alternative to specifying axis (``labels, axis=1`` + is equivalent to ``columns=labels``). + level: + For MultiIndex, level from which the labels will be removed. Returns: - DataFrame without the removed column labels. + DataFrame: DataFrame without the removed column labels. Raises: - KeyError: - If any of the labels is not found in the selected axis. + KeyError: If any of the labels is not found in the selected axis. """ raise NotImplementedError("abstract method") @@ -155,15 +223,14 @@ def rename( will be left as-is. Extra labels listed don't throw an error. Args: - columns: + columns (Mapping): Dict-like from old column labels to new column labels. Returns: - DataFrame with the renamed axis labels. + DataFrame: DataFrame with the renamed axis labels. Raises: - KeyError: - If any of the labels is not found. + KeyError: If any of the labels is not found. """ raise NotImplementedError("abstract method") @@ -171,14 +238,16 @@ def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: """ Set the name of the axis for the index. - Note: currently only accepts a single string parameter (the new name of the index) + .. Note:: + + Currently only accepts a single string parameter (the new name of the index). Args: - mapper : str + mapper str: Value to set the axis name attribute. Returns: - DataFrame with the new index name + DataFrame: DataFrame with the new index name """ raise NotImplementedError("abstract method") @@ -197,12 +266,11 @@ def set_index( Args: keys: A label. This parameter can be a single column key. - drop: + drop : Delete columns to be used as the new index. Returns: - DataFrame: - Changed row labels. + DataFrame: Changed row labels. """ raise NotImplementedError("abstract method") @@ -211,12 +279,12 @@ def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: Rearrange index levels using input order. May not drop or duplicate levels. Args: - order: list of int or list of str + order (list of int or list of str): List representing new level order. Reference level by number (position) or by key (label). Returns: - DataFrame + DataFrame: DataFrame of rearranged index. """ raise NotImplementedError("abstract method") @@ -225,12 +293,12 @@ def droplevel(self, level): Return DataFrame with requested index / column level(s) removed. Args: - level: int, str, or list-like + level (int, str, or list-like): If a string is given, must be the name of a level If list-like, elements must be names or positional indexes of levels. Returns: - DataFrame with requested index / column level(s) removed. + DataFrame: DataFrame with requested index / column level(s) removed. """ raise NotImplementedError("abstract method") @@ -244,12 +312,12 @@ def reset_index( Reset the index of the DataFrame, and use the default one instead. Args: - drop: + drop (bool, default False): Do not try to insert index into dataframe columns. This resets the index to the default integer index. Returns: - DataFrame with the new index. + DataFrame: DataFrame with the new index. """ raise NotImplementedError("abstract method") @@ -265,10 +333,10 @@ def drop_duplicates( are ignored. Args: - subset : column label or sequence of labels, optional + subset (column label or sequence of labels, optional): Only consider certain columns for identifying duplicates, by default use all of the columns. - keep : {'first', 'last', ``False``}, default 'first' + keep ({'first', 'last', ``False``}, default 'first'): Determines which duplicates (if any) to keep. - 'first' : Drop duplicates except for the first occurrence. @@ -276,7 +344,7 @@ def drop_duplicates( - ``False`` : Drop all duplicates. Returns: - DataFrame with duplicates removed + DataFrame: DataFrame with duplicates removed """ raise NotImplementedError("abstract method") @@ -287,10 +355,10 @@ def duplicated(self, subset=None, keep="first"): Considering certain columns is optional. Args: - subset : column label or sequence of labels, optional + subset (column label or sequence of labels, optional): Only consider certain columns for identifying duplicates, by default use all of the columns. - keep : {'first', 'last', False}, default 'first' + keep ({'first', 'last', False}, default 'first'): Determines which duplicates (if any) to mark. - ``first`` : Mark duplicates as ``True`` except for the first occurrence. @@ -311,7 +379,7 @@ def dropna( """Remove missing values. Returns: - DataFrame with NA entries dropped from it. + DataFrame: DataFrame with NA entries dropped from it. """ raise NotImplementedError("abstract method") @@ -323,18 +391,23 @@ def sort_values( by: str | Sequence[str], *, ascending: bool | Sequence[bool] = True, + kind: str = "quicksort", na_position="last", ) -> DataFrame: """Sort by the values along row axis. Args: - by: + by (str or Sequence[str]): Name or list of names to sort by. - ascending: + ascending (bool or Sequence[bool], default True): Sort ascending vs. descending. Specify list for multiple sort orders. If this is a list of bools, must match the length of the by. - na_position: + kind (str, default `quicksort`): + Choice of sorting algorithm. Accepts 'quicksort’, ‘mergesort’, + ‘heapsort’, ‘stable’. Ignored except when determining whether to + sort stably. 'mergesort' or 'stable' will result in stable reorder. + na_position ({'first', 'last'}, default `last`): ``{'first', 'last'}``, default 'last' Puts NaNs at the beginning if `first`; `last` puts NaNs at the end. @@ -356,8 +429,9 @@ def sort_index( # ---------------------------------------------------------------------- # Arithmetic Methods - def le(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'less than or equal to' of dataframe and other, element-wise (binary operator `<=`). + def eq(self, other, axis: str | int = "columns") -> DataFrame: + """ + Get equal to of DataFrame and other, element-wise (binary operator `eq`). Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison operators. @@ -366,24 +440,20 @@ def le(self, other, axis: str | int = "columns") -> DataFrame: (rows or columns) and level for comparison. Args: - other: scalar, sequence, Series, or DataFrame + other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. - axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): Whether to compare by the index (0 or 'index') or columns (1 or 'columns'). Returns: - DataFrame of bool. The result of the comparison. - - Notes: - Mismatched indices will be unioned together. `NaN` values in - floating point columns are considered different - (i.e. `NaN` != `NaN`). + Result of the comparison. """ raise NotImplementedError("abstract method") - def lt(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'less than' of dataframe and other, element-wise (binary operator `<`). + def ne(self, other, axis: str | int = "columns") -> DataFrame: + """ + Get not equal to of DataFrame and other, element-wise (binary operator `ne`). Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison operators. @@ -392,24 +462,44 @@ def lt(self, other, axis: str | int = "columns") -> DataFrame: (rows or columns) and level for comparison. Args: - other: scalar, sequence, Series, or DataFrame + other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. - axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): Whether to compare by the index (0 or 'index') or columns (1 or 'columns'). - Returns: - DataFrame of bool. The result of the comparison. + DataFrame: Result of the comparison. + """ + raise NotImplementedError("abstract method") - Notes: + def le(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'less than or equal to' of dataframe and other, element-wise (binary operator `<=`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + .. note:: Mismatched indices will be unioned together. `NaN` values in floating point columns are considered different (i.e. `NaN` != `NaN`). + + Args: + other (scalar, sequence, Series, or DataFrame): + Any single or multiple element data structure, or list-like object. + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame: DataFrame of bool. The result of the comparison. """ raise NotImplementedError("abstract method") - def ge(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'greater than or equal to' of dataframe and other, element-wise (binary operator `>=`). + def lt(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'less than' of DataFrame and other, element-wise (binary operator `<`). Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison operators. @@ -417,25 +507,51 @@ def ge(self, other, axis: str | int = "columns") -> DataFrame: Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis (rows or columns) and level for comparison. + .. note:: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + Args: - other: scalar, sequence, Series, or DataFrame + other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. - axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): Whether to compare by the index (0 or 'index') or columns (1 or 'columns'). Returns: - DataFrame of bool. The result of the comparison. + DataFrame: DataFrame of bool. The result of the comparison. + """ + raise NotImplementedError("abstract method") - Notes: + def ge(self, other, axis: str | int = "columns") -> DataFrame: + """Get 'greater than or equal to' of DataFrame and other, element-wise (binary operator `>=`). + + Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison + operators. + + Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis + (rows or columns) and level for comparison. + + .. note:: Mismatched indices will be unioned together. `NaN` values in floating point columns are considered different (i.e. `NaN` != `NaN`). + + Args: + other (scalar, sequence, Series, or DataFrame): + Any single or multiple element data structure, or list-like object. + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): + Whether to compare by the index (0 or 'index') or columns + (1 or 'columns'). + + Returns: + DataFrame: DataFrame of bool. The result of the comparison. """ raise NotImplementedError("abstract method") def gt(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'greater than' of dataframe and other, element-wise (binary operator `>`). + """Get 'greater than' of DataFrame and other, element-wise (binary operator `>`). Among flexible wrappers (`eq`, `ne`, `le`, `lt`, `ge`, `gt`) to comparison operators. @@ -443,260 +559,250 @@ def gt(self, other, axis: str | int = "columns") -> DataFrame: Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis (rows or columns) and level for comparison. + .. note:: + Mismatched indices will be unioned together. `NaN` values in + floating point columns are considered different + (i.e. `NaN` != `NaN`). + Args: - other: scalar, sequence, Series, or DataFrame + other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. - axis: {{0 or 'index', 1 or 'columns'}}, default 'columns' + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): Whether to compare by the index (0 or 'index') or columns (1 or 'columns'). Returns: - DataFrame of bool. The result of the comparison. - - Notes: - Mismatched indices will be unioned together. `NaN` values in - floating point columns are considered different - (i.e. `NaN` != `NaN`). + DataFrame: DataFrame of bool: The result of the comparison. """ raise NotImplementedError("abstract method") def add(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'addition' of dataframe and other, element-wise (binary operator `+`). + """Get addition of DataFrame and other, element-wise (binary operator `+`). Equivalent to ``dataframe + other``. With reverse version, `radd`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def sub(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'subtraction' of dataframe and other, element-wise (binary operator `-`). + """Get subtraction of DataFrame and other, element-wise (binary operator `-`). Equivalent to ``dataframe - other``. With reverse version, `rsub`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def rsub(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'subtraction' of dataframe and other, element-wise (binary operator `-`). + """Get subtraction of DataFrame and other, element-wise (binary operator `-`). Equivalent to ``other - dataframe``. With reverse version, `sub`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def mul(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'multiplication' of dataframe and other, element-wise (binary operator `*`). + """Get multiplication of DataFrame and other, element-wise (binary operator `*`). Equivalent to ``dataframe * other``. With reverse version, `rmul`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def truediv(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'floating division' of dataframe and other, element-wise (binary operator `/`). + """Get floating division of DataFrame and other, element-wise (binary operator `/`). Equivalent to ``dataframe / other``. With reverse version, `rtruediv`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'floating division' of dataframe and other, element-wise (binary operator `/`). + """Get floating division of DataFrame and other, element-wise (binary operator `/`). Equivalent to ``other / dataframe``. With reverse version, `truediv`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. """ raise NotImplementedError("abstract method") def floordiv(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'integer division' of dataframe and other, element-wise (binary operator `//`). + """Get integer division of DataFrame and other, element-wise (binary operator `//`). Equivalent to ``dataframe // other``. With reverse version, `rfloordiv`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def rfloordiv(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'integer division' of dataframe and other, element-wise (binary operator `//`). + """Get integer division of DataFrame and other, element-wise (binary operator `//`). Equivalent to ``other // dataframe``. With reverse version, `rfloordiv`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def mod(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'modulo' of dataframe and other, element-wise (binary operator `%`). + """Get modulo of DataFrame and other, element-wise (binary operator `%`). Equivalent to ``dataframe % other``. With reverse version, `rmod`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: other: Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") def rmod(self, other, axis: str | int = "columns") -> DataFrame: - """Get 'modulo' of dataframe and other, element-wise (binary operator `%`). + """Get modulo of DataFrame and other, element-wise (binary operator `%`). Equivalent to ``other % dataframe``. With reverse version, `mod`. Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + .. note:: + Mismatched indices will be unioned together. + Args: - other: + other (float, int, or Series): Any single or multiple element data structure, or list-like object. - axis: - ``{{0 or 'index', 1 or 'columns'}}``. Whether to compare by the - index (0 or 'index') or columns. (1 or 'columns'). For Series - input, axis to match Series index on. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. - - Notes: - Mismatched indices will be unioned together. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError("abstract method") @@ -719,20 +825,20 @@ def groupby( groups. Args: - by: + by (str, Sequence[str]): A label or list of labels may be passed to group by the columns in ``self``. Notice that a tuple is interpreted as a (single) key. - level : int, level name, or sequence of such, default None + level (int, level name, or sequence of such, default None): If the axis is a MultiIndex (hierarchical), group by a particular level or levels. Do not specify both ``by`` and ``level``. - as_index: + as_index (bool, default True): Default True. Return object with group labels as the index. Only relevant for DataFrame input. ``as_index=False`` is effectively "SQL-style" grouped output. This argument has no effect on filtrations such as ``head()``, ``tail()``, ``nth()`` and in transformations. - dropna: + dropna (bool, default True): Default True. If True, and if group keys contain NA values, NA values together with row/column will be dropped. If False, NA values will also be treated as the key in groups. @@ -748,31 +854,31 @@ def groupby( def map(self, func, na_action: Optional[str] = None) -> DataFrame: """Apply a function to a Dataframe elementwise. + This method applies a function that accepts and returns a scalar + to every element of a DataFrame. + .. note:: In pandas 2.1.0, DataFrame.applymap is deprecated and renamed to DataFrame.map. - This method applies a function that accepts and returns a scalar - to every element of a DataFrame. - Args: func: Python function wrapped by ``remote_function`` decorator, returns a single value from a single value. - na_action: + na_action (Optional[str], default None): ``{None, 'ignore'}``, default None. If ‘ignore’, propagate NaN values, without passing them to func. Returns: - Transformed DataFrame. + DataFrame: Transformed DataFrame. """ raise NotImplementedError("abstract method") # ---------------------------------------------------------------------- # Merging / joining methods - def join(self, other, *, how: str) -> DataFrame: + def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame: """Join columns of another DataFrame. Join columns with `other` DataFrame on index @@ -780,22 +886,20 @@ def join(self, other, *, how: str) -> DataFrame: Args: other: DataFrame with an Index similar to the Index of this one. - how: - ``{'left', 'right', 'outer', 'inner'}, default 'left'`` - + on: + Column in the caller to join on the index in other, otherwise + joins index-on-index. Like an Excel VLOOKUP operation. + how ({'left', 'right', 'outer', 'inner'}, default 'left'`): How to handle the operation of the two objects. - - * left: use calling frame's index (or column if on is specified) - * right: use `other`'s index. - * outer: form union of calling frame's index (or column if on is - specified) with `other`'s index, and sort it. - lexicographically. - * inner: form intersection of calling frame's index (or column if - on is specified) with `other`'s index, preserving the order - of the calling's one. + ``left``: use calling frame's index (or column if on is specified) + ``right``: use `other`'s index. ``outer``: form union of calling + frame's index (or column if on is specified) with `other`'s index, + and sort it lexicographically. ``inner``: form intersection of + calling frame's index (or column if on is specified) with `other`'s + index, preserving the order of the calling's one. Returns: - A dataframe containing columns from both the caller and `other`. + DataFrame: A dataframe containing columns from both the caller and `other`. """ raise NotImplementedError("abstract method") @@ -824,7 +928,6 @@ def merge( allowed. .. warning:: - If both key columns contain rows where the key is a null value, those rows will be matched against each other. This is different from usual SQL join behaviour and can lead to unexpected results. @@ -835,15 +938,14 @@ def merge( how: ``{'left', 'right', 'outer', 'inner'}, default 'inner'`` Type of merge to be performed. - - * left: use only keys from left frame, similar to a SQL left outer join; - preserve key order. - * right: use only keys from right frame, similar to a SQL right outer join; - preserve key order. - * outer: use union of keys from both frames, similar to a SQL full outer - join; sort keys lexicographically. - * inner: use intersection of keys from both frames, similar to a SQL inner - join; preserve the order of the left keys. + ``left``: use only keys from left frame, similar to a SQL left outer join; + preserve key order. + ``right``: use only keys from right frame, similar to a SQL right outer join; + preserve key order. + ``outer``: use union of keys from both frames, similar to a SQL full outer + join; sort keys lexicographically. + ``inner``: use intersection of keys from both frames, similar to a SQL inner + join; preserve the order of the left keys. on: Column join on. It must be found in both DataFrames. Either on or left_on + right_on @@ -867,7 +969,7 @@ def merge( no suffix. At least one of the values must not be None. Returns: - A DataFrame of the two merged objects. + DataFrame: A DataFrame of the two merged objects. """ raise NotImplementedError("abstract method") @@ -883,7 +985,7 @@ def any(self, *, bool_only: bool = False): non-empty). Args: - bool_only: + bool_only (bool. default False): Include only boolean columns. Returns: @@ -895,12 +997,12 @@ def all(self, *, bool_only: bool = False): """ Return whether all elements are True, potentially over an axis. - Returns True unless there at least one element within a series or - along a Dataframe axis that is False or equivalent (e.g. zero or + Returns True unless there at least one element within a Series or + along a DataFrame axis that is False or equivalent (e.g. zero or empty). Args: - bool_only: + bool_only (bool. default False): Include only boolean columns. Returns: @@ -913,7 +1015,7 @@ def prod(self, *, numeric_only: bool = False): Return the product of the values over the requested axis. Args: - numeric_only: + numeric_only (bool. default False): Include only float, int, boolean columns. Returns: @@ -928,7 +1030,7 @@ def min(self, *, numeric_only: bool = False): equivalent of the ``numpy.ndarray`` method ``argmin``. Args: - numeric_only: + numeric_only (bool, default False): Default False. Include only float, int, boolean columns. Returns: @@ -943,7 +1045,7 @@ def max(self, *, numeric_only: bool = False): the equivalent of the ``numpy.ndarray`` method ``argmax``. Args: - numeric_only: + numeric_only (bool. default False): Default False. Include only float, int, boolean columns. Returns: @@ -957,7 +1059,7 @@ def sum(self, *, numeric_only: bool = False): This is equivalent to the method ``numpy.sum``. Args: - numeric_only: + numeric_only (bool. default False): Default False. Include only float, int, boolean columns. Returns: @@ -969,7 +1071,7 @@ def mean(self, *, numeric_only: bool = False): """Return the mean of the values over the requested axis. Args: - numeric_only: + numeric_only (bool. default False): Default False. Include only float, int, boolean columns. Returns: @@ -977,13 +1079,28 @@ def mean(self, *, numeric_only: bool = False): """ raise NotImplementedError("abstract method") + def median(self, *, numeric_only: bool = False, exact: bool = False): + """Return the median of the values over the requested axis. + + Args: + numeric_only (bool. default False): + Default False. Include only float, int, boolean columns. + exact (bool. default False): + Default False. Get the exact median instead of an approximate + one. Note: ``exact=True`` not yet supported. + + Returns: + Series + """ + raise NotImplementedError("abstract method") + def var(self, *, numeric_only: bool = False): """Return unbiased variance over requested axis. Normalized by N-1 by default. Args: - numeric_only: + numeric_only (bool. default False): Default False. Include only float, int, boolean columns. Returns: @@ -997,7 +1114,7 @@ def std(self, *, numeric_only: bool = False): Normalized by N-1 by default. Args: - numeric_only: + numeric_only (bool. default False): Default False. Include only float, int, boolean columns. Returns: @@ -1013,7 +1130,7 @@ def count(self, *, numeric_only: bool = False): on `pandas.options.mode.use_inf_as_na`) are considered NA. Args: - numeric_only : bool, default False + numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. Returns: @@ -1039,7 +1156,7 @@ def cummin(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative minimum. Returns: - Return cumulative minimum of DataFrame. + DataFrame: Return cumulative minimum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1049,7 +1166,7 @@ def cummax(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative maximum. Returns: - Return cumulative maximum of DataFrame. + DataFrame: Return cumulative maximum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1059,7 +1176,7 @@ def cumsum(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative sum. Returns: - Return cumulative sum of DataFrame. + DataFrame: Return cumulative sum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1069,7 +1186,49 @@ def cumprod(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative product. Returns: - Return cumulative product of DataFrame. + DataFrame: Return cumulative product of DataFrame. + """ + raise NotImplementedError("abstract method") + + def agg(self, func): + """ + Aggregate using one or more operations over the specified axis. + + Args: + func (function): + Function to use for aggregating the data. + Accepted combinations are: string function name, list of + function names, e.g. ``['sum', 'mean']``. + + Returns: + Series or DataFrame: Aggregated results + """ + raise NotImplementedError("abstract method") + + def describe(self): + """ + Generate descriptive statistics. + + Descriptive statistics include those that summarize the central + tendency, dispersion and shape of a + dataset's distribution, excluding ``NaN`` values. + + Only supports numeric columns. + + .. note:: + Percentile values are approximates only. + + Returns: + Summary statistics of the Series or Dataframe provided. + + + Notes + ----- + For numeric data, the result's index will include ``count``, + ``mean``, ``std``, ``min``, ``max`` as well as lower, ``50`` and + upper percentiles. By default the lower percentile is ``25`` and the + upper percentile is ``75``. The ``50`` percentile is the + same as the median. """ raise NotImplementedError("abstract method") @@ -1107,18 +1266,18 @@ def value_counts( Return a Series containing counts of unique rows in the DataFrame. Args: - subset : label or list of labels, optional + subset (label or list of labels, optional): Columns to use when counting unique combinations. - normalize : bool, default False + normalize (bool, default False): Return proportions rather than frequencies. - sort : bool, default True + sort (bool, default True): Sort by frequencies. - ascending : bool, default False + ascending (bool, default False): Sort in ascending order. - dropna : bool, default True + dropna (bool, default True): Don’t include counts of rows that contain NA values. Returns: - Series + Series: Series containing counts of unique rows in the DataFrame """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index b6ee19b46ac..4843c971da2 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -19,7 +19,8 @@ class NDFrame(indexing.IndexingMixin): def ndim(self) -> int: """Return an int representing the number of axes / array dimensions. - Return 1 if Series. Otherwise return 2 if DataFrame. + Returns: + int: Return 1 if Series. Otherwise return 2 if DataFrame. """ raise NotImplementedError("abstract method") @@ -27,8 +28,9 @@ def ndim(self) -> int: def size(self) -> int: """Return an int representing the number of elements in this object. - Return the number of rows if Series. Otherwise return the number of - rows times number of columns if DataFrame. + Returns: + int: Return the number of rows if Series. Otherwise return the number of + rows times number of columns if DataFrame. """ raise NotImplementedError("abstract method") @@ -42,6 +44,7 @@ def abs(self): Returns: Series/DataFrame containing the absolute value of each element. + Returns a Series/DataFrame containing the absolute value of each element. """ raise NotImplementedError("abstract method") @@ -49,18 +52,18 @@ def astype(self, dtype): """ Cast a pandas object to a specified dtype ``dtype``. - Parameters - ---------- - dtype : str, data type, Series or Mapping of column name -> data type - Use a str, numpy.dtype, pandas.ExtensionDtype or Python type to - cast entire pandas object to the same type. Alternatively, use a - mapping, e.g. {col: dtype, ...}, where col is a column label and dtype is - a numpy.dtype or Python type to cast one or more of the DataFrame's - columns to column-specific types. + Args: + dtype (str or pandas.ExtensionDtype): + A dtype supported by BigQuery DataFrame include 'boolean','Float64','Int64', + 'string', 'tring[pyarrow]','timestamp[us, tz=UTC][pyarrow]', + 'timestamp[us][pyarrow]','date32[day][pyarrow]','time64[us][pyarrow]' + A pandas.ExtensionDtype include pandas.BooleanDtype(), pandas.Float64Dtype(), + pandas.Int64Dtype(), pandas.StringDtype(storage="pyarrow"), + pd.ArrowDtype(pa.date32()), pd.ArrowDtype(pa.time64("us")), + pd.ArrowDtype(pa.timestamp("us")), pd.ArrowDtype(pa.timestamp("us", tz="UTC")). - Returns - ------- - same type as caller + Returns: + same type as caller """ raise NotImplementedError("abstract method") @@ -70,17 +73,17 @@ def astype(self, dtype): @property def empty(self) -> bool: - """Indicator whether Series/DataFrame is empty. + """Indicates whether Series/DataFrame is empty. True if Series/DataFrame is entirely empty (no items), meaning any of the axes are of length 0. - Returns: - If Series/DataFrame is empty, return True, if not return False. - - Note: + .. note:: If Series/DataFrame contains only NA values, it is still not considered empty. + + Returns: + bool: If Series/DataFrame is empty, return True, if not return False. """ raise NotImplementedError("abstract method") @@ -97,28 +100,28 @@ def to_json( index: bool = True, lines: bool = False, ) -> str | None: - """Convert the object to a JSON string, written to GCS. + """Convert the object to a JSON string, written to Cloud Storage. Note NaN's and None will be converted to null and datetime objects will be converted to UNIX timestamps. + .. note:: + Only ``orient='records'`` and ``lines=True`` is supported so far. + Args: - path_or_buf: - A destination URI of GCS files(s) to store the extracted dataframe - in format of ``gs:///``. + path_or_buf (str): + A destination URI of Cloud Storage files(s) to store the extracted + dataframe in format of ``gs:///``. + Must contain a wildcard `*` character. If the data size is more than 1GB, you must use a wildcard to export the data into multiple files and the size of the files varies. None, file-like objects or local file paths not yet supported. - orient: + orient ({`split`, `records`, `index`, `columns`, `values`, `table`}, default 'columns): Indication of expected JSON string format. - .. note:: - - In BigQuery DataFrame, only `orient='records'` is supported so far. - * Series: - default is 'index' @@ -141,30 +144,25 @@ def to_json( - 'table' : dict like {{'schema': {{schema}}, 'data': {{data}}}} Describing the data, where data component is like ``orient='records'``. + index (bool, default True): + If True, write row names (index). - lines: + lines (bool, default False): If 'orient' is 'records' write out line-delimited json format. Will throw ValueError if incorrect 'orient' since others are not list-like. - .. note:: - - BigQuery DataFrames only supports ``lines=True`` so far. - - index: - If True, write row names (index). - Returns: - None. String output not yet supported. + None: String output not yet supported. """ raise NotImplementedError("abstract method") def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: - """Write object to a comma-separated values (csv) file on GCS. + """Write object to a comma-separated values (csv) file on Cloud Storage. Args: - path_or_buf: - A destination URI of GCS files(s) to store the extracted dataframe + path_or_buf (str): + A destination URI of Cloud Storage files(s) to store the extracted dataframe in format of ``gs:///``. If the data size is more than 1GB, you must use a wildcard to @@ -173,11 +171,11 @@ def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: None, file-like objects or local file paths not yet supported. - index: + index (bool, default True): If True, write row names (index). Returns: - None. String output not yet supported. + None: String output not yet supported. """ raise NotImplementedError("abstract method") @@ -208,11 +206,11 @@ def add_prefix(self, prefix: str, axis: int | str | None = None): For DataFrame, the column labels are prefixed. Args: - prefix: + prefix (str): The string to add before each label. - axis: + axis (int or str or None, default None): ``{{0 or 'index', 1 or 'columns', None}}``, default None. Axis - to add prefix on + to add prefix on. Returns: New Series or DataFrame with updated labels. @@ -250,7 +248,7 @@ def head(self, n: int = 5): If n is larger than the number of rows, this function returns all rows. Args: - n: + n (int, default 5): Default 5. Number of rows to select. Returns: @@ -271,7 +269,8 @@ def tail(self, n: int = 5): If n is larger than the number of rows, this function returns all rows. Args: - n: int, default 5. Number of rows to select. + n (int, default 5): + Number of rows to select. Returns: The last `n` rows of the caller object. @@ -290,12 +289,12 @@ def sample( You can use `random_state` for reproducibility. Args: - n: + n (Optional[int], default None): Number of items from axis to return. Cannot be used with `frac`. Default = 1 if `frac` = None. - frac: + frac (Optional[float], default None): Fraction of axis items to return. Cannot be used with `n`. - random_state: + random_state (Optional[int], default None): Seed for random number generator. Returns: @@ -360,7 +359,7 @@ def notna(self) -> NDFrame: NA values get mapped to False values. Returns: - Mask of bool values for each element that indicates whether an + NDFrame: Mask of bool values for each element that indicates whether an element is not an NA value. """ raise NotImplementedError("abstract method") @@ -376,11 +375,11 @@ def shift( Shifts the index without realigning the data. Args: - periods: + periods int: Number of periods to shift. Can be positive or negative. Returns: - Copy of input object, shifted. + NDFrame: Copy of input object, shifted. """ raise NotImplementedError("abstract method") @@ -398,34 +397,27 @@ def rank( By default, equal values are assigned a rank that is the average of the ranks of those values. - Parameters - ---------- - method : {'average', 'min', 'max', 'first', 'dense'}, default 'average' - How to rank the group of records that have the same value (i.e. ties): - - * average: average rank of the group - * min: lowest rank in the group - * max: highest rank in the group - * first: ranks assigned in order they appear in the array - * dense: like 'min', but rank always increases by 1 between groups. - - numeric_only : bool, default False - For DataFrame objects, rank only numeric columns if set to True. + Args: + method ({'average', 'min', 'max', 'first', 'dense'}, default 'average'): + How to rank the group of records that have the same value (i.e. ties): + `average`: average rank of the group, `min`: lowest rank in the group + max`: highest rank in the group, `first`: ranks assigned in order they + appear in the array, `dense`: like 'min', but rank always increases by + 1 between groups. - na_option : {'keep', 'top', 'bottom'}, default 'keep' - How to rank NaN values: + numeric_only (bool, default False): + For DataFrame objects, rank only numeric columns if set to True. - * keep: assign NaN rank to NaN values - * top: assign lowest rank to NaN values - * bottom: assign highest rank to NaN values + na_option ({'keep', 'top', 'bottom'}, default 'keep'): + How to rank NaN values: `keep`: assign NaN rank to NaN values, + , `top`: assign lowest rank to NaN values, `bottom`: assign highest + rank to NaN values. - ascending : bool, default True - Whether or not the elements should be ranked in ascending order. + ascending (bool, default True): + Whether or not the elements should be ranked in ascending order. - Returns - ------- - same type as caller - Return a Series or DataFrame with data ranks as values. + Returns: + same type as caller: Return a Series or DataFrame with data ranks as values. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index 5b5a9f206b3..6975e6edf2b 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -17,26 +17,23 @@ class GroupBy: def any(self): """ - Return True if any value in the group is truthful, else False. + Return True if any value in the group is true, else False. - - Returns - ------- - Series or DataFrame - DataFrame or Series of boolean values, where a value is True if any element - is True within its respective group, False otherwise. + Returns: + Series or DataFrame: DataFrame or Series of boolean values, + where a value is True if any element is True within its + respective group, False otherwise. """ raise NotImplementedError("abstract property") def all(self): """ - Return True if all values in the group are truthful, else False. + Return True if all values in the group are true, else False. - Returns - ------- - Series or DataFrame - DataFrame or Series of boolean values, where a value is True if all elements - are True within its respective group, False otherwise. + Returns: + Series or DataFrame: DataFrame or Series of boolean values, + where a value is True if all elements are True within its + respective group, False otherwise. """ raise NotImplementedError("abstract property") @@ -44,10 +41,8 @@ def count(self): """ Compute count of group, excluding missing values. - Returns - ------- - Series or DataFrame - Count of values within each group. + Returns: + Series or DataFrame: Count of values within each group. """ raise NotImplementedError("abstract property") @@ -58,14 +53,33 @@ def mean( """ Compute mean of groups, excluding missing values. - Parameters - ---------- - numeric_only : bool, default False - Include only float, int, boolean columns. + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. - Returns - ------- - pandas.Series or pandas.DataFrame + Returns: + pandas.Series or pandas.DataFrame: Mean of groups. + """ + raise NotImplementedError("abstract property") + + def median( + self, + numeric_only: bool = False, + *, + exact: bool = False, + ): + """ + Compute median of groups, excluding missing values. + + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + exact (bool, default False): + Calculate the exact median instead of an approximation. Note: + ``exact=True`` not yet supported. + + Returns: + pandas.Series or pandas.DataFrame: Median of groups. """ raise NotImplementedError("abstract property") @@ -79,15 +93,12 @@ def std( For multiple groupings, the result index will be a MultiIndex. - Parameters - ---------- - numeric_only : bool, default False - Include only `float`, `int` or `boolean` data. + Args: + numeric_only (bool, default False): + Include only `float`, `int` or `boolean` data. - Returns - ------- - Series or DataFrame - Standard deviation of values within each group. + Returns: + Series or DataFrame: Standard deviation of values within each group. """ raise NotImplementedError("abstract property") @@ -101,15 +112,13 @@ def var( For multiple groupings, the result index will be a MultiIndex. - Parameters - ---------- - numeric_only : bool, default False - Include only `float`, `int` or `boolean` data. + Args: + numeric_only (bool, default False): + Include only `float`, `int` or `boolean` data. - Returns - ------- - Series or DataFrame - Variance of values within each group. + Returns: + Series or DataFrame + Variance of values within each group. """ raise NotImplementedError("abstract property") @@ -121,17 +130,15 @@ def sum( """ Compute sum of group values. - Parameters - ---------- - numeric_only : bool, default False - Include only float, int, boolean columns. - min_count : int, default 0 - The required number of valid values to perform the operation. If fewer - than ``min_count`` non-NA values are present the result will be NA. - Returns - ------- - Series or DataFrame - Computed sum of values within each group. + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + min_count (int, default 0): + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + Returns: + Series or DataFrame: Computed sum of values within each group. """ raise NotImplementedError("abstract property") @@ -139,17 +146,15 @@ def prod(self, numeric_only: bool = False, min_count: int = 0): """ Compute prod of group values. - Parameters - ---------- - numeric_only : bool, default False - Include only float, int, boolean columns. - min_count : int, default 0 - The required number of valid values to perform the operation. If fewer - than ``min_count`` non-NA values are present the result will be NA. - Returns - ------- - Series or DataFrame - Computed prod of values within each group. + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + min_count (int, default 0): + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + Returns: + Series or DataFrame: Computed prod of values within each group. """ raise NotImplementedError("abstract property") @@ -161,17 +166,15 @@ def min( """ Compute min of group values. - Parameters - ---------- - numeric_only : bool, default False - Include only float, int, boolean columns. - min_count : int, default 0 - The required number of valid values to perform the operation. If fewer - than ``min_count`` non-NA values are present the result will be NA. - Returns - ------- - Series or DataFrame - Computed min of values within each group. + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + min_count (int, default 0): + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + Returns: + Series or DataFrame: Computed min of values within each group. """ raise NotImplementedError("abstract property") @@ -183,17 +186,15 @@ def max( """ Compute max of group values. - Parameters - ---------- - numeric_only : bool, default False - Include only float, int, boolean columns. - min_count : int, default 0 - The required number of valid values to perform the operation. If fewer - than ``min_count`` non-NA values are present the result will be NA. - Returns - ------- - Series or DataFrame - Computed max of values within each group. + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + min_count (int, default 0): + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + Returns: + Series or DataFrame: Computed max of values within each group. """ raise NotImplementedError("abstract property") @@ -201,15 +202,12 @@ def cumcount(self, ascending: bool = True): """ Number each item in each group from 0 to the length of that group - 1. - Parameters - ---------- - ascending : bool, default True - If False, number in reverse, from length of group - 1 to 0. + Args: + ascending (bool, default True): + If False, number in reverse, from length of group - 1 to 0. - Returns - ------- - Series - Sequence number of each element within each group. + Returns: + Series: Sequence number of each element within each group. """ raise NotImplementedError("abstract property") @@ -217,9 +215,8 @@ def cumprod(self, *args, **kwargs): """ Cumulative product for each group. - Returns - ------- - Series or DataFrame + Returns: + Series or DataFrame: Cumulative product for each group. """ raise NotImplementedError("abstract property") @@ -227,9 +224,8 @@ def cumsum(self, *args, **kwargs): """ Cumulative sum for each group. - Returns - ------- - Series or DataFrame + Returns: + Series or DataFrame: Cumulative sum for each group. """ raise NotImplementedError("abstract property") @@ -237,9 +233,8 @@ def cummin(self, *args, numeric_only: bool = False, **kwargs): """ Cumulative min for each group. - Returns - ------- - Series or DataFrame + Returns: + Series or DataFrame: Cumulative min for each group. """ raise NotImplementedError("abstract property") @@ -247,9 +242,8 @@ def cummax(self, *args, numeric_only: bool = False, **kwargs): """ Cumulative max for each group. - Returns - ------- - Series or DataFrame + Returns: + Series or DataFrame: Cumulative max for each group. """ raise NotImplementedError("abstract property") @@ -259,35 +253,30 @@ def diff(self): Calculates the difference of each element compared with another element in the group (default is element in previous row). - Returns - ------- - Series or DataFrame - First differences. + Returns: + Series or DataFrame: First differences. """ raise NotImplementedError("abstract property") def shift(self, periods: int = 1): """ Shift each group by periods observations. - If freq is passed, the index will be increased using the periods and the freq. - Parameters - ---------- - periods : int, default 1 - Number of periods to shift. - Returns - ------- - Series or DataFrame - Object shifted within each group. + Args: + periods (int, default 1): + Number of periods to shift. + + Returns: + Series or DataFrame: Object shifted within each group. """ raise NotImplementedError("abstract property") def rolling(self, *args, **kwargs): """ - Return a rolling grouper, providing rolling functionality per group. + Returns a rolling grouper, providing rolling functionality per group. Args: - min_periods : int, default None + min_periods (int, default None): Minimum number of observations in window required to have a value; otherwise, result is ``np.nan``. @@ -298,14 +287,16 @@ def rolling(self, *args, **kwargs): to the size of the window. Returns: - Return a new grouper with our rolling appended. + Series or DataFrame: Return a new grouper with our rolling appended. """ raise NotImplementedError("abstract property") def expanding(self, *args, **kwargs): """ - Return an expanding grouper, providing expanding - functionality per group. + Provides expanding functionality. + + Returns: + Series or DataFrame: A expanding grouper, providing expanding functionality per group. """ raise NotImplementedError("abstract property") diff --git a/third_party/bigframes_vendored/pandas/core/indexes/accessor.py b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py index 5d431a8a949..d59886e8aad 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py @@ -1,6 +1,6 @@ class DatetimeProperties: """ - Accessor object for datetimelike properties of the Series values. + Accessor object for datetime-like properties of the Series values. """ @property @@ -19,20 +19,19 @@ def dayofweek(self): values (using the `dt` accessor) or DatetimeIndex. Returns: - Series or Index - Containing integers indicating the day number. + Series or Index: Containing integers indicating the day number. """ raise NotImplementedError("abstract method") @property def date(self): - """Returns numpy array of python :class:`datetime.date` objects. + """Returns numpy array of Python :class:`datetime.date` objects. Namely, the date part of Timestamps without time and timezone information. - warning: + .. warning:: This method returns a Series whereas pandas returns a numpy array. """ @@ -69,7 +68,7 @@ def time(self): The time part of the Timestamps. - warning: + .. warning:: This method returns a Series whereas pandas returns a numpy array. """ @@ -80,7 +79,7 @@ def time(self): def quarter(self): """The quarter of the date. - warning: + .. warning:: This method returns a Series whereas pandas returns a numpy array. """ diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py index bdd05a01007..ebad5eb9181 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/base.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -9,7 +9,7 @@ class Index: @property def name(self): - """Return Index name.""" + """Returns Index name.""" raise NotImplementedError("abstract method") @property @@ -18,3 +18,19 @@ def shape(self): Return a tuple of the shape of the underlying data. """ raise NotImplementedError("abstract method") + + def to_numpy(self, dtype): + """ + A NumPy ndarray representing the values in this Series or Index. + + Args: + dtype: + The dtype to pass to :meth:`numpy.asarray`. + **kwargs: + Additional keywords passed through to the ``to_numpy`` method + of the underlying array (for extension arrays). + + Returns: + numpy.ndarray + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/indexing.py b/third_party/bigframes_vendored/pandas/core/indexing.py index 87fb34b32a3..d5b9f3c079f 100644 --- a/third_party/bigframes_vendored/pandas/core/indexing.py +++ b/third_party/bigframes_vendored/pandas/core/indexing.py @@ -21,7 +21,7 @@ def iloc(self): - A slice object with ints, e.g. ``1:7``. - **Not supported yet** A boolean array. - **Not supported yet** A ``callable`` function with one argument (the - calling Series or DataFrame) and that returns valid output for + calling Series or DataFrame) that returns valid output for indexing (one of the above). This is useful in method chains, when you don't have a reference to the calling object, but would like to base your selection on some value. @@ -43,30 +43,24 @@ def loc(self): Allowed inputs are: - - **Not supported yet** A single label, e.g. ``5`` or ``'a'``, (note + - A single label, e.g. ``5`` or ``'a'``, (note that ``5`` is interpreted as a *label* of the index, and **never** as an integer position along the index). - - **Not supported yet** A list or array of labels, e.g. ``['a', 'b', 'c']``. - - **Not supported yet** A slice object with labels, e.g. ``'a':'f'``. - - .. warning:: Note that contrary to usual python slices, **both** the - start and the stop are included - - - **Not supported yet** A boolean array of the same length as the axis being sliced, + - A list of labels, e.g. ``['a', 'b', 'c']``. + - A boolean series of the same length as the axis being sliced, e.g. ``[True, False, True]``. - - An alignable boolean Series. The index of the key will be aligned before - masking. - - **Not supported yet** An alignable Index. The Index of the returned + - An alignable Index. The index of the returned selection will be the input. + - **Not supported yet** An alignable boolean Series. The index of the key will be aligned before + masking. + - **Not supported yet** A slice object with labels, e.g. ``'a':'f'``. + Note: contrary to usual python slices, **both** the start and the stop are included. - **Not supported yet** A ``callable`` function with one argument (the - calling Series or DataFrame) - and that returns valid output for indexing (one of the above) + calling Series or DataFrame) that returns valid output for indexing + (one of the above). Raises: - KeyError: If any items are not found. - IndexingError: - If an indexed key is passed and its index is unalignable to the - frame index. + NotImplementError: if the inputs are not supported. """ raise NotImplementedError("abstract methdod") diff --git a/third_party/bigframes_vendored/pandas/core/reshape/__init__.py b/third_party/bigframes_vendored/pandas/core/reshape/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/reshape/concat.py b/third_party/bigframes_vendored/pandas/core/reshape/concat.py index 98b43c7dfba..6a5a9fdde92 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/concat.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/concat.py @@ -8,6 +8,7 @@ def concat( objs, *, + axis=0, join: str = "outer", ignore_index: bool = False, ): @@ -25,6 +26,8 @@ def concat( objs: Objects to concatenate. Any None objects will be dropped silently unless they are all None in which case a ValueError will be raised. + axis : {0/'index', 1/'columns'}, default 0 + The axis to concatenate along. join: {'inner', 'outer'}, default 'outer' How to handle indexes on other axis (or axes). ignore_index : bool, default False diff --git a/third_party/bigframes_vendored/pandas/core/reshape/tile.py b/third_party/bigframes_vendored/pandas/core/reshape/tile.py new file mode 100644 index 00000000000..9381ad4552a --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/reshape/tile.py @@ -0,0 +1,65 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/tile.py +""" +Quantilization functions and related stuff +""" +from __future__ import annotations + + +def cut( + x, + bins, + *, + labels=None, +): + """ + Bin values into discrete intervals. + + Use `cut` when you need to segment and sort data values into bins. This + function is also useful for going from a continuous variable to a + categorical variable. For example, `cut` could convert ages to groups of + age ranges. Supports binning into an equal number of bins, or a + pre-specified array of bins. + + ``labels=False`` implies you just want the bins back. + + Examples: + + .. code-block:: + + import bigframes.pandas as pd + + pd.options.display.progress_bar = None + s = pd.Series([0, 1, 1, 2]) + pd.cut(s, bins=4, labels=False) + + 0 0 + 1 1 + 2 1 + 3 3 + dtype: Int64 + + Args: + x (Series): + The input Series to be binned. Must be 1-dimensional. + bins (int): + The criteria to bin by. + + int : Defines the number of equal-width bins in the range of `x`. The + range of `x` is extended by .1% on each side to include the minimum + and maximum values of `x`. + labels (None): + Specifies the labels for the returned bins. Must be the same length as + the resulting bins. If False, returns only integer indicators of the + bins. This affects the type of the output container (see below). + If True, raises an error. When `ordered=False`, labels must be + provided. + + Returns: + Series: A Series representing the respective bin for each value + of `x`. The type depends on the value of `labels`. + sequence of scalars : returns a Series for Series `x` or a + Categorical for all other inputs. The values stored within + are whatever the type in the sequence is. + False : returns an ndarray of integers. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index a254fcf842d..215f7ec4e08 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -21,7 +21,7 @@ class Series(NDFrame): # type: ignore[misc] @property def dt(self): """ - Accessor object for datetimelike properties of the Series values. + Accessor object for datetime-like properties of the Series values. """ raise NotImplementedError("abstract property") @@ -58,10 +58,9 @@ def name(self) -> Hashable: to form a DataFrame. It is also used whenever displaying the Series using the interpreter. - Returns - ------- - label (hashable object) - The name of the Series, also the column name if part of a DataFrame. + Returns: + hashable object: The name of the Series, also the column name + if part of a DataFrame. """ raise NotImplementedError("abstract property") @@ -78,24 +77,21 @@ def reset_index( when the index is meaningless and needs to be reset to the default before another operation. - Parameters - ---------- - drop : bool, default False - Just reset the index, without inserting it as a column in - the new DataFrame. - name : object, optional - The name to use for the column containing the original Series - values. Uses ``self.name`` by default. This argument is ignored - when `drop` is True. + Args: + drop (bool, default False): + Just reset the index, without inserting it as a column in + the new DataFrame. + name (object, optional): + The name to use for the column containing the original Series + values. Uses ``self.name`` by default. This argument is ignored + when `drop` is True. - Returns - ------- - Series or DataFrame or None - When `drop` is False (the default), a DataFrame is returned. - The newly created columns will come first in the DataFrame, - followed by the original Series values. - When `drop` is True, a `Series` is returned. - In either case, if ``inplace=True``, no value is returned. + Returns: + Series or DataFrame or None; When `drop` is False (the default), + a DataFrame is returned. The newly created columns will come first + in the DataFrame, followed by the original Series values. + When `drop` is True, a `Series` is returned. + In either case, if ``inplace=True``, no value is returned. """ raise NotImplementedError("abstract method") @@ -125,36 +121,34 @@ def to_string( """ Render a string representation of the Series. - Parameters - ---------- - buf : StringIO-like, optional - Buffer to write to. - na_rep : str, optional - String representation of NaN to use, default 'NaN'. - float_format : one-parameter function, optional - Formatter function to apply to columns' elements if they are - floats, default None. - header : bool, default True - Add the Series header (index name). - index : bool, optional - Add index (row) labels, default True. - length : bool, default False - Add the Series length. - dtype : bool, default False - Add the Series dtype. - name : bool, default False - Add the Series name if not None. - max_rows : int, optional - Maximum number of rows to show before truncating. If None, show - all. - min_rows : int, optional - The number of rows to display in a truncated repr (when number - of rows is above `max_rows`). + Args: + buf (StringIO-like, optional): + Buffer to write to. + na_rep (str, optional): + String representation of NaN to use, default 'NaN'. + float_format (one-parameter function, optional): + Formatter function to apply to columns' elements if they are + floats, default None. + header (bool, default True): + Add the Series header (index name). + index (bool, optional): + Add index (row) labels, default True. + length (bool, default False): + Add the Series length. + dtype (bool, default False): + Add the Series dtype. + name (bool, default False): + Add the Series name if not None. + max_rows (int, optional): + Maximum number of rows to show before truncating. If None, show + all. + min_rows (int, optional): + The number of rows to display in a truncated repr (when number + of rows is above `max_rows`). - Returns - ------- - str or None - String representation of Series if ``buf=None``, otherwise None. + Returns: + str or None: String representation of Series if ``buf=None``, + otherwise None. """ formatter = fmt.SeriesFormatter( self, @@ -183,18 +177,16 @@ def to_markdown( """ Print {klass} in Markdown-friendly format. - Parameters - ---------- - buf : str, Path or StringIO-like, optional, default None - Buffer to write to. If None, the output is returned as a string. - mode : str, optional - Mode in which file is opened, "wt" by default. - index : bool, optional, default True - Add index (row) labels. - Returns - ------- - str - {klass} in Markdown-friendly format. + Args: + buf (str, Path or StringIO-like, optional, default None): + Buffer to write to. If None, the output is returned as a string. + mode (str, optional): + Mode in which file is opened, "wt" by default. + index (bool, optional, default True): + Add index (row) labels. + + Returns: + str: {klass} in Markdown-friendly format. """ raise NotImplementedError("abstract method") @@ -202,18 +194,15 @@ def to_dict(self, into: type[dict] = dict) -> Mapping: """ Convert Series to {label -> value} dict or dict-like object. - Parameters - ---------- - into : class, default dict - The collections.abc.Mapping subclass to use as the return - object. Can be the actual class or an empty - instance of the mapping type you want. If you want a - collections.defaultdict, you must pass it initialized. + Args: + into (class, default dict): + The collections.abc.Mapping subclass to use as the return + object. Can be the actual class or an empty + instance of the mapping type you want. If you want a + collections.defaultdict, you must pass it initialized. - Returns - ------- - collections.abc.Mapping - Key-value representation of Series. + Returns: + collections.abc.Mapping: Key-value representation of Series. """ raise NotImplementedError("abstract method") @@ -221,10 +210,8 @@ def to_frame(self) -> DataFrame: """ Convert Series to DataFrame. - Returns - ------- - DataFrame - DataFrame representation of Series. + Returns: + DataFrame: DataFrame representation of Series. """ raise NotImplementedError("abstract method") @@ -242,12 +229,11 @@ def to_excel(self, excel_writer, sheet_name): Note that creating an `ExcelWriter` object with a file name that already exists will result in the contents of the existing file being erased. - Parameters - ---------- - excel_writer : path-like, file-like, or ExcelWriter object - File path or existing ExcelWriter. - sheet_name : str, default 'Sheet1' - Name of sheet which will contain DataFrame. + Args: + excel_writer (path-like, file-like, or ExcelWriter object): + File path or existing ExcelWriter. + sheet_name (str, default 'Sheet1'): + Name of sheet to contain Series. """ raise NotImplementedError("abstract method") @@ -255,23 +241,20 @@ def to_latex(self, buf=None, columns=None, header=True, index=True, **kwargs): """ Render object to a LaTeX tabular, longtable, or nested table. - Parameters - ---------- - buf : str, Path or StringIO-like, optional, default None - Buffer to write to. If None, the output is returned as a string. - columns : list of label, optional - The subset of columns to write. Writes all columns by default. - header : bool or list of str, default True - Write out the column names. If a list of strings is given, - it is assumed to be aliases for the column names. - index : bool, default True - Write row names (index). - + Args: + buf (str, Path or StringIO-like, optional, default None): + Buffer to write to. If None, the output is returned as a string. + columns (list of label, optional): + The subset of columns to write. Writes all columns by default. + header (bool or list of str, default True): + Write out the column names. If a list of strings is given, + it is assumed to be aliases for the column names. + index (bool, default True): + Write row names (index). - Returns - ------- - str or None - If buf is None, returns the result as a string. Otherwise returns None. + Returns: + str or None: If buf is None, returns the result as a string. + Otherwise returns None. """ raise NotImplementedError("abstract method") @@ -281,11 +264,10 @@ def tolist(self) -> list: These are each a scalar type, which is a Python scalar (for str, int, float) or a pandas scalar - (for Timestamp/Timedelta/Interval/Period) + (for Timestamp/Timedelta/Interval/Period). - Returns - ------- - list + Returns: + list: list of the values """ raise NotImplementedError("abstract method") @@ -295,25 +277,24 @@ def to_numpy(self, dtype, copy=False, na_value=None): """ A NumPy ndarray representing the values in this Series or Index. - Parameters - ---------- - dtype : str or numpy.dtype, optional - The dtype to pass to :meth:`numpy.asarray`. - copy : bool, default False - Whether to ensure that the returned value is not a view on - another array. Note that ``copy=False`` does not *ensure* that - ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that - a copy is made, even if not strictly necessary. - na_value : Any, optional - The value to use for missing values. The default value depends - on `dtype` and the type of the array. - **kwargs - Additional keywords passed through to the ``to_numpy`` method - of the underlying array (for extension arrays). + Args: + dtype (str or numpy.dtype, optional): + The dtype to pass to :meth:`numpy.asarray`. + copy (bool, default False): + Whether to ensure that the returned value is not a view on + another array. Note that ``copy=False`` does not *ensure* that + ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that + a copy is made, even if not strictly necessary. + na_value (Any, optional): + The value to use for missing values. The default value depends + on `dtype` and the type of the array. + ``**kwargs``: + Additional keywords passed through to the ``to_numpy`` method + of the underlying array (for extension arrays). - Returns - ------- - numpy.ndarray + Returns: + numpy.ndarray: A NumPy ndarray representing the values in this + Series or Index. """ raise NotImplementedError("abstract method") @@ -321,12 +302,11 @@ def to_pickle(self, path, **kwargs): """ Pickle (serialize) object to file. - Parameters - ---------- - path : str, path object, or file-like object - String, path object (implementing ``os.PathLike[str]``), or file-like - object implementing a binary ``write()`` function. File path where - the pickled object will be stored. + Args: + path (str, path object, or file-like object): + String, path object (implementing ``os.PathLike[str]``), or file-like + object implementing a binary ``write()`` function. File path where + the pickled object will be stored. """ raise NotImplementedError("abstract method") @@ -334,11 +314,10 @@ def to_xarray(self): """ Return an xarray object from the pandas object. - Returns - ------- - xarray.DataArray or xarray.Dataset - Data in the pandas structure converted to Dataset if the object is - a DataFrame, or a DataArray if the object is a Series. + Returns: + xarray.DataArray or xarray.Dataset: Data in the pandas structure + converted to Dataset if the object is a DataFrame, or a DataArray if + the object is a Series. """ raise NotImplementedError("abstract method") @@ -356,43 +335,24 @@ def to_json( Note NaN's and None will be converted to null and datetime objects will be converted to UNIX timestamps. - Parameters - ---------- - path_or_buf : str, path object, file-like object, or None, default None - String, path object (implementing os.PathLike[str]), or file-like - object implementing a write() function. If None, the result is - returned as a string. - orient: - Indication of expected JSON string format. - - * Series: - - - default is 'index' - - allowed values are: {{'split', 'records', 'index', 'table'}}. - - * DataFrame: - - - default is 'columns' - - allowed values are: {{'split', 'records', 'index', 'columns', - 'values', 'table'}}. - - * The format of the JSON string: - - - 'split' : dict like {{'index' -> [index], 'columns' -> [columns], - 'data' -> [values]}} - - 'records' : list like [{{column -> value}}, ... , {{column -> value}}] - - 'index' : dict like {{index -> {{column -> value}}}} - - 'columns' : dict like {{column -> {{index -> value}}}} - - 'values' : just the values array - - 'table' : dict like {{'schema': {{schema}}, 'data': {{data}}}} - + Args: + path_or_buf (str, path object, file-like object, or None, default None): + String, path object (implementing os.PathLike[str]), or file-like + object implementing a write() function. If None, the result is + returned as a string. + orient ({"split", "records", "index", "columns", "values", "table"}, default "columns"): + Indication of expected JSON string format. + 'split' : dict like {{'index' -> [index], 'columns' -> [columns],'data' -> [values]}} + 'records' : list like [{{column -> value}}, ... , {{column -> value}}] + 'index' : dict like {{index -> {{column -> value}}}} + 'columns' : dict like {{column -> {{index -> value}}}} + 'values' : just the values array + 'table' : dict like {{'schema': {{schema}}, 'data': {{data}}}} Describing the data, where data component is like ``orient='records'``. - Returns - ------- - None or str - If path_or_buf is None, returns the resulting json format as a - string. Otherwise returns None. + Returns: + None or str: If path_or_buf is None, returns the resulting json format as a + string. Otherwise returns None. """ raise NotImplementedError("abstract method") @@ -400,20 +360,32 @@ def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: """ Write object to a comma-separated values (csv) file. - Parameters - ---------- - path_or_buf : str, path object, file-like object, or None, default None - String, path object (implementing os.PathLike[str]), or file-like - object implementing a write() function. If None, the result is - returned as a string. If a non-binary file object is passed, it should - be opened with `newline=''`, disabling universal newlines. If a binary - file object is passed, `mode` might need to contain a `'b'`. + Args: + path_or_buf (str, path object, file-like object, or None, default None): + String, path object (implementing os.PathLike[str]), or file-like + object implementing a write() function. If None, the result is + returned as a string. If a non-binary file object is passed, it should + be opened with `newline=''`, disabling universal newlines. If a binary + file object is passed, `mode` might need to contain a `'b'`. - Returns - ------- - None or str - If path_or_buf is None, returns the resulting csv format as a - string. Otherwise returns None. + Returns: + None or str: If path_or_buf is None, returns the resulting csv format + as a string. Otherwise returns None. + """ + raise NotImplementedError("abstract method") + + def agg(self, func): + """ + Aggregate using one or more operations over the specified axis. + + Args: + func (function): + Function to use for aggregating the data. + Accepted combinations are: string function name, list of + function names, e.g. ``['sum', 'mean']``. + + Returns: + scalar or Series: Aggregated results """ raise NotImplementedError("abstract method") @@ -421,10 +393,9 @@ def count(self): """ Return number of non-NA/null observations in the Series. - Returns - ------- - int or Series (if level specified) - Number of non-null values in the Series. + Returns: + int or Series (if level specified): Number of non-null values in the + Series. """ raise NotImplementedError("abstract method") @@ -434,9 +405,8 @@ def nunique(self) -> int: Excludes NA values by default. - Returns - ------- - int + Returns: + int: number of unique elements in the object. """ raise NotImplementedError("abstract method") @@ -448,10 +418,8 @@ def mode(self) -> Series: Always returns Series even if only one value is returned. - Returns - ------- - Series - Modes of the Series in sorted order. + Returns: + Series: Modes of the Series in sorted order. """ raise NotImplementedError("abstract method") @@ -464,15 +432,15 @@ def drop_duplicates( Return Series with duplicate values removed. Args: - keep : {'first', 'last', ``False``}, default 'first' + keep ({'first', 'last', ``False``}, default 'first'): Method to handle dropping duplicates: - - 'first' : Drop duplicates except for the first occurrence. - - 'last' : Drop duplicates except for the last occurrence. - - ``False`` : Drop all duplicates. + 'first' : Drop duplicates except for the first occurrence. + 'last' : Drop duplicates except for the last occurrence. + ``False`` : Drop all duplicates. Returns: - Series with duplicates dropped or None if ``inplace=True``. + Series: Series with duplicates dropped or None if ``inplace=True``. """ raise NotImplementedError("abstract method") @@ -485,17 +453,17 @@ def duplicated(self, keep="first") -> Series: last occurrence of duplicates can be indicated. Args: - keep : {'first', 'last', False}, default 'first' + keep ({'first', 'last', False}, default 'first'): Method to handle dropping duplicates: - - 'first' : Mark duplicates as ``True`` except for the first - occurrence. - - 'last' : Mark duplicates as ``True`` except for the last - occurrence. - - ``False`` : Mark all duplicates as ``True``. + 'first' : Mark duplicates as ``True`` except for the first + occurrence. + 'last' : Mark duplicates as ``True`` except for the last + occurrence. + ``False`` : Mark all duplicates as ``True``. Returns: - Series indicating whether each value has occurred in the + Series: Series indicating whether each value has occurred in the preceding values. """ raise NotImplementedError("abstract method") @@ -504,16 +472,13 @@ def round(self, decimals: int = 0) -> Series: """ Round each value in a Series to the given number of decimals. - Parameters - ---------- - decimals : int, default 0 - Number of decimal places to round to. If decimals is negative, - it specifies the number of positions to the left of the decimal point. + Args: + decimals (int, default 0): + Number of decimal places to round to. If decimals is negative, + it specifies the number of positions to the left of the decimal point. - Returns - ------- - Series - Rounded values of the Series. + Returns: + Series: Rounded values of the Series. """ raise NotImplementedError("abstract method") @@ -524,10 +489,13 @@ def diff(self) -> Series: Calculates the difference of a {klass} element compared with another element in the {klass} (default is element in previous row). - Returns - ------- - {klass} - First differences of the Series. + Args: + periods (int, default 1): + Periods to shift for calculating difference, accepts negative + values. + + Returns: + {klass}: First differences of the Series. """ raise NotImplementedError("abstract method") @@ -541,23 +509,22 @@ def dot(self, other) -> Series | np.ndarray: It can also be called using `self @ other` in Python >= 3.5. - Parameters - ---------- - other : Series, DataFrame or array-like - The other object to compute the dot product with its columns. + .. note:: + The Series and other has to share the same index if other is a Series + or a DataFrame. + BigQuery Dataframes does not validate this property and will produce + incorrect results if indices are not equal. - Returns - ------- - scalar, Series or numpy.ndarray - Return the dot product of the Series and other if other is a - Series, the Series of the dot product of Series and each rows of - other if other is a DataFrame or a numpy.ndarray between the Series - and each columns of the numpy array. + Args: + other (Series): + The other object to compute the dot product with its columns. + + Returns: + scalar, Series or numpy.ndarray: Return the dot product of the Series + and other if other is a Series, the Series of the dot product of + Series and each rows of other if other is a DataFrame or a + numpy.ndarray between the Series and each columns of the numpy array. - Notes - ----- - The Series and other has to share the same index if other is a Series - or a DataFrame. """ raise NotImplementedError("abstract method") @@ -579,6 +546,7 @@ def sort_values( *, axis: Axis = 0, ascending: bool | int | Sequence[bool] | Sequence[int] = True, + kind: str = "quicksort", na_position: str = "last", ) -> Series | None: """ @@ -587,20 +555,21 @@ def sort_values( Sort a Series in ascending or descending order by some criterion. - Parameters - ---------- - axis : {0 or 'index'} - Unused. Parameter needed for compatibility with DataFrame. - ascending : bool or list of bools, default True - If True, sort values in ascending order, otherwise descending. - na_position : {'first' or 'last'}, default 'last' - Argument 'first' puts NaNs at the beginning, 'last' puts NaNs at - the end. + Args: + axis (0 or 'index'): + Unused. Parameter needed for compatibility with DataFrame. + ascending (bool or list of bools, default True): + If True, sort values in ascending order, otherwise descending. + kind (str, default to 'quicksort'): + Choice of sorting algorithm. Accepts 'quicksort’, ‘mergesort’, + ‘heapsort’, ‘stable’. Ignored except when determining whether to + sort stably. 'mergesort' or 'stable' will result in stable reorder + na_position ({'first' or 'last'}, default 'last'): + Argument 'first' puts NaNs at the beginning, 'last' puts NaNs at + the end. - Returns - ------- - Series or None - Series ordered by values or None if ``inplace=True``. + Returns: + Series or None: Series ordered by values or None if ``inplace=True``. """ raise NotImplementedError("abstract method") @@ -617,21 +586,19 @@ def sort_index( Returns a new Series sorted by label if `inplace` argument is ``False``, otherwise updates the original series and returns None. - Parameters - ---------- - axis : {0 or 'index'} - Unused. Parameter needed for compatibility with DataFrame. - ascending : bool or list-like of bools, default True - Sort ascending vs. descending. When the index is a MultiIndex the - sort direction can be controlled for each level individually. - na_position : {'first', 'last'}, default 'last' - If 'first' puts NaNs at the beginning, 'last' puts NaNs at the end. - Not implemented for MultiIndex. + Args: + axis ({0 or 'index'}): + Unused. Parameter needed for compatibility with DataFrame. + ascending (bool or list-like of bools, default True): + Sort ascending vs. descending. When the index is a MultiIndex the + sort direction can be controlled for each level individually. + na_position ({'first', 'last'}, default 'last'): + If 'first' puts NaNs at the beginning, 'last' puts NaNs at the end. + Not implemented for MultiIndex. - Returns - ------- - Series or None - The original Series sorted by the labels or None if ``inplace=True``. + Returns: + Series or None: The original Series sorted by the labels or None if + ``inplace=True``. """ @@ -643,25 +610,21 @@ def nlargest( """ Return the largest `n` elements. - Parameters - ---------- - n : int, default 5 - Return this many descending sorted values. - keep : {'first', 'last', 'all'}, default 'first' - When there are duplicate values that cannot all fit in a - Series of `n` elements: - - - ``first`` : return the first `n` occurrences in order - of appearance. - - ``last`` : return the last `n` occurrences in reverse - order of appearance. - - ``all`` : keep all occurrences. This can result in a Series of - size larger than `n`. + Args: + n (int, default 5): + Return this many descending sorted values. + keep ({'first', 'last', 'all'}, default 'first'): + When there are duplicate values that cannot all fit in a + Series of `n` elements: + ``first`` : return the first `n` occurrences in order + of appearance. + ``last`` : return the last `n` occurrences in reverse + order of appearance. + ``all`` : keep all occurrences. This can result in a Series of + size larger than `n`. - Returns - ------- - Series - The `n` largest values in the Series, sorted in decreasing order. + Returns: + Series: The `n` largest values in the Series, sorted in decreasing order. """ raise NotImplementedError("abstract method") @@ -669,25 +632,22 @@ def nsmallest(self, n: int = 5, keep: str = "first") -> Series: """ Return the smallest `n` elements. - Parameters - ---------- - n : int, default 5 - Return this many ascending sorted values. - keep : {'first', 'last', 'all'}, default 'first' - When there are duplicate values that cannot all fit in a - Series of `n` elements: - - - ``first`` : return the first `n` occurrences in order - of appearance. - - ``last`` : return the last `n` occurrences in reverse - order of appearance. - - ``all`` : keep all occurrences. This can result in a Series of - size larger than `n`. + Args: + n (int, default 5): + Return this many ascending sorted values. + keep ({'first', 'last', 'all'}, default 'first'): + When there are duplicate values that cannot all fit in a + Series of `n` elements: + + ``first`` : return the first `n` occurrences in order + of appearance. + ``last`` : return the last `n` occurrences in reverse + order of appearance. + ``all`` : keep all occurrences. This can result in a Series of + size larger than `n`. - Returns - ------- - Series - The `n` smallest values in the Series, sorted in increasing order. + Returns: + Series: The `n` smallest values in the Series, sorted in increasing order. """ raise NotImplementedError("abstract method") @@ -704,14 +664,13 @@ def apply( Can be ufunc (a NumPy function that applies to the entire Series) or a Python function that only works on single values. - Parameters - ---------- - func : function - Python function or NumPy ufunc to apply. - Returns - ------- - Series or DataFrame - If func returns a Series object the result will be a DataFrame. + Args: + func (function): + Python function or NumPy ufunc to apply. + + Returns: + Series or DataFrame: If func returns a Series object the result + will be a DataFrame. """ raise NotImplementedError("abstract method") @@ -731,49 +690,45 @@ def groupby( used to group large amounts of data and compute operations on these groups. - Parameters - ---------- - by : mapping, function, label, pd.Grouper or list of such - Used to determine the groups for the groupby. - If ``by`` is a function, it's called on each value of the object's - index. If a dict or Series is passed, the Series or dict VALUES - will be used to determine the groups (the Series' values are first - aligned; see ``.align()`` method). If a list or ndarray of length - equal to the selected axis is passed (see the `groupby user guide - `_), - the values are used as-is to determine the groups. A label or list - of labels may be passed to group by the columns in ``self``. - Notice that a tuple is interpreted as a (single) key. - axis : {0 or 'index', 1 or 'columns'}, default 0 - Split along rows (0) or columns (1). For `Series` this parameter - is unused and defaults to 0. - level : int, level name, or sequence of such, default None - If the axis is a MultiIndex (hierarchical), group by a particular - level or levels. Do not specify both ``by`` and ``level``. - as_index : bool, default True - Return object with group labels as the - index. Only relevant for DataFrame input. as_index=False is - effectively "SQL-style" grouped output. This argument has no effect - on filtrations (see the `filtrations in the user guide - `_), - such as ``head()``, ``tail()``, ``nth()`` and in transformations - (see the `transformations in the user guide - `_). - dropna : bool, default True - If True, and if group keys contain NA values, NA values together - with row/column will be dropped. - If False, NA values will also be treated as the key in groups. + Args: + by (mapping, function, label, pd.Grouper or list of such, default None): + Used to determine the groups for the groupby. + If ``by`` is a function, it's called on each value of the object's + index. If a dict or Series is passed, the Series or dict VALUES + will be used to determine the groups (the Series' values are first + aligned; see ``.align()`` method). If a list or ndarray of length + equal to the selected axis is passed (see the `groupby user guide + `_), + the values are used as-is to determine the groups. A label or list + of labels may be passed to group by the columns in ``self``. + Notice that a tuple is interpreted as a (single) key. + axis ({0 or 'index', 1 or 'columns'}, default 0): + Split along rows (0) or columns (1). For `Series` this parameter + is unused and defaults to 0. + level (int, level name, or sequence of such, default None): + If the axis is a MultiIndex (hierarchical), group by a particular + level or levels. Do not specify both ``by`` and ``level``. + as_index (bool, default True): + Return object with group labels as the + index. Only relevant for DataFrame input. as_index=False is + effectively "SQL-style" grouped output. This argument has no effect + on filtrations (see the "filtrations in the user guide" + ``_), + such as ``head()``, ``tail()``, ``nth()`` and in transformations + (see the "transformations in the user guide" + ``_). + dropna : bool, default True + If True, and if group keys contain NA values, NA values together + with row/column will be dropped. + If False, NA values will also be treated as the key in groups. - Returns - ------- - SeriesGroupBy - Returns a groupby object that contains information about the groups. + Returns: + SeriesGroupBy: Returns a groupby object that contains information about the groups. """ raise NotImplementedError("abstract method") def drop( - self, - labels=None, + self, labels=None, *, axis=0, index=None, columns=None, level=None ) -> Series | None: """ Return Series with specified index labels removed. @@ -782,10 +737,18 @@ def drop( When using a multi-index, labels on different levels can be removed by specifying the level. - Parameters - ---------- - labels : single label or list-like - Index labels to drop. + Args: + labels (single label or list-like): + Index labels to drop. + axis: + Unused. Parameter needed for compatibility with DataFrame. + index: + Redundant for application on Series, but 'index' can be used instead + of 'labels'. + columns: + No change is made to the Series; use 'index' or 'labels' instead. + level: + For MultiIndex, level for which the labels will be removed. Returns ------- @@ -806,7 +769,7 @@ def reorder_levels(self, order: Sequence) -> Series: May not drop or duplicate levels. Args: - order: list of int representing new level order + order (list of int representing new level order): Reference level by number or key. Returns: @@ -819,7 +782,7 @@ def droplevel(self, level): Return Series with requested index / column level(s) removed. Args: - level: int, str, or list-like + level (int, str, or list-like): If a string is given, must be the name of a level If list-like, elements must be names or positional indexes of levels. @@ -833,6 +796,16 @@ def fillna( self, value=None, ) -> Series | None: + """ + Fill NA/NaN values using the specified method. + + Args: + value (scalar, dict, Series, or DataFrame, default None): + Value to use to fill holes (e.g. 0). + + Returns: + Series or None: Object with missing values filled or None. + """ raise NotImplementedError("abstract method") def between( @@ -848,24 +821,33 @@ def between( corresponding Series element is between the boundary values `left` and `right`. NA values are treated as `False`. - Parameters - ---------- - left : scalar or list-like - Left boundary. - right : scalar or list-like - Right boundary. - inclusive : {"both", "neither", "left", "right"} - Include boundaries. Whether to set each bound as closed or open. + Args: + left (scalar or list-like): + Left boundary. + right (scalar or list-like): + Right boundary. + inclusive ({"both", "neither", "left", "right"}): + Include boundaries. Whether to set each bound as closed or open. - Returns - ------- - Series - Series representing whether each element is between left and + Returns: + Series: Series representing whether each element is between left and right (inclusive). """ raise NotImplementedError("abstract method") + def cumprod(self): + """ + Return cumulative product over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative + product. + + Returns: + Return cumulative sum of scalar or Series. + """ + raise NotImplementedError("abstract method") + def cumsum(self): """ Return cumulative sum over a DataFrame or Series axis. @@ -873,16 +855,13 @@ def cumsum(self): Returns a DataFrame or Series of the same size containing the cumulative sum. - Parameters - ---------- - axis : {{0 or 'index', 1 or 'columns'}}, default 0 - The index or the name of the axis. 0 is equivalent to None or 'index'. - For `Series` this parameter is unused and defaults to 0. + Args: + axis ({0 or 'index', 1 or 'columns'}, default 0): + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. - Returns - ------- - scalar or Series - Return cumulative sum of scalar or Series. + Returns: + scalar or Series: Return cumulative sum of scalar or Series. """ raise NotImplementedError("abstract method") @@ -893,16 +872,13 @@ def cummax(self): Returns a DataFrame or Series of the same size containing the cumulative maximum. - Parameters - ---------- - axis : {{0 or 'index', 1 or 'columns'}}, default 0 - The index or the name of the axis. 0 is equivalent to None or 'index'. - For `Series` this parameter is unused and defaults to 0. + Args: + axis ({{0 or 'index', 1 or 'columns'}}, default 0): + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. - Returns - ------- - scalar or Series - Return cumulative maximum of scalar or Series. + Returns: + scalar or Series: Return cumulative maximum of scalar or Series. """ raise NotImplementedError("abstract method") @@ -913,57 +889,54 @@ def cummin(self): Returns a DataFrame or Series of the same size containing the cumulative minimum. - Parameters - ---------- - axis : {{0 or 'index', 1 or 'columns'}}, default 0 - The index or the name of the axis. 0 is equivalent to None or 'index'. - For `Series` this parameter is unused and defaults to 0. - skipna : bool, default True - Exclude NA/null values. If an entire row/column is NA, the result - will be NA. - *args, **kwargs - Additional keywords have no effect but might be accepted for - compatibility with NumPy. + Args: + axis ({0 or 'index', 1 or 'columns'}, default 0): + The index or the name of the axis. 0 is equivalent to None or 'index'. + For `Series` this parameter is unused and defaults to 0. + skipna (bool, default True): + Exclude NA/null values. If an entire row/column is NA, the result + will be NA. + `*args`, `**kwargs`: + Additional keywords have no effect but might be accepted for + compatibility with NumPy. - Returns - ------- - scalar or Series - Return cumulative minimum of scalar or Series. + Returns: + scalar or Series: Return cumulative minimum of scalar or Series. """ raise NotImplementedError("abstract method") def eq(self, other) -> Series: - """Return Equal of series and other, element-wise (binary operator eq). + """Return equal of Series and other, element-wise (binary operator eq). Equivalent to ``other == series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def ne(self, other) -> Series: - """Return Not equal of series and other, element-wise (binary operator ne). + """Return not equal of Series and other, element-wise (binary operator ne). Equivalent to ``other != series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the comparison. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def le(self, other) -> Series: - """Get 'less than or equal to' of series and other, element-wise (binary operator `<=`). + """Get 'less than or equal to' of Series and other, element-wise (binary operator `<=`). Equivalent to ``series <= other``, but with support to substitute a fill_value for missing data in either one of the inputs. @@ -978,61 +951,76 @@ def le(self, other) -> Series: raise NotImplementedError("abstract method") def lt(self, other) -> Series: - """Get 'less than' of series and other, element-wise (binary operator `<`). + """Get 'less than' of Series and other, element-wise (binary operator `<`). - Equivalent to ``series < other``, but with support to substitute a fill_value for - missing data in either one of the inputs. + Equivalent to ``series < other``, but with support to substitute a fill_value for + missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): - Returns: - Series. The result of the comparison. + Returns: + Series: The result of the operation. """ raise NotImplementedError("abstract method") def ge(self, other) -> Series: - """Get 'greater than or equal to' of series and other, element-wise (binary operator `>=`). + """Get 'greater than or equal to' of Series and other, element-wise (binary operator `>=`). Equivalent to ``series >= other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the comparison. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def gt(self, other) -> Series: - """Get 'less than or equal to' of series and other, element-wise (binary operator `<=`). + """Get 'less than or equal to' of Series and other, element-wise (binary operator `<=`). Equivalent to ``series <= other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def add(self, other) -> Series: - """Return Addition of series and other, element-wise (binary operator add). + """Return addition of Series and other, element-wise (binary operator add). Equivalent to ``series + other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): + + Returns: + Series: The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def radd(self, other) -> Series: + """Return addition of Series and other, element-wise (binary operator radd). + + Equivalent to ``other + series``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1041,136 +1029,180 @@ def sub( self, other, ) -> Series: - """Return Subtraction of series and other, element-wise (binary operator sub). + """Return subtraction of Series and other, element-wise (binary operator sub). Equivalent to ``series - other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def rsub(self, other) -> Series: - """Return Subtraction of series and other, element-wise (binary operator rsub). + """Return subtraction of Series and other, element-wise (binary operator rsub). Equivalent to ``other - series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def mul(self, other) -> Series: - """Return Multiplication of series and other, element-wise (binary operator mul). + """Return multiplication of Series and other, element-wise (binary operator mul). Equivalent to ``other * series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") + def rmul(self, other) -> Series: + """Return multiplication of Series and other, element-wise (binary operator mul). + + Equivalent to ``series * others``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other (Series, or scalar value): + + Returns: + Series: The result of the operation. + """ + raise NotImplementedError("abstract method") + def truediv(self, other) -> Series: - """Return Floating division of series and other, element-wise (binary operator truediv). + """Return floating division of Series and other, element-wise (binary operator truediv). Equivalent to ``series / other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def rtruediv(self, other) -> Series: - """Return Floating division of series and other, element-wise (binary operator rtruediv). + """Return floating division of Series and other, element-wise (binary operator rtruediv). Equivalent to ``other / series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def floordiv(self, other) -> Series: - """Return Integer division of series and other, element-wise (binary operator floordiv). + """Return integer division of Series and other, element-wise (binary operator floordiv). Equivalent to ``series // other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def rfloordiv(self, other) -> Series: - """Return Integer division of series and other, element-wise (binary operator rfloordiv). + """Return integer division of Series and other, element-wise (binary operator rfloordiv). Equivalent to ``other // series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def mod(self, other) -> Series: - """Return Modulo of series and other, element-wise (binary operator mod). + """Return modulo of Series and other, element-wise (binary operator mod). Equivalent to ``series % other``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: - other: Series, or scalar value + other (Series, or scalar value): Returns: - Series. The result of the operation. + Series: The result of the operation. """ raise NotImplementedError("abstract method") def rmod(self, other) -> Series: - """Get Modulo of series and other, element-wise (binary operator `rmod`). + """Get modulo of Series and other, element-wise (binary operator `rmod`). Equivalent to ``other % series``, but with support to substitute a fill_value for missing data in either one of the inputs. + Args: + other (Series, or scalar value): + + Returns: + Series: The result of the operation. + + """ + raise NotImplementedError("abstract method") + + def divmod(self, other) -> Series: + """Return integer division and modulo of Series and other, element-wise (binary operator divmod). + + Equivalent to divmod(series, other). + Args: other: Series, or scalar value Returns: - Series. The result of the operation. + 2-Tuple of Series. The result of the operation. The result is always + consistent with (floordiv, mod) (though pandas may not). + + """ + raise NotImplementedError("abstract method") + + def rdivmod(self, other) -> Series: + """Return integer division and modulo of Series and other, element-wise (binary operator rdivmod). + + Equivalent to other divmod series. + + Args: + other: Series, or scalar value + + Returns: + 2-Tuple of Series. The result of the operation. The result is always + consistent with (rfloordiv, rmod) (though pandas may not). """ raise NotImplementedError("abstract method") @@ -1181,15 +1213,12 @@ def all( """ Return whether all elements are True, potentially over an axis. - Returns True unless there at least one element within a series or along a - Dataframe axis that is False or equivalent (e.g. zero or empty). + Returns True unless there at least one element within a Series or along a + DataFrame axis that is False or equivalent (e.g. zero or empty). - - Returns - ------- - scalar or Series - If level is specified, then, Series is returned; otherwise, scalar - is returned. + Returns: + scalar or Series: If level is specified, then, Series is returned; + otherwise, scalar is returned. """ raise NotImplementedError("abstract method") @@ -1202,12 +1231,9 @@ def any( Returns False unless there is at least one element within a series or along a Dataframe axis that is True or equivalent (e.g. non-zero or non-empty). - - Returns - ------- - scalar or Series - If level is specified, then, Series is returned; otherwise, scalar - is returned. + Returns: + scalar or Series: If level is specified, then, Series is returned; + otherwise, scalar is returned. """ raise NotImplementedError("abstract method") @@ -1221,9 +1247,8 @@ def max( of the ``numpy.ndarray`` method ``argmax``. - Returns - ------- - scalar or scalar + Returns: + scalar or scalar """ raise NotImplementedError("abstract method") @@ -1236,10 +1261,8 @@ def min( If you want the index of the minimum, use ``idxmin``. This is the equivalent of the ``numpy.ndarray`` method ``argmin``. - - Returns - ------- - scalar or scalar + Returns: + scalar or scalar """ raise NotImplementedError("abstract method") @@ -1249,7 +1272,7 @@ def std( """ Return sample standard deviation over requested axis. - Normalized by N-1 by default. This can be changed using the ddof argument. + Normalized by N-1 by default. Returns @@ -1264,12 +1287,10 @@ def var( """ Return unbiased variance over requested axis. - Normalized by N-1 by default. This can be changed using the ddof argument. + Normalized by N-1 by default. - - Returns - ------- - scalar or Series (if level specified) + Returns: + scalar or Series (if level specified) """ raise NotImplementedError("abstract method") @@ -1291,6 +1312,19 @@ def mean(self): """ raise NotImplementedError("abstract method") + def median(self, *, exact: bool = False): + """Return the median of the values over the requested axis. + + Args: + exact (bool. default False): + Default False. Get the exact median instead of an approximate + one. Note: ``exact=True`` not yet supported. + + Returns: + scalar + """ + raise NotImplementedError("abstract method") + def prod(self): """Return the product of the values over the requested axis. @@ -1299,66 +1333,71 @@ def prod(self): """ raise NotImplementedError("abstract method") + def skew(self): + """Return unbiased skew over requested axis. + + Normalized by N-1. + + Returns: + scalar + """ + raise NotImplementedError("abstract method") + def kurt(self): """Return unbiased kurtosis over requested axis. Kurtosis obtained using Fisher’s definition of kurtosis (kurtosis of normal == 0.0). Normalized by N-1. - Returns - ------- - scalar or scalar + Returns: + scalar or scalar: Unbiased kurtosis over requested axis. """ raise NotImplementedError("abstract method") def where(self, cond, other): """Replace values where the condition is False. - Parameters - ---------- - cond: bool Series/DataFrame, array-like, or callable - Where cond is True, keep the original value. Where False, replace - with corresponding value from other. If cond is callable, it is - computed on the Series/DataFrame and should return boolean - Series/DataFrame or array. The callable must not change input - Series/DataFrame (though pandas doesn’t check it). - other: scalar, Series/DataFrame, or callable - Entries where cond is False are replaced with corresponding value - from other. If other is callable, it is computed on the - Series/DataFrame and should return scalar or Series/DataFrame. - The callable must not change input Series/DataFrame (though pandas - doesn’t check it). If not specified, entries will be filled with - the corresponding NULL value (np.nan for numpy dtypes, pd.NA for - extension dtypes). + Args: + cond (bool Series/DataFrame, array-like, or callable): + Where cond is True, keep the original value. Where False, replace + with corresponding value from other. If cond is callable, it is + computed on the Series/DataFrame and returns boolean + Series/DataFrame or array. The callable must not change input + Series/DataFrame (though pandas doesn’t check it). + other (scalar, Series/DataFrame, or callable): + Entries where cond is False are replaced with corresponding value + from other. If other is callable, it is computed on the + Series/DataFrame and returns scalar or Series/DataFrame. + The callable must not change input Series/DataFrame (though pandas + doesn’t check it). If not specified, entries will be filled with + the corresponding NULL value (np.nan for numpy dtypes, pd.NA for + extension dtypes). - Returns - ------- - Series + Returns: + Series """ raise NotImplementedError("abstract method") def mask(self, cond, other): """Replace values where the condition is True. - Parameters - ---------- - cond: bool Series/DataFrame, array-like, or callable - Where cond is False, keep the original value. Where True, replace - with corresponding value from other. If cond is callable, it is - computed on the Series/DataFrame and should return boolean - Series/DataFrame or array. The callable must not change input - Series/DataFrame (though pandas doesn’t check it). - other: scalar, Series/DataFrame, or callable - Entries where cond is True are replaced with corresponding value - from other. If other is callable, it is computed on the - Series/DataFrame and should return scalar or Series/DataFrame. - The callable must not change input Series/DataFrame (though pandas - doesn’t check it). If not specified, entries will be filled with - the corresponding NULL value (np.nan for numpy dtypes, pd.NA for - extension dtypes). + Args: + cond (bool Series/DataFrame, array-like, or callable): + Where cond is False, keep the original value. Where True, replace + with corresponding value from other. If cond is callable, it is + computed on the Series/DataFrame and should return boolean + Series/DataFrame or array. The callable must not change input + Series/DataFrame (though pandas doesn’t check it). + other (scalar, Series/DataFrame, or callable): + Entries where cond is True are replaced with corresponding value + from other. If other is callable, it is computed on the + Series/DataFrame and should return scalar or Series/DataFrame. + The callable must not change input Series/DataFrame (though pandas + doesn’t check it). If not specified, entries will be filled with + the corresponding NULL value (np.nan for numpy dtypes, pd.NA for + extension dtypes). - Returns - ------- - Series + Returns: + Series """ raise NotImplementedError("abstract method") @@ -1369,18 +1408,15 @@ def clip(self): singular values or array like, and in the latter case the clipping is performed element-wise in the specified axis. - Parameters - ---------- - - lower: float or array-like, default None - Minimum threshold value. All values below this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. + Args: + lower (float or array-like, default None): + Minimum threshold value. All values below this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. - upper: float or array-like, default None - Maximum threshold value. All values above this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. + upper (float or array-like, default None): + Maximum threshold value. All values above this threshold will be set to it. A missing threshold (e.g NA) will not clip the value. - Returns - ------- - Series + Returns: + Series. """ raise NotImplementedError("abstract method") @@ -1390,10 +1426,8 @@ def argmax(self): If the minimum is achieved in multiple locations, the first row position is returned. - Returns - ------- - Series - Row position of the maximum value. + Returns: + Series: Row position of the maximum value. """ raise NotImplementedError("abstract method") @@ -1403,10 +1437,8 @@ def argmin(self): If the maximum is achieved in multiple locations, the first row position is returned. - Returns - ------- - Series - Row position of the minimum value. + Returns: + Series: Row position of the minimum value. """ raise NotImplementedError("abstract method") @@ -1420,18 +1452,15 @@ def rename(self, index, **kwargs) -> Series | None: Alternatively, change ``Series.name`` with a scalar value. - Parameters - ---------- - index : scalar, hashable sequence, dict-like or function optional - Functions or dict-like are transformations to apply to - the index. - Scalar or hashable sequence-like will alter the ``Series.name`` - attribute. + Args: + index (scalar, hashable sequence, dict-like or function optional): + Functions or dict-like are transformations to apply to + the index. + Scalar or hashable sequence-like will alter the ``Series.name`` + attribute. - Returns - ------- - Series - Series with index labels + Returns: + Series: Series with index labels """ raise NotImplementedError("abstract method") @@ -1440,14 +1469,12 @@ def rename_axis(self, mapper, **kwargs): """ Set the name of the axis for the index or columns. - Parameters - ---------- - mapper : scalar, list-like, optional - Value to set the axis name attribute. + Args: + mapper (scalar, list-like, optional): + Value to set the axis name attribute. - Returns - ------- - Series + Returns: + Series: Series with the name of the axis set. """ raise NotImplementedError("abstract method") @@ -1459,39 +1486,35 @@ def rolling( """ Provide rolling window calculations. - Parameters - ---------- - window : int, timedelta, str, offset, or BaseIndexer subclass - Size of the moving window. + Args: + window (int, timedelta, str, offset, or BaseIndexer subclass): + Size of the moving window. - If an integer, the fixed number of observations used for - each window. + If an integer, the fixed number of observations used for + each window. - If a timedelta, str, or offset, the time period of each window. Each - window will be a variable sized based on the observations included in - the time-period. This is only valid for datetimelike indexes. - To learn more about the offsets & frequency strings, please see `this link - `__. + If a timedelta, str, or offset, the time period of each window. Each + window will be a variable sized based on the observations included in + the time-period. This is only valid for datetime-like indexes. + To learn more about the offsets & frequency strings, please see `this link + `__. - If a BaseIndexer subclass, the window boundaries - based on the defined ``get_window_bounds`` method. Additional rolling - keyword arguments, namely ``min_periods``, ``center``, ``closed`` and - ``step`` will be passed to ``get_window_bounds``. + If a BaseIndexer subclass, the window boundaries + based on the defined ``get_window_bounds`` method. Additional rolling + keyword arguments, namely ``min_periods``, ``center``, ``closed`` and + ``step`` will be passed to ``get_window_bounds``. - min_periods : int, default None - Minimum number of observations in window required to have a value; - otherwise, result is ``np.nan``. + min_periods (int, default None): + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. - For a window that is specified by an offset, ``min_periods`` will default to 1. + For a window that is specified by an offset, ``min_periods`` will default to 1. - For a window that is specified by an integer, ``min_periods`` will default - to the size of the window. + For a window that is specified by an integer, ``min_periods`` will default + to the size of the window. - Returns - ------- - ``Window`` subclass if a ``win_type`` is passed - - ``Rolling`` subclass if ``win_type`` is not passed + Returns: + ``Window`` subclass if a ``win_type`` is passed.``Rolling`` subclass if ``win_type`` is not passed """ raise NotImplementedError("abstract method") @@ -1499,14 +1522,12 @@ def expanding(self, min_periods=1): """ Provide expanding window calculations. - Parameters - ---------- - min_periods : int, default 1 - Minimum number of observations in window required to have a value; - otherwise, result is ``np.nan``. + Args: + min_periods (int, default 1): + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. - Returns - ------- + Returns: ``Expanding`` subclass """ raise NotImplementedError("abstract method") @@ -1527,18 +1548,18 @@ def value_counts( Excludes NA values by default. Args: - normalize : bool, default False + normalize (bool, default False): If True then the object returned will contain the relative frequencies of the unique values. - sort : bool, default True + sort (bool, default True): Sort by frequencies. - ascending : bool, default False + ascending (bool, default False): Sort in ascending order. - dropna : bool, default True + dropna (bool, default True): Don't include counts of NaN. Returns: - Series + Series: Series containing counts of unique values. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/strings/accessor.py b/third_party/bigframes_vendored/pandas/core/strings/accessor.py index 81b9bd3d146..e464843c779 100644 --- a/third_party/bigframes_vendored/pandas/core/strings/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/strings/accessor.py @@ -1,3 +1,7 @@ +import re +import typing + + class StringMethods: """ Vectorized string functions for Series and Index. @@ -7,6 +11,29 @@ class StringMethods: R's stringr package. """ + def extract(self, pat: str, flags: int = 0): + """ + Extract capture groups in the regex `pat` as columns in a DataFrame. + + For each subject string in the Series, extract groups from the + first match of regular expression `pat`. + + Args: + pat: + Regular expression pattern with capturing groups. + flags: + Flags from the ``re`` module, e.g. ``re.IGNORECASE``, that + modify regular expression matching for things like case, + spaces, etc. For more details, see :mod:`re`. + + Returns: + A DataFrame with one row for each subject string, and one + column for each group. Any capture group names in regular + expression pat will be used for column names; otherwise + capture group numbers will be used. + """ + raise NotImplementedError("abstract method") + def find(self, sub, start: int = 0, end=None): """Return lowest indexes in each strings in the Series/Index. @@ -17,15 +44,14 @@ def find(self, sub, start: int = 0, end=None): Args: sub: Substring being searched. - start: + start (int, default 0): Left edge index. - end: + end (None): Right edge index. Returns: - Series or Index of int. + bigframes.series.Series: Series with lowest indexes in each strings. """ - raise NotImplementedError("abstract method") def len(self): @@ -35,9 +61,8 @@ def len(self): (such as a dictionary). Returns: - Series or Index of int - A Series or Index of integer values indicating the length of each - element in the Series or Index. + bigframes.series.Series: A Series or Index of integer values indicating + the length of each element in the Series or Index. """ raise NotImplementedError("abstract method") @@ -48,7 +73,7 @@ def lower(self): Equivalent to :meth:`str.lower`. Returns: - Series or Index of object + bigframes.series.Series: Series with lowercase. """ raise NotImplementedError("abstract method") @@ -57,16 +82,16 @@ def slice(self, start=None, stop=None): """Slice substrings from each element in the Series or Index. Args: - start : int, optional + start (int, optional): Start position for slice operation. - stop : int, optional + stop (int, optional): Stop position for slice operation. - step : int, optional + step (int, optional): Step size for slice operation. Returns: - Series or Index of object - Series or Index from sliced substring from original string object. + bigframes.series.Series:: Series or Index from sliced + substring from original string object. """ raise NotImplementedError("abstract method") @@ -80,7 +105,8 @@ def strip(self): Equivalent to :meth:`str.strip`. Returns: - Series or Index of object + bigframes.series.Series: Series or Index without leading + and trailing characters. """ raise NotImplementedError("abstract method") @@ -91,7 +117,7 @@ def upper(self): Equivalent to :meth:`str.upper`. Returns: - Series or Index of object + bigframes.series.Series: Series with uppercase strings. """ raise NotImplementedError("abstract method") @@ -104,9 +130,8 @@ def isnumeric(self): has zero characters, ``False`` is returned for that check. Returns: - Series or Index of bool - Series or Index of boolean values with the same length as the original - Series/Index. + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. """ raise NotImplementedError("abstract method") @@ -120,7 +145,7 @@ def rstrip(self): Equivalent to :meth:`str.rstrip`. Returns: - Series or Index of object + bigframes.series.Series: Series without trailing characters. """ raise NotImplementedError("abstract method") @@ -134,7 +159,7 @@ def lstrip(self): Equivalent to :meth:`str.lstrip`. Returns: - Series or Index of object` + bigframes.series.Series: Series without leading characters. """ raise NotImplementedError("abstract method") @@ -147,9 +172,8 @@ def repeat(self, repeats: int): Same value for all (int) or different value per (sequence). Returns: - Series or pandas.Index - Series or Index of repeated string objects specified by - input parameter repeats. + bigframes.series.Series: Series or Index of repeated string + objects specified by input parameter repeats. """ raise NotImplementedError("abstract method") @@ -160,7 +184,7 @@ def capitalize(self): Equivalent to :meth:`str.capitalize`. Returns: - Series or Index of object + bigframes.series.Series: Series with captitalized strings. """ raise NotImplementedError("abstract method") @@ -172,16 +196,121 @@ def cat(self, others, *, join): and elements of `others` element-wise. Args: - others : Series + others (Series): - join : {'left', 'outer'}, default 'left' + join ({'left', 'outer'}, default 'left'): Determines the join-style between the calling Series and any Series in `others` (objects without an index need to match the length of the calling Series). To disable alignment, use `.values` on any Series/Index/DataFrame in `others`. Returns: - Series + bigframes.series.Series: Series with concatenated strings. """ raise NotImplementedError("abstract method") + + def contains(self, pat, case: bool = True, flags: int = 0, *, regex: bool = True): + """ + Test if pattern or regex is contained within a string of a Series or Index. + + Return boolean Series or Index based on whether a given pattern or regex is + contained within a string of a Series or Index. + + Args: + pat (str, re.Pattern): + Character sequence or regular expression. + case (bool, default True): + If True, case sensitive. + flags (int, default 0): + Flags to pass through to the re module, e.g. re.IGNORECASE. + regex (bool, default True): + If True, assumes the pat is a regular expression. + If False, treats the pat as a literal string. + + Returns: + bigframes.series.Series: A Series or Index of boolean values indicating + whether the given pattern is contained within the string of each + element of the Series or Index. + """ + raise NotImplementedError("abstract method") + + def replace( + self, + pat: typing.Union[str, re.Pattern], + repl: str, + *, + case: typing.Optional[bool] = None, + flags: int = 0, + regex: bool = False, + ): + """ + Replace each occurrence of pattern/regex in the Series/Index. + + Equivalent to :meth:`str.replace` or :func:`re.sub`, depending on + the regex value. + + Args: + pat (str, re.Pattern): + String can be a character sequence or regular expression. + repl (str): + Replacement string. + case (default None): + Determines if replace is case sensitive: + + - If True, case sensitive (the default if `pat` is a string) + - Set to False for case insensitive + - Cannot be set if `pat` is a compiled regex. + flags (int, default 0): + Regex module flags, e.g. re.IGNORECASE. Cannot be set if `pat` is a compiled + regex. + regex (bool: default False): + Determines if the passed-in pattern is a regular expression: + + - If True, assumes the passed-in pattern is a regular expression. + - If False, treats the pattern as a literal string + - Cannot be set to False if `pat` is a compiled regex or `repl` is + a callable. + + Returns: + bigframes.series.Series: A copy of the object with all matching occurrences + of `pat` replaced by `repl`. + + """ + raise NotImplementedError("abstract method") + + def startswith( + self, + pat: typing.Union[str, tuple[str, ...]], + ): + """ + Test if the start of each string element matches a pattern. + + Args: + pat (str, tuple[str, ...]): + Character sequence or tuple of strings. Regular expressions are not + accepted. + + Returns: + bigframes.series.Series: A Series of booleans indicating whether the given + pattern matches the start of each string element. + """ + raise NotImplementedError("abstract method") + + def endswith( + self, + pat: typing.Union[str, tuple[str, ...]], + ): + """ + Test if the end of each string element matches a pattern. + + Args: + pat (str, tuple[str, ...]): + Character sequence or tuple of strings. Regular expressions are not + accepted. + + Returns: + bigframes.series.Series: A Series of booleans indicating whether the given + pattern matches the end of each string element. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/common.py b/third_party/bigframes_vendored/pandas/io/common.py index 057ed96721d..506984e64da 100644 --- a/third_party/bigframes_vendored/pandas/io/common.py +++ b/third_party/bigframes_vendored/pandas/io/common.py @@ -17,9 +17,10 @@ def dedup_names( but a custom pattern may be supported in the future. Examples - -------- - >>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False) + ``` + dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False) ['x', 'y', 'x.1', 'x.2'] + ``` """ names = list(names) # so we can index counts: DefaultDict[Hashable, int] = defaultdict(int) diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 4f4b1b8199f..9425ead0e34 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -15,23 +15,23 @@ def read_gbq( col_order: Iterable[str] = (), max_results: Optional[int] = None, ): - """Loads DataFrame from Google BigQuery. + """Loads DataFrame from BigQuery. Args: - query: + query (str): A SQL string to be executed or a BigQuery table to be read. The table must be specified in the format of `project.dataset.tablename` or `dataset.tablename`. - index_col: + index_col (Iterable[str] or str): Name of result column(s) to use for index in results DataFrame. - col_order: + col_order (Iterable[str]): List of BigQuery column names in the desired order for results DataFrame. - max_results: + max_results (Optional[int], default None): If set, limit the maximum number of rows to fetch from the query results. Returns: - A DataFrame representing results of the query or table. + bigframes.dataframe.DataFrame: A DataFrame representing results of the query or table. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/parquet.py b/third_party/bigframes_vendored/pandas/io/parquet.py index ee2e9a65f22..6f0a2b3cb46 100644 --- a/third_party/bigframes_vendored/pandas/io/parquet.py +++ b/third_party/bigframes_vendored/pandas/io/parquet.py @@ -8,15 +8,18 @@ def read_parquet( self, path: str, ): - r"""Load a parquet object from the file path (local or GCS), returning a DataFrame. + r"""Load a Parquet object from the file path (local or Cloud Storage), returning a DataFrame. - Args: - path: - Local or GCS path to parquet file. - - Note: + .. note:: This method will not guarantee the same ordering as the file. Instead, set a serialized index column as the index and sort by that in the resulting DataFrame. + + Args: + path (str): + Local or Cloud Storage path to Parquet file. + + Returns: + bigframes.dataframe.DataFrame: A BigQuery DataFrames. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/parsers/readers.py b/third_party/bigframes_vendored/pandas/io/parsers/readers.py index e5ed39b54aa..e01eb734fb8 100644 --- a/third_party/bigframes_vendored/pandas/io/parsers/readers.py +++ b/third_party/bigframes_vendored/pandas/io/parsers/readers.py @@ -32,27 +32,30 @@ def read_csv( encoding: Optional[str] = None, **kwargs, ): - r"""Loads DataFrame from comma-separated values (csv) file locally or from GCS. + """Loads DataFrame from comma-separated values (csv) file locally or from + Cloud Storage. The CSV file data will be persisted as a temporary BigQuery table, which can be automatically recycled after the Session is closed. - Note: using `engine="bigquery"` will not guarantee the same ordering as the - file. Instead, set a serialized index column as the index and sort by - that in the resulting DataFrame. + .. note:: + using `engine="bigquery"` will not guarantee the same ordering as the + file. Instead, set a serialized index column as the index and sort by + that in the resulting DataFrame. Args: - filepath_or_buffer: a string path including GCS and local file. - - sep: the separator for fields in a CSV file. For the BigQuery engine, the separator + filepath_or_buffer (str): + a string path including Cloud Storage and local file. + sep (Optional[str], default ","): + the separator for fields in a CSV file. For the BigQuery engine, the separator can be any ISO-8859-1 single-byte character. To use a character in the range 128-255, you must encode the character as UTF-8. Both engines support `sep="\t"` to specify tab character as separator. Default engine supports having any number of spaces as separator by specifying `sep="\s+"`. Separators longer than 1 character are interpreted as regular expressions by the default engine. BigQuery engine only supports single character separators. - - header: row number to use as the column names. + header (Optional[int], default 0): + row number to use as the column names. - ``None``: Instructs autodetect that there are no headers and data should be read starting from the first row. - ``0``: If using `engine="bigquery"`, Autodetect tries to detect headers in the @@ -68,19 +71,19 @@ def read_csv( contains column names unless the `names` argument is specified. If `names` is provided, row N+1 will be ignored, row N+2 will be read as data, and column names are inferred from `names`. - - names: a list of column names to use. If the file contains a header row and you + names (default None): + a list of column names to use. If the file contains a header row and you want to pass this parameter, then `header=0` should be passed as well so the first (header) row is ignored. Only to be used with default engine. - - index_col: column(s) to use as the row labels of the DataFrame, either given as + index_col (default None): + column(s) to use as the row labels of the DataFrame, either given as string name or column index. `index_col=False` can be used with the default engine only to enforce that the first column is not used as the index. Using column index instead of column name is only supported with the default engine. The BigQuery engine only supports having a single column name as the `index_col`. Neither engine supports having a multi-column index. - - usecols: list of column names to use. The BigQuery engine only supports having a list + usecols (default None): + List of column names to use): The BigQuery engine only supports having a list of string column names. Column indices and callable functions are only supported with the default engine. Using the default engine, the column names in `usecols` can be defined to correspond to column names provided with the `names` parameter @@ -89,22 +92,22 @@ def read_csv( column names provided with the BigQuery engine will be consistent in the resulting dataframe. If using a callable function with the default engine, only column names that evaluate to True by the callable function will be in the resulting dataframe. - - dtype: data type for data or columns. Only to be used with default engine. - - engine: type of engine to use. If `engine="bigquery"` is specified, then BigQuery's - load API will be used. Otherwise, the engine will be passed to `pandas.read_csv`. - - encoding: the character encoding of the data. The default encoding is `UTF-8` for both + dtype (data type for data or columns): + Data type for data or columns. Only to be used with default engine. + engine (Optional[Dict], default None): + Type of engine to use. If `engine="bigquery"` is specified, then BigQuery's load API will be used. + Otherwise, the engine will be passed to `pandas.read_csv`. + encoding (Optional[str], default to None): + encoding the character encoding of the data. The default encoding is `UTF-8` for both engines. The default engine acceps a wide range of encodings. Refer to Python documentation for a comprehensive list, https://docs.python.org/3/library/codecs.html#standard-encodings The BigQuery engine only supports `UTF-8` and `ISO-8859-1`. - - **kwargs: keyword arguments. + **kwargs: + keyword arguments. Returns: - A BigQuery DataFrames. + bigframes.dataframe.DataFrame: A BigQuery DataFrames. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py index 925bb76e1a1..03958f75950 100644 --- a/third_party/bigframes_vendored/sklearn/base.py +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -13,7 +13,7 @@ class BaseEstimator: """Base class for all estimators. - Notes: + .. note:: All estimators should specify all the parameters that can be set at the class level in their ``__init__`` as explicit keyword arguments (no ``*args`` or ``**kwargs``). @@ -54,12 +54,12 @@ def get_params(self, deep: bool = True) -> Dict[str, Any]: """Get parameters for this estimator. Args: - deep: + deep (bool, default True): Default ``True``. If True, will return the parameters for this estimator and contained subobjects that are estimators. Returns: - A dictionary of parameter names mapped to their values. + Dictionary: A dictionary of parameter names mapped to their values. """ out: Dict = dict() for key in self._get_param_names(): @@ -84,15 +84,15 @@ def score(self, X, y): each label set be correctly predicted. Args: - X: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples, n_features). Test samples. - y: + y (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples,) or (n_samples, n_outputs). True labels for `X`. Returns: - A DataFrame of the evaluation result. + bigframes.dataframe.DataFrame: A DataFrame of the evaluation result. """ raise NotImplementedError("abstract method") @@ -106,19 +106,19 @@ def score(self, X, y): """Return the evaluation metrics of the model. Args: - X: - DataFrame of shape (n_samples, n_features). Test samples. For + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Test samples. For some estimators this may be a precomputed kernel matrix or a list of generic objects instead with shape ``(n_samples, n_samples_fitted)``, where ``n_samples_fitted`` is the number of samples used in the fitting for the estimator. - y: - DataFrame of shape (n_samples,) or (n_samples, n_outputs). True + y (bigframes.dataframe.DataFrame or bigframes.series.Series: + Series or DataFrame of shape (n_samples,) or (n_samples, n_outputs). True values for `X`. Returns: - A DataFrame of the evaluation result. + bigframes.dataframe.DataFrame: A DataFrame of the evaluation result. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index 8f67cab1c31..bddb82c7ba0 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -24,13 +24,13 @@ def predict(self, X): """Predict the closest cluster each sample in X belongs to. Args: - X: - DataFrame of shape (n_samples, n_features). The data matrix for + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). The data matrix for which we want to get the predictions. Returns: - DataFrame of shape (n_samples,), containing the class labels for - each sample. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing the + class labels for each sample. """ raise NotImplementedError("abstract method") @@ -39,7 +39,7 @@ class KMeans(_BaseKMeans): """K-Means clustering. Args: - n_clusters: int, default=8 + n_clusters (int, default 8): The number of clusters to form as well as the number of centroids to generate. Default to 8. """ @@ -53,19 +53,51 @@ def fit( """Compute k-means clustering. Args: - X: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples, n_features). Training data. - y: Ignored + y (default None): Not used, present here for API consistency by convention. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. - transforms: - An optional list of SQL expressions to apply over top of the - model inputs as preprocessing. This preprocessing will be - automatically reapplied to new input data (e.g. in .predict), - and may contain steps (like ML.STANDARD_SCALER) that fit to the - training data. Returns: - Fitted Estimator. + KMeans: Fitted Estimator. + """ + raise NotImplementedError("abstract method") + + def predict( + self, + X, + ): + """Predict the closest cluster each sample in X belongs to. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + DataFrame of shape (n_samples, n_features). New data to predict. + y: (default None) + Not used, present here for API consistency by convention. + + Returns: + bigframes.dataframe.DataFrame: DataFrame of the cluster each sample belongs to. + """ + raise NotImplementedError("abstract method") + + def score( + self, + X, + y=None, + ): + """Metrics of the model. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + DataFrame of shape (n_samples, n_features). New Data. + y (default None) + Not used, present here for API consistency by convention. + + Returns: + bigframes.dataframe.DataFrame: DataFrame of the metrics. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py index dadfa5d0131..bc8bc3980ac 100644 --- a/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py +++ b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py @@ -35,13 +35,13 @@ def fit( """Fit all transformers using X. Args: - X: - DataFrame of shape (n_samples, n_features). Training vector, + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Series or DataFrame of shape (n_samples, n_features). Training vector, where `n_samples` is the number of samples and `n_features` is the number of features. Returns: - Fitted estimator. + ColumnTransformer: Fitted estimator. """ raise NotImplementedError("abstract method") @@ -52,10 +52,10 @@ def transform( """Transform X separately by each transformer, concatenate results. Args: - X: - The DataFrame to be transformed by subset. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Series or DataFrame to be transformed by subset. Returns: - Transformed result. + bigframes.dataframe.DataFrame: Transformed result. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py index 6f0d17b2e6d..619c13f35d7 100644 --- a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py +++ b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py @@ -30,24 +30,54 @@ class PCA(BaseEstimator, metaclass=ABCMeta): truncated SVD. Args: - n_components: Optional[int] - Number of components to keep. if n_components is not set all components are kept. + n_components (Optional[int], default 3): + Number of components to keep. if n_components is not set all components + are kept. """ - def fit( - self, - X, - ): + def fit(self, X, y=None): """Fit the model according to the given training data. Args: - X: - DataFrame of shape (n_samples, n_features). Training vector, + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Training vector, where `n_samples` is the number of samples and `n_features` is the number of features. + y (default None): + Ignored. + + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. + Returns: - Fitted estimator. + PCA: Fitted estimator. """ raise NotImplementedError("abstract method") + + def score(self, X=None, y=None): + """Return the metrics of the model. + + Args: + X (default None): + Ignored. + + y (default None): + Ignored. + Returns: + bigframes.dataframe.DataFrame: DataFrame that represents model metrics. + """ + raise NotImplementedError("abstract method") + + def predict(self, X): + """Predict the closest cluster for each sample in X. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or a DataFrame to predict. + + Returns: + bigframes.dataframe.DataFrame: predicted DataFrames.""" + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py index c91821c7622..73f4684dc30 100644 --- a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py +++ b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py @@ -46,12 +46,17 @@ def fit(self, X, y): Args: X: - DataFrame of shape (n_samples, n_features). Training data. + Series or DataFrame of shape (n_samples, n_features). Training data. y: - DataFrame of shape (n_samples,) or (n_samples, n_targets). + Series or DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. + + Returns: Fitted Estimator. """ @@ -71,7 +76,7 @@ def predict(self, X): Args: X: - DataFrame of shape (n_samples, n_features). The data matrix for + Series or DataFrame of shape (n_samples, n_features). The data matrix for which we want to get the predictions. Returns: @@ -137,7 +142,7 @@ def predict(self, X): Args: X: - DataFrame of shape (n_samples, n_features). The data matrix for + Series or DataFrame of shape (n_samples, n_features). The data matrix for which we want to get the predictions. Returns: diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py index e67edac1aa4..65e895298df 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_base.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -30,11 +30,11 @@ def predict(self, X): """Predict using the linear model. Args: - X: - DataFrame of shape (n_samples, n_features). Samples. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Samples. Returns: - DataFrame of shape (n_samples,). Returns predicted values. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,). Returns predicted values. """ raise NotImplementedError("abstract method") @@ -44,13 +44,13 @@ def predict(self, X): """Predict class labels for samples in X. Args: - X: - DataFrame of shape (n_samples, n_features). The data matrix for + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). The data matrix for which we want to get the predictions. Returns: - DataFrame of shape (n_samples,), containing the class labels for - each sample. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing + the class labels for each sample. """ raise NotImplementedError("abstract method") @@ -63,7 +63,7 @@ class LinearRegression(RegressorMixin, LinearModel): the dataset, and the targets predicted by the linear approximation. Args: - fit_intercept: + fit_intercept (default True): Default ``True``. Whether to calculate the intercept for this model. If set to False, no intercept will be used in calculations (i.e. data is expected to be centered). @@ -78,21 +78,18 @@ def fit( """Fit linear model. Args: - X: - DataFrame of shape (n_samples, n_features). Training data. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Training data. - y: - DataFrame of shape (n_samples,) or (n_samples, n_targets). + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. - transforms: - An optional list of SQL expressions to apply over top of the - model inputs as preprocessing. This preprocessing will be - automatically reapplied to new input data (e.g. in .predict), - and may contain steps (like ML.STANDARD_SCALER) that fit to the - training data. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. Returns: - Fitted Estimator. + LinearRegression: Fitted Estimator. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py index 6de5bf65e25..8525e570689 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py @@ -24,33 +24,12 @@ class LogisticRegression(LinearClassifierMixin, BaseEstimator): """Logistic Regression (aka logit, MaxEnt) classifier. Args: - fit_intercept: + fit_intercept (default True): Default True. Specifies if a constant (a.k.a. bias or intercept) should be added to the decision function. - auto_class_weights: + auto_class_weights (default False): Default False. If True, balance class labels using weights for each class in inverse proportion to the frequency of that class. - - References: - L-BFGS-B -- Software for Large-scale Bound-constrained Optimization - Ciyou Zhu, Richard Byrd, Jorge Nocedal and Jose Luis Morales. - http://users.iems.northwestern.edu/~nocedal/lbfgsb.html - - LIBLINEAR -- A Library for Large Linear Classification - https://www.csie.ntu.edu.tw/~cjlin/liblinear/ - - SAG -- Mark Schmidt, Nicolas Le Roux, and Francis Bach - Minimizing Finite Sums with the Stochastic Average Gradient - https://hal.inria.fr/hal-00860051/document - - SAGA -- Defazio, A., Bach F. & Lacoste-Julien S. (2014). - "SAGA: A Fast Incremental Gradient Method With Support - for Non-Strongly Convex Composite Objectives" (Arxiv <1407.0202>) - - Hsiang-Fu Yu, Fang-Lan Huang, Chih-Jen Lin (2011). Dual coordinate descent - methods for logistic regression and maximum entropy models. - Machine Learning 85(1-2):41-75. - https://www.csie.ntu.edu.tw/~cjlin/papers/maxent_dual.pdf """ def fit( @@ -62,22 +41,20 @@ def fit( """Fit the model according to the given training data. Args: - X: - DataFrame of shape (n_samples, n_features). Training vector, + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Training vector, where `n_samples` is the number of samples and `n_features` is the number of features. - y: + y (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples,). Target vector relative to X. - transforms: - An optional list of SQL expressions to apply over top of the - model inputs as preprocessing. This preprocessing will be - automatically reapplied to new input data (e.g. in .predict), - and may contain steps (like ML.STANDARD_SCALER) that fit to the - training data. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. + Returns: - Fitted estimator. + LogisticRegression: Fitted Estimator. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/metrics/_classification.py b/third_party/bigframes_vendored/sklearn/metrics/_classification.py index 05fb5f6fb6e..6d9692ac8d6 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_classification.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_classification.py @@ -25,18 +25,17 @@ def accuracy_score(y_true, y_pred, normalize=True) -> float: """Accuracy classification score. Args: - y_true: DataFrame of shape (n_samples,) + y_true (Series or DataFrame of shape (n_samples,)): Ground truth (correct) labels. - y_pred: DataFrame of shape (n_samples,) + y_pred (Series or DataFrame of shape (n_samples,)): Predicted labels, as returned by a classifier. - normalize: bool, default=True + normalize (bool, default True): Default to True. If ``False``, return the number of correctly classified samples. Otherwise, return the fraction of correctly classified samples. Returns: - score: float. - If ``normalize == True``, return the fraction of correctly + float: If ``normalize == True``, return the fraction of correctly classified samples (float), else returns the number of correctly classified samples (int). """ @@ -58,17 +57,16 @@ def confusion_matrix( :math:`C_{1,1}` and false positives is :math:`C_{0,1}`. Args: - y_true: DataFrame of shape (n_samples,) + y_true (Series or DataFrame of shape (n_samples,)): Ground truth (correct) target values. - y_pred: DataFrame of shape (n_samples,) + y_pred (Series or DataFrame of shape (n_samples,)): Estimated targets as returned by a classifier. Returns: - C: DataFrame of shape (n_samples, n_features). - Confusion matrix whose i-th row and j-th - column entry indicates the number of - samples with true label being i-th class - and predicted label being j-th class. + DataFrame of shape (n_samples, n_features): Confusion matrix whose + i-th row and j-th column entry indicates the number of + samples with true label being i-th class and predicted label + being j-th class. """ raise NotImplementedError("abstract method") @@ -87,19 +85,18 @@ def recall_score( The best value is 1 and the worst value is 0. Args: - y_true: DataFrame of shape (n_samples,) + y_true (Series or DataFrame of shape (n_samples,)): Ground truth (correct) target values. - y_pred: DataFrame of shape (n_samples,) + y_pred (Series or DataFrame of shape (n_samples,)): Estimated targets as returned by a classifier. - average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ - default='binary' + average ({'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ + default='binary'): This parameter is required for multiclass/multilabel targets. Possible values are 'None', 'micro', 'macro', 'samples', 'weighted', 'binary'. Returns: - recall: float (if average is not None) or Series of float of shape \ - (n_unique_labels,). - Recall of the positive class in binary classification or weighted + float (if average is not None) or Series of float of shape n_unique_labels,): Recall + of the positive class in binary classification or weighted average of the recall of each class for the multiclass task. """ raise NotImplementedError("abstract method") @@ -120,9 +117,9 @@ def precision_score( The best value is 1 and the worst value is 0. Args: - y_true: DataFrame of shape (n_samples,) + y_true: Series or DataFrame of shape (n_samples,) Ground truth (correct) target values. - y_pred: DataFrame of shape (n_samples,) + y_pred: Series or DataFrame of shape (n_samples,) Estimated targets as returned by a classifier. average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ default='binary' @@ -148,16 +145,16 @@ def f1_score( The F1 score can be interpreted as a harmonic mean of the precision and recall, where an F1 score reaches its best value at 1 and worst score at 0. The relative contribution of precision and recall to the F1 score are - equal. The formula for the F1 score is: F1 = 2 * (precision * recall) / (precision + recall) + equal. The formula for the F1 score is: F1 = 2 * (precision * recall) / (precision + recall). In the multi-class and multi-label case, this is the average of the F1 score of each class with weighting depending on the ``average`` parameter. Args: - y_true: DataFrame of shape (n_samples,) + y_true: Series or DataFrame of shape (n_samples,) Ground truth (correct) target values. - y_pred: DataFrame of shape (n_samples,) + y_pred: Series or DataFrame of shape (n_samples,) Estimated targets as returned by a classifier. average: {'micro', 'macro', 'samples', 'weighted', 'binary'} or None, \ default='binary' diff --git a/third_party/bigframes_vendored/sklearn/metrics/_ranking.py b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py index b7625171b56..693996070f0 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_ranking.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py @@ -26,15 +26,14 @@ def auc(x, y) -> float: :func:`average_precision_score`. Args: - x : DataFrame of shape (n_samples,) + x (Series or DataFrame of shape (n_samples,)): X coordinates. These must be either monotonic increasing or monotonic decreasing. - y : DataFrame of shape (n_samples,) + y (Series or DataFrame of shape (n_samples,)): Y coordinates. Returns: - auc : float - Area Under the Curve. + float: Area Under the Curve. """ raise NotImplementedError("abstract method") @@ -44,11 +43,11 @@ def roc_auc_score(y_true, y_score) -> float: from prediction scores. Args: - y_true: DataFrame of shape (n_samples,) + y_true (Series or DataFrame of shape (n_samples,)): True labels or binary label indicators. The binary and multiclass cases expect labels with shape (n_samples,) while the multilabel case expects binary label indicators with shape (n_samples, n_classes). - y_score: DataFrame of shape (n_samples,) + y_score (Series or DataFrame of shape (n_samples,)): Target scores. * In the binary case, it corresponds to an array of shape `(n_samples,)`. Both probability estimates and non-thresholded @@ -59,8 +58,7 @@ def roc_auc_score(y_true, y_score) -> float: corresponds to the output of `estimator.decision_function(X, y)`. Returns: - auc: float. - Area Under the Curve score. + float: Area Under the Curve score. """ raise NotImplementedError("abstract method") @@ -73,10 +71,10 @@ def roc_curve( """Compute Receiver operating characteristic (ROC). Args: - y_true: DataFrame of shape (n_samples,) + y_true: Series or DataFrame of shape (n_samples,) True binary labels. If labels are not either {-1, 1} or {0, 1}, then pos_label should be explicitly given. - y_score: DataFrame of shape (n_samples,) + y_score: Series or DataFrame of shape (n_samples,) Target scores, can either be probability estimates of the positive class, confidence values, or non-thresholded measure of decisions (as returned by "decision_function" on some classifiers). diff --git a/third_party/bigframes_vendored/sklearn/metrics/_regression.py b/third_party/bigframes_vendored/sklearn/metrics/_regression.py index 44ceab48a6a..b90c4158878 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_regression.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_regression.py @@ -41,13 +41,12 @@ def r2_score(y_true, y_pred, force_finite=True) -> float: predictions) respectively. Args: - y_true: DataFrame of shape (n_samples,) + y_true (Series or DataFrame of shape (n_samples,)): Ground truth (correct) target values. - y_pred: DataFrame of shape (n_samples,) + y_pred (Series or DataFrame of shape (n_samples,)): Estimated target values. Returns: - z: float. - The :math:`R^2` score. + float: The :math:`R^2` score. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/pipeline.py b/third_party/bigframes_vendored/sklearn/pipeline.py index 7374edf990a..f8bbae86df6 100644 --- a/third_party/bigframes_vendored/sklearn/pipeline.py +++ b/third_party/bigframes_vendored/sklearn/pipeline.py @@ -18,13 +18,13 @@ class Pipeline(BaseEstimator, metaclass=ABCMeta): """Pipeline of transforms with a final estimator. Sequentially apply a list of transforms and a final estimator. - Intermediate steps of the pipeline must be 'transforms', that is, they + Intermediate steps of the pipeline must be `transforms`, that is, they must implement `fit` and `transform` methods. The final estimator only needs to implement `fit`. The purpose of the pipeline is to assemble several steps that can be cross-validated together while setting different parameters. This simplifies code, and allows deploying an estimator - and peprocessing together, e.g. with Pipeline.to_gbq(...) + and peprocessing together, e.g. with `Pipeline.to_gbq(...).` """ def fit( @@ -38,14 +38,14 @@ def fit( data. Finally, fit the transformed data using the final estimator. Args: - X: - A BigQuery DataFrames representing training data. Must match the + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A DataFrame or Series representing training data. Must match the input requirements of the first step of the pipeline. - y: - A BigQuery DataFrames representing training targets, if applicable. + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + A DataFrame or Series representing training targets, if applicable. Returns: - Pipeline with fitted steps. + Pipeline: Pipeline with fitted steps. """ raise NotImplementedError("abstract method") @@ -58,14 +58,14 @@ def score(self, X, y): `score` method. Only valid if the final estimator implements `score`. Args: - X: - A BigQuery DataFrames as evaluation data. - y: - A BigQuery DataFrames as evaluation labels. - - Returns: - A BigQuery DataFrames representing the result of calling - `score` on the final estimator. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A DataFrame or Series as evaluation data. + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + A DataFrame or Series as evaluation labels. + + Returns:; + DataFrame: A DataFrame representing the result + of calling `score` on the final estimator. """ raise NotImplementedError("abstract method") @@ -74,10 +74,11 @@ def predict(self, X): """Predict the pipeline result for each sample in X. Args: - X: - A BigQuery DataFrames to predict. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + A DataFrame or Series to predict. Returns: - A BigQuery DataFrames Dataframe representing predicted result. + bigframes.dataframe.DataFrame: A Dataframe representing + predicted result. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index 7a44c4e66d3..c57d1f22304 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -32,12 +32,32 @@ class StandardScaler(BaseEstimator): def fit(self, X): """Compute the mean and std to be used for later scaling. + Examples: + + .. code-block:: + + from bigframes.ml.preprocessing import StandardScaler + + enc = StandardScaler() + X = [['Male', 1], ['Female', 3], ['Female', 2]] + enc.fit(X) + + Examples: + + .. code-block:: + + from bigframes.ml import StandardScaler + + enc = StandardScaler() + X = [['Male', 1], ['Female', 3], ['Female', 2]] + enc.fit(X) + Args: - X: - A dataframe with training data. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Dataframe or Series with training data. Returns: - Fitted scaler. + StandardScaler: Fitted scaler. """ raise NotImplementedError("abstract method") @@ -45,9 +65,10 @@ def transform(self, X): """Perform standardization by centering and scaling. Args: - X: - The DataFrame to be transformed. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. Returns: - Transformed result.""" + bigframes.dataframe.DataFrame: Transformed result. + """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py index 6f0565ac136..a6c32d91c1d 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -14,18 +14,52 @@ class OneHotEncoder(BaseEstimator): encoding scheme. Note that this method deviates from Scikit-Learn; instead of producing sparse - binary columns, the encoding is a single column of STRUCT + binary columns, the encoding is a single column of `STRUCT`. + + Args: + drop (Optional[Literal["most_frequent"]], default None): + Specifies a methodology to use to drop one of the categories per feature. + This is useful in situations where perfectly collinear features cause problems, + such as when feeding the resulting data into an unregularized linear regression model. + However, dropping one category breaks the symmetry of the original representation + and can therefore induce a bias in downstream models, for instance for penalized + linear classification or regression models. + Default None: retain all the categories. + "most_frequent": Drop the most frequent category found in the string expression. + Selecting this value causes the function to use dummy encoding. + min_frequency (Optional[int], default None): + Specifies the minimum frequency below which a category will be considered infrequent. + Default None. + int: categories with a smaller cardinality will be considered infrequent as index 0. + max_categories (Optional[int], default None): + Specifies an upper limit to the number of output features for each input feature + when considering infrequent categories. If there are infrequent categories, + max_categories includes the category representing the infrequent categories along with the frequent categories. + Default None, set limit to 1,000,000. """ def fit(self, X): """Fit OneHotEncoder to X. + Examples: + + Given a dataset with two features, we let the encoder find the unique + values per feature and transform the data to a binary one-hot encoding. + + .. code-block:: + + from bigframes.ml.preprocessing import OneHotEncoder + + enc = OneHotEncoder() + X = [['Male', 1], ['Female', 3], ['Female', 2]] + enc.fit(X) + Args: - X: - A dataframe with training data. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series with training data. Returns: - Fitted encoder. + OneHotEncoder: Fitted encoder. """ raise NotImplementedError("abstract method") @@ -33,9 +67,10 @@ def transform(self, X): """Transform X using one-hot encoding. Args: - X: - The DataFrame to be transformed. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. Returns: - Transformed result.""" + bigframes.dataframe.DataFrame: The result is categorized as index: number, value: number. + Where index is the position of the dict that seeing the category, and value is 0 or 1.""" raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/xgboost/sklearn.py b/third_party/bigframes_vendored/xgboost/sklearn.py index 268484b2c23..fcb5d2ec593 100644 --- a/third_party/bigframes_vendored/xgboost/sklearn.py +++ b/third_party/bigframes_vendored/xgboost/sklearn.py @@ -12,11 +12,11 @@ def predict(self, X): """Predict using the XGB model. Args: - X: - DataFrame of shape (n_samples, n_features). Samples. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Samples. Returns: - DataFrame of shape (n_samples,). Returns predicted values. + DataFrame of shape (n_samples,): Returns predicted values. """ raise NotImplementedError("abstract method") @@ -28,15 +28,19 @@ def fit(self, X, y): pass ``xgb_model`` argument. Args: - X: - DataFrame of shape (n_samples, n_features). Training data. + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). Training data. - y: + y (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. + transforms (Optional[List[str]], default None): + Do not use. Internal param to be deprecated. + Use bigframes.ml.pipeline instead. + Returns: - Fitted Estimator. + XGBModel: Fitted Estimator. """ raise NotImplementedError("abstract method") @@ -53,46 +57,47 @@ class XGBRegressor(XGBModel, XGBRegressorBase): XGBoost regression model. Args: - num_parallel_tree: Optional[int] + num_parallel_tree (Optional[int]): Number of parallel trees constructed during each iteration. Default to 1. - booster: Optional[str] + booster (Optional[str]): Specify which booster to use: gbtree or dart. Default to "gbtree". - dart_normalized_type": Optional[str] + dart_normalized_type (Optional[str]): Type of normalization algorithm for DART booster. Possible values: "TREE", "FOREST". Default to "TREE". - tree_method: Optional[str] + tree_method (Optional[str]): Specify which tree method to use. Default to "auto". If this parameter is set to - default, XGBoost will choose the most conservative option available. - min_child_weight : Optional[float] + default, XGBoost will choose the most conservative option available. Possible values: ""exact", "approx", + "hist". + min_child_weight (Optional[float]): Minimum sum of instance weight(hessian) needed in a child. Default to 1. - colsample_bytree : Optional[float] + colsample_bytree (Optional[float]): Subsample ratio of columns when constructing each tree. Default to 1.0. - colsample_bylevel : Optional[float] + colsample_bylevel (Optional[float]): Subsample ratio of columns for each level. Default to 1.0. - colsample_bynode : Optional[float] + colsample_bynode (Optional[float]): Subsample ratio of columns for each split. Default to 1.0. - gamma : Optional[float] + gamma (Optional[float]): (min_split_loss) Minimum loss reduction required to make a further partition on a leaf node of the tree. Default to 0.0. - max_depth : Optional[int] + max_depth (Optional[int]): Maximum tree depth for base learners. Default to 6. - subsample : Optional[float] + subsample (Optional[float]): Subsample ratio of the training instance. Default to 1.0. - reg_alpha : Optional[float] + reg_alpha (Optional[float]): L1 regularization term on weights (xgb's alpha). Default to 0.0. - reg_lambda : Optional[float] + reg_lambda (Optional[float]): L2 regularization term on weights (xgb's lambda). Default to 1.0. - early_stop: Optional[bool] + early_stop (Optional[bool]): Whether training should stop after the first iteration. Default to True. - learning_rate: Optional[float] + learning_rate (Optional[float]): Boosting learning rate (xgb's "eta"). Default to 0.3. - max_iterations: Optional[int] + max_iterations (Optional[int]): Maximum number of rounds for boosting. Default to 20. - min_rel_progress: Optional[float] + min_rel_progress (Optional[float]): Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. - enable_global_explain: Optional[bool] + enable_global_explain (Optional[bool]): Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. - xgboost_version: Optional[str] - Specifies the Xgboost version for model training. Default to "0.9". + xgboost_version (Optional[str]): + Specifies the Xgboost version for model training. Default to "0.9". Possible values: "0.9", "1.1". """ @@ -101,45 +106,45 @@ class XGBClassifier(XGBModel, XGBClassifierMixIn, XGBClassifierBase): XGBoost classifier model. Args: - num_parallel_tree: Optional[int] + num_parallel_tree (Optional[int]): Number of parallel trees constructed during each iteration. Default to 1. - booster: Optional[str] + booster (Optional[str]): Specify which booster to use: gbtree or dart. Default to "gbtree". - dart_normalized_type": Optional[str] + dart_normalized_type (Optional[str]): Type of normalization algorithm for DART booster. Possible values: "TREE", "FOREST". Default to "TREE". - tree_method: Optional[str] + tree_method (Optional[str]): Specify which tree method to use. Default to "auto". If this parameter is set to default, XGBoost will choose the most conservative option available. Possible values: ""exact", "approx", "hist". - min_child_weight : Optional[float] + min_child_weight (Optional[float]): Minimum sum of instance weight(hessian) needed in a child. Default to 1. - colsample_bytree : Optional[float] + colsample_bytree (Optional[float]): Subsample ratio of columns when constructing each tree. Default to 1.0. - colsample_bylevel : Optional[float] + colsample_bylevel (Optional[float]): Subsample ratio of columns for each level. Default to 1.0. - colsample_bynode : Optional[float] + colsample_bynode (Optional[float]): Subsample ratio of columns for each split. Default to 1.0. - gamma : Optional[float] + gamma (Optional[float]): (min_split_loss) Minimum loss reduction required to make a further partition on a leaf node of the tree. Default to 0.0. - max_depth : Optional[int] + max_depth (Optional[int]): Maximum tree depth for base learners. Default to 6. - subsample : Optional[float] + subsample (Optional[float]): Subsample ratio of the training instance. Default to 1.0. - reg_alpha : Optional[float] + reg_alpha (Optional[float]): L1 regularization term on weights (xgb's alpha). Default to 0.0. - reg_lambda : Optional[float] + reg_lambda (Optional[float]): L2 regularization term on weights (xgb's lambda). Default to 1.0. - early_stop: Optional[bool] + early_stop (Optional[bool]): Whether training should stop after the first iteration. Default to True. - learning_rate: Optional[float] + learning_rate (Optional[float]): Boosting learning rate (xgb's "eta"). Default to 0.3. - max_iterations: Optional[int] + max_iterations (Optional[int]): Maximum number of rounds for boosting. Default to 20. - min_rel_progress: Optional[float] + min_rel_progress (Optional[float]): Minimum relative loss improvement necessary to continue training when early_stop is set to True. Default to 0.01. - enable_global_explain: Optional[bool] + enable_global_explain (Optional[bool]): Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. - xgboost_version: Optional[str] + xgboost_version (Optional[str]): Specifies the Xgboost version for model training. Default to "0.9". Possible values: "0.9", "1.1". """ From 3e143d793e300f1cf27e2b1822170d69876eaa22 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 14 Aug 2023 10:36:04 -0500 Subject: [PATCH 04/80] chore: release 0.1.1 (#5) Change-Id: I10199358657195d60fb414360f1bff4a94490bad --- CHANGELOG.md | 7 +++++++ README.rst | 4 ++-- bigframes/version.py | 2 +- docs/user_guide/bigframes.pandas/remote_functions.rst | 5 +++-- setup.py | 8 +++++--- 5 files changed, 18 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d966220bd2..19bf3a117b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ [1]: https://pypi.org/project/bigframes/#history +## 0.1.1 (2023-08-14) + +### Documentation + +* Correct link to code repository in `setup.py` and use correct terminology for + `console.cloud.google.com` links. + ## 0.1.0 (2023-08-11) ### Features diff --git a/README.rst b/README.rst index 6f51dfde23e..c6dbb059576 100644 --- a/README.rst +++ b/README.rst @@ -117,7 +117,7 @@ BigQuery connections are created in the same location as the BigQuery DataFrames session, using the name you provide in the custom function definition. To view and manage connections, do the following: -1. Go to `BigQuery Studio `__. +1. Go to `BigQuery in the Google Cloud Console `__. 2. Select the project in which you created the remote function. 3. In the Explorer pane, expand that project and then expand External connections. @@ -127,7 +127,7 @@ the location used by the BigQuery DataFrames session. For example, ``bigframes_temp_us_central1``. To view and manage remote functions, do the following: -1. Go to `BigQuery Studio `__. +1. Go to `BigQuery in the Google Cloud Console `__. 2. Select the project in which you created the remote function. 3. In the Explorer pane, expand that project, expand the dataset in which you created the remote function, and then expand Routines. diff --git a/bigframes/version.py b/bigframes/version.py index df6dc3c17c8..c387f59d2c6 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.1.0" +__version__ = "0.1.1" diff --git a/docs/user_guide/bigframes.pandas/remote_functions.rst b/docs/user_guide/bigframes.pandas/remote_functions.rst index abfe6a10caa..7540ba8a28c 100644 --- a/docs/user_guide/bigframes.pandas/remote_functions.rst +++ b/docs/user_guide/bigframes.pandas/remote_functions.rst @@ -104,8 +104,9 @@ Using Existing Functions ======================== If you have already defined a custom function in BigQuery, either in the -BigQuery Studio or with the `remote_function` decorator above or otherwise, you -may use it with BigQuery DataFrames with the `read_gbq_function` method. +BigQuery Google Cloud Console or with the `remote_function` decorator above or +otherwise, you may use it with BigQuery DataFrames with the `read_gbq_function` +method. More details are available via the `help` command: diff --git a/setup.py b/setup.py index 3e36cd3ecfb..20f080b1668 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,9 @@ # Package metadata. name = "bigframes" -description = "Scalable DataFrames with BigQuery" +description = ( + "BigQuery DataFrames -- scalable analytics and machine learning with BigQuery" +) # Should be one of: # 'Development Status :: 3 - Alpha' @@ -85,9 +87,9 @@ description=description, long_description=readme, author="Google LLC", - author_email="googleapis-packages@google.com", + author_email="bigframes-feedback@google.com", license="Apache 2.0", - url="https://github.com/googleapis/python-bigquery", + url="https://github.com/googleapis/python-bigquery-dataframes", classifiers=[ release_status, "Intended Audience :: Developers", From 76f4daa45db9be12abb112683f2bd2f89589ff86 Mon Sep 17 00:00:00 2001 From: Genesis929 <142538604+Genesis929@users.noreply.github.com> Date: Thu, 17 Aug 2023 13:17:39 -0700 Subject: [PATCH 05/80] chore: release 0.2.0 (#6) * chore: release 0.2.0 Change-Id: Id9f7843e3c5238e82a2514b5f9b81d7335460c6e * Update CHANGELOG Change-Id: Ic77a9acfdbffe3bf97103d164f606f85fc9fe4db --- .kokoro/docker/docs/Dockerfile | 88 + .kokoro/docker/docs/fetch_gpg_keys.sh | 45 + .kokoro/docs/common.cfg | 61 + .kokoro/docs/docs-presubmit.cfg | 23 + .kokoro/docs/docs.cfg | 1 + .kokoro/publish-docs.sh | 46 + .kokoro/trampoline_v2.sh | 487 ++++ .trampolinerc | 58 + CHANGELOG.md | 19 + bigframes/core/__init__.py | 7 +- bigframes/core/blocks.py | 10 +- bigframes/core/groupby/__init__.py | 122 +- bigframes/core/indexes/index.py | 6 +- bigframes/core/io.py | 5 + bigframes/core/utils.py | 11 + bigframes/dataframe.py | 73 +- bigframes/formatting_helpers.py | 32 +- bigframes/ml/cluster.py | 9 + bigframes/ml/core.py | 17 +- bigframes/ml/sql.py | 7 + bigframes/operations/__init__.py | 36 + bigframes/operations/base.py | 6 +- bigframes/operations/strings.py | 25 +- bigframes/pandas/__init__.py | 5 + bigframes/remote_function.py | 55 +- bigframes/series.py | 41 +- bigframes/session.py | 37 +- bigframes/version.py | 2 +- docs/templates/toc.yml | 130 + .../03 - Using ML - ML fundamentals.ipynb | 2425 ++++++++++++----- noxfile.py | 25 +- tests/system/conftest.py | 32 + tests/system/small/ml/conftest.py | 41 +- tests/system/small/ml/test_cluster.py | 66 + tests/system/small/ml/test_core.py | 78 +- tests/system/small/operations/test_strings.py | 83 + tests/system/small/test_dataframe.py | 20 +- tests/system/small/test_dataframe_io.py | 6 +- tests/system/small/test_groupby.py | 105 + tests/system/small/test_remote_function.py | 70 +- tests/system/small/test_series.py | 93 + tests/system/small/test_session.py | 21 + tests/unit/core/test_io.py | 49 + tests/unit/ml/test_sql.py | 5 + .../pandas/core/groupby/__init__.py | 45 +- .../bigframes_vendored/pandas/core/series.py | 45 + .../pandas/core/strings/accessor.py | 113 + .../sklearn/cluster/_kmeans.py | 21 + 48 files changed, 3883 insertions(+), 924 deletions(-) create mode 100644 .kokoro/docker/docs/Dockerfile create mode 100644 .kokoro/docker/docs/fetch_gpg_keys.sh create mode 100644 .kokoro/docs/common.cfg create mode 100644 .kokoro/docs/docs-presubmit.cfg create mode 100644 .kokoro/docs/docs.cfg create mode 100644 .kokoro/publish-docs.sh create mode 100644 .kokoro/trampoline_v2.sh create mode 100644 .trampolinerc create mode 100644 docs/templates/toc.yml create mode 100644 tests/unit/core/test_io.py diff --git a/.kokoro/docker/docs/Dockerfile b/.kokoro/docker/docs/Dockerfile new file mode 100644 index 00000000000..d300bee2607 --- /dev/null +++ b/.kokoro/docker/docs/Dockerfile @@ -0,0 +1,88 @@ +# 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. + +from ubuntu:22.04 + +ENV DEBIAN_FRONTEND noninteractive + +# Ensure local Python is preferred over distribution Python. +ENV PATH /usr/local/bin:$PATH + +# Install dependencies. +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + apt-transport-https \ + build-essential \ + ca-certificates \ + curl \ + dirmngr \ + git \ + gpg-agent \ + graphviz \ + libbz2-dev \ + libdb5.3-dev \ + libexpat1-dev \ + libffi-dev \ + liblzma-dev \ + libreadline-dev \ + libsnappy-dev \ + libssl-dev \ + libsqlite3-dev \ + portaudio19-dev \ + python3-distutils \ + redis-server \ + software-properties-common \ + ssh \ + sudo \ + tcl \ + tcl-dev \ + tk \ + tk-dev \ + uuid-dev \ + wget \ + zlib1g-dev \ + && add-apt-repository universe \ + && apt-get update \ + && apt-get -y install jq \ + && apt-get clean autoclean \ + && apt-get autoremove -y \ + && rm -rf /var/lib/apt/lists/* \ + && rm -f /var/cache/apt/archives/*.deb + +###################### Install python 3.9.13 and 3.10.5 + +# Download python 3.9.13 and 3.10.5 +RUN wget https://www.python.org/ftp/python/3.9.13/Python-3.9.13.tgz +RUN wget https://www.python.org/ftp/python/3.9.13/Python-3.10.5.tgz + +# Extract files +RUN tar -xvf Python-3.9.13.tgz +RUN tar -xvf Python-3.10.5.tgz + +# Install python 3.9.13 and 3.10.5 +RUN ./Python-3.9.13/configure --enable-optimizations +RUN ./Python-3.10.5/configure --enable-optimizations +RUN make altinstall + +###################### Install pip +RUN wget -O /tmp/get-pip.py 'https://bootstrap.pypa.io/get-pip.py' \ + && python3 /tmp/get-pip.py \ + && rm /tmp/get-pip.py + +# Test pip +RUN python3 -m pip +RUN python3.9 -m pip +RUN python3.10 -m pip + +CMD ["python3.10"] diff --git a/.kokoro/docker/docs/fetch_gpg_keys.sh b/.kokoro/docker/docs/fetch_gpg_keys.sh new file mode 100644 index 00000000000..c4a92a33eaa --- /dev/null +++ b/.kokoro/docker/docs/fetch_gpg_keys.sh @@ -0,0 +1,45 @@ +#!/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 +# +# 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. + +# A script to fetch gpg keys with retry. +# Avoid jinja parsing the file. +# + +function retry { + if [[ "${#}" -le 1 ]]; then + echo "Usage: ${0} retry_count commands.." + exit 1 + fi + local retries=${1} + local command="${@:2}" + until [[ "${retries}" -le 0 ]]; do + $command && return 0 + if [[ $? -ne 0 ]]; then + echo "command failed, retrying" + ((retries--)) + fi + done + return 1 +} + +# 3.6.9, 3.7.5 (Ned Deily) +retry 3 gpg --keyserver ha.pool.sks-keyservers.net --recv-keys \ + 0D96DF4D4110E5C43FBFB17F2D347EA6AA65421D + +# 3.8.0 (Łukasz Langa) +retry 3 gpg --keyserver ha.pool.sks-keyservers.net --recv-keys \ + E3FF2839C048B25C084DEBE9B26995E310250568 + +# diff --git a/.kokoro/docs/common.cfg b/.kokoro/docs/common.cfg new file mode 100644 index 00000000000..3e36916024a --- /dev/null +++ b/.kokoro/docs/common.cfg @@ -0,0 +1,61 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# 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" + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-lib-docs" +} +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/publish-docs.sh" +} + +env_vars: { + key: "V2_STAGING_BUCKET" + # Push google cloud library docs to the Cloud RAD bucket `docs-staging-v2` + value: "docs-staging-v2" +} + +# It will upload the docker image after successful builds. +env_vars: { + key: "TRAMPOLINE_IMAGE_UPLOAD" + value: "true" +} + +# It will always build the docker image. +env_vars: { + key: "TRAMPOLINE_DOCKERFILE" + value: ".kokoro/docker/docs/Dockerfile" +} + +# Fetch the token needed for reporting release status to GitHub +before_action { + fetch_keystore { + keystore_resource { + keystore_config_id: 73713 + keyname: "yoshi-automation-github-key" + } + } +} + +before_action { + fetch_keystore { + keystore_resource { + keystore_config_id: 73713 + keyname: "docuploader_service_account" + } + } +} diff --git a/.kokoro/docs/docs-presubmit.cfg b/.kokoro/docs/docs-presubmit.cfg new file mode 100644 index 00000000000..43ec87185e4 --- /dev/null +++ b/.kokoro/docs/docs-presubmit.cfg @@ -0,0 +1,23 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "V2_STAGING_BUCKET" + value: "gcloud-python-test" +} + +# We only upload the image in the main `docs` build. +env_vars: { + key: "TRAMPOLINE_IMAGE_UPLOAD" + value: "false" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/build.sh" +} + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "docfx" +} diff --git a/.kokoro/docs/docs.cfg b/.kokoro/docs/docs.cfg new file mode 100644 index 00000000000..18a4c35325b --- /dev/null +++ b/.kokoro/docs/docs.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto diff --git a/.kokoro/publish-docs.sh b/.kokoro/publish-docs.sh new file mode 100644 index 00000000000..3fc3905ac1c --- /dev/null +++ b/.kokoro/publish-docs.sh @@ -0,0 +1,46 @@ +#!/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 + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=1 + +export PATH="${HOME}/.local/bin:${PATH}" + +# Install nox +python3 -m pip install --require-hashes -r .kokoro/requirements.txt +python3 -m nox --version + +# docfx yaml files +nox -s docfx + +# create metadata. +python3 -m docuploader create-metadata \ + --name=$(jq --raw-output '.name // empty' .repo-metadata.json) \ + --version=$(python3 setup.py --version) \ + --language=$(jq --raw-output '.language // empty' .repo-metadata.json) \ + --distribution-name=$(python3 setup.py --name) \ + --product-page=$(jq --raw-output '.product_documentation // empty' .repo-metadata.json) \ + --github-repository=$(jq --raw-output '.repo // empty' .repo-metadata.json) \ + --issue-tracker=$(jq --raw-output '.issue_tracker // empty' .repo-metadata.json) + +cat docs.metadata + +# Replace toc.yml template file +mv docs/templates/toc.yml docs/_build/html/docfx_yaml/toc.yml + +# upload docs +python3 -m docuploader upload docs/_build/html/docfx_yaml --metadata-file docs.metadata --destination-prefix docfx --staging-bucket "${V2_STAGING_BUCKET}" diff --git a/.kokoro/trampoline_v2.sh b/.kokoro/trampoline_v2.sh new file mode 100644 index 00000000000..59a7cf3a937 --- /dev/null +++ b/.kokoro/trampoline_v2.sh @@ -0,0 +1,487 @@ +#!/usr/bin/env 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 +# +# 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). +# +# gsutil cp gs://cloud-devrel-kokoro-resources/python-docs-samples/secrets_viewer_service_account.json /dev/shm +# gsutil 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/.trampolinerc b/.trampolinerc new file mode 100644 index 00000000000..4c16e33fee0 --- /dev/null +++ b/.trampolinerc @@ -0,0 +1,58 @@ +# 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. + +# Add required env vars here. +required_envvars+=( +) + +# Add env vars which are passed down into the container here. +pass_down_envvars+=( + "NOX_SESSION" + ############### + "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/CHANGELOG.md b/CHANGELOG.md index 19bf3a117b6..9ab2b05d571 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,25 @@ [1]: https://pypi.org/project/bigframes/#history +## 0.2.0 (2023-08-17) + +### Features +* Add KMeans.cluster_centers_. +* Allow column labels to be any type handled by bq df, column labels can be integers now. +* Add dataframegroupby.agg(). +* Add Series Property is_monotonic_increasing and is_monotonic_decreasing. +* Add match, fullmatch, get, pad str methods. +* Add series isin function. + +### Bug Fixes +* Update ML package to use sessions for queries. +* Optimize `read_gbq` with `index_col` set to cluster by `index_col`. +* Raise ValueError if the location mismatched. +* `read_gbq` no longer uses 'time travel' with query inputs. + +### Documentation +* Add docstring to _uniform_sampling to avoid user using it. + ## 0.1.1 (2023-08-14) ### Documentation diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 075e27b0c23..9f392ce1492 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -304,6 +304,11 @@ def reversed(self) -> ArrayValue: return expr_builder.build() def _uniform_sampling(self, fraction: float) -> ArrayValue: + """Sampling the table on given fraction. + + .. warning:: + The row numbers of result is non-deterministic, avoid to use. + """ table = self.to_ibis_expr( ordering_mode="order_by", expose_hidden_cols=True, fraction=fraction ) @@ -849,7 +854,7 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal def unpivot_single_row( self, - row_labels: typing.Sequence[typing.Optional[str]], + row_labels: typing.Sequence[typing.Hashable], unpivot_columns: typing.Sequence[typing.Tuple[str, typing.Sequence[str]]], *, index_col_id: str = "index", diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 5ef92630727..f696b8287bc 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -47,7 +47,7 @@ import bigframes.operations.aggregations as agg_ops # Type constraint for wherever column labels are used -Label = typing.Optional[str] +Label = typing.Hashable # Bytes to Megabyte Conversion _BYTES_TO_KILOBYTES = 1024 @@ -585,8 +585,8 @@ def with_column_labels(self, value: typing.Iterable[Label]) -> Block: def with_index_labels(self, value: typing.Sequence[Label]) -> Block: if len(value) != len(self.index_columns): raise ValueError( - f"The index labels size `{len(value)} ` should equal to the index" - + f"columns size: {len(self.value_columns)}." + f"The index labels size `{len(value)} ` should equal to the index " + + f"columns size: {len(self.index_columns)}." ) return Block( self._expr, @@ -1126,10 +1126,6 @@ def block_from_local(data, session=None, use_index=True) -> Block: f"multi-indices not supported. {constants.FEEDBACK_LINK}" ) index_label = pd_data.index.name - if (index_label is not None) and (not isinstance(index_label, str)): - raise NotImplementedError( - f"Only string index names supported. {constants.FEEDBACK_LINK}" - ) index_id = guid.generate_guid() pd_data = pd_data.reset_index(names=index_id) diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 288dcdd5b0b..10bee4f56d9 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -15,14 +15,13 @@ from __future__ import annotations import typing - -import pandas as pd -import typing_extensions +import warnings import bigframes.constants as constants import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.ordering as order +import bigframes.core.utils as utils import bigframes.core.window as windows import bigframes.dataframe as df import bigframes.dtypes as dtypes @@ -75,7 +74,7 @@ def __getitem__( typing.Sequence[blocks.Label], ], ): - if _is_list_like(key): + if utils.is_list_like(key): keys = list(key) else: keys = [key] @@ -168,6 +167,76 @@ def cummax(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: def cumprod(self, *args, **kwargs) -> df.DataFrame: return self._apply_window_op(agg_ops.product_op, numeric_only=True) + def agg(self, func=None, **kwargs) -> df.DataFrame: + column_labels = [] + if func: + warnings.warn( + "DataFrameGroupby aggregate produces single-level column labels only currently. Subject to change in future versions." + ) + if isinstance(func, str): + aggregations = [ + (col_id, agg_ops.AGGREGATIONS_LOOKUP[func]) + for col_id in self._aggregated_columns() + ] + elif utils.is_dict_like(func): + aggregations = [] + for label, funcs_for_id in func.items(): + col_id = self._resolve_label(label) + func_list = ( + funcs_for_id + if utils.is_list_like(funcs_for_id) + else [funcs_for_id] + ) + for f in func_list: + aggregations.append((col_id, agg_ops.AGGREGATIONS_LOOKUP[f])) + # Pandas creates multi-index here instead + column_labels.append(f"{label}_{f}") + elif utils.is_list_like(func): + aggregations = [ + (col_id, agg_ops.AGGREGATIONS_LOOKUP[f]) + for col_id in self._aggregated_columns() + for f in func + ] + column_labels = [ + f"{self._block.col_id_to_label[col_id]}_{f}" + for col_id in self._aggregated_columns() + for f in func + ] + else: + raise NotImplementedError( + f"Aggregate with {func} not supported. {constants.FEEDBACK_LINK}" + ) + else: + aggregations = [] + for k, v in kwargs.items(): + if not isinstance(k, str): + raise NotImplementedError( + f"Only string aggregate names supported. {constants.FEEDBACK_LINK}" + ) + if not hasattr(v, "column") or not hasattr(v, "aggfunc"): + import bigframes.pandas as bpd + + raise NotImplementedError( + f"kwargs values must be {bpd.NamedAgg.__qualname__}" + ) + col_id = self._resolve_label(v.column) + aggregations.append((col_id, agg_ops.AGGREGATIONS_LOOKUP[v.aggfunc])) + column_labels.append(k) + + agg_block, _ = self._block.aggregate( + by_column_ids=self._by_col_ids, + aggregations=aggregations, + as_index=self._as_index, + dropna=self._dropna, + ) + + if column_labels: + agg_block = agg_block.with_column_labels(column_labels) + + return df.DataFrame(agg_block) + + aggregate = agg + def _raise_on_non_numeric(self, op: str): if not all( dtype in dtypes.NUMERIC_BIGFRAMES_TYPES for dtype in self._block.dtypes @@ -222,6 +291,15 @@ def _apply_window_op( block = block.select_columns(columns) return df.DataFrame(block) + def _resolve_label(self, label: blocks.Label) -> str: + """Resolve label to column id.""" + col_ids = self._block.label_to_col_id.get(label, ()) + if len(col_ids) > 1: + raise ValueError(f"Label {label} is ambiguous") + if len(col_ids) == 0: + raise ValueError(f"Label {label} does not match any columns") + return col_ids[0] + class SeriesGroupBy(vendored_pandas_groupby.SeriesGroupBy): __doc__ = vendored_pandas_groupby.GroupBy.__doc__ @@ -231,7 +309,7 @@ def __init__( block: blocks.Block, value_column: str, by_col_ids: typing.Sequence[str], - value_name: typing.Optional[str] = None, + value_name: blocks.Label = None, dropna=True, ): # TODO(tbergeron): Support more group-by expression types @@ -278,6 +356,36 @@ def var(self, *args, **kwargs) -> series.Series: def prod(self, *args) -> series.Series: return self._aggregate(agg_ops.product_op) + def agg(self, func=None) -> typing.Union[df.DataFrame, series.Series]: + column_names: list[str] = [] + if isinstance(func, str): + aggregations = [(self._value_column, agg_ops.AGGREGATIONS_LOOKUP[func])] + column_names = [func] + elif utils.is_list_like(func): + aggregations = [ + (self._value_column, agg_ops.AGGREGATIONS_LOOKUP[f]) for f in func + ] + column_names = list(func) + else: + raise NotImplementedError( + f"Aggregate with {func} not supported. {constants.FEEDBACK_LINK}" + ) + + agg_block, _ = self._block.aggregate( + by_column_ids=self._by_col_ids, + aggregations=aggregations, + dropna=self._dropna, + ) + + if column_names: + agg_block = agg_block.with_column_labels(column_names) + + if len(aggregations) > 1: + return df.DataFrame(agg_block) + return series.Series(agg_block) + + aggregate = agg + def cumsum(self, *args, **kwargs) -> series.Series: return self._apply_window_op( agg_ops.sum_op, @@ -374,7 +482,3 @@ def _apply_window_op( skip_null_groups=self._dropna, ) return series.Series(block.select_column(result_id)) - - -def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: - return pd.api.types.is_list_like(obj) diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index bb1fcc4b175..4eb37e6d926 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -37,7 +37,7 @@ def __init__(self, data: blocks.BlockHolder): self._data = data @property - def name(self) -> typing.Optional[str]: + def name(self) -> blocks.Label: return self.names[0] @name.setter @@ -105,11 +105,11 @@ def _expr(self) -> core.ArrayValue: return self._block.expr @property - def name(self) -> typing.Optional[str]: + def name(self) -> blocks.Label: return self._block._index_labels[0] @property - def names(self) -> typing.Sequence[typing.Optional[str]]: + def names(self) -> typing.Sequence[blocks.Label]: return self._block._index_labels @property diff --git a/bigframes/core/io.py b/bigframes/core/io.py index 6b7ed52488d..3c2e5a25f5a 100644 --- a/bigframes/core/io.py +++ b/bigframes/core/io.py @@ -75,6 +75,11 @@ def create_snapshot_sql( if table_ref.dataset_id.upper() == "_SESSION": return f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" + # If we have an anonymous query results table, it can't be modified and + # there isn't any BigQuery time travel. + if table_ref.dataset_id.startswith("_"): + return f"SELECT * FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}`" + return textwrap.dedent( f""" SELECT * diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py index 8ad58fb166e..a3300029054 100644 --- a/bigframes/core/utils.py +++ b/bigframes/core/utils.py @@ -13,6 +13,9 @@ # limitations under the License. import typing +import pandas as pd +import typing_extensions + def get_axis_number(axis: typing.Union[str, int, None]) -> typing.Literal[0, 1]: if axis in {0, "index", "rows", None}: @@ -20,3 +23,11 @@ def get_axis_number(axis: typing.Union[str, int, None]) -> typing.Literal[0, 1]: elif axis in {1, "columns"}: return 1 raise ValueError(f"Not a valid axis: {axis}") + + +def is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: + return pd.api.types.is_list_like(obj) + + +def is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]: + return pd.api.types.is_dict_like(obj) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index cb062f09249..d4b6e470255 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -34,7 +34,6 @@ import google.cloud.bigquery as bigquery import numpy import pandas -import typing_extensions import bigframes import bigframes._config.display_options as display_options @@ -114,7 +113,7 @@ def __init__( # Dict of Series elif ( - _is_dict_like(data) + utils.is_dict_like(data) and len(data) >= 1 and any(isinstance(data[key], bf_series.Series) for key in data.keys()) ): @@ -178,7 +177,13 @@ def __dir__(self): return dir(type(self)) + self._block.column_labels def _ipython_key_completions_(self) -> List[str]: - return list([label for label in self._block.column_labels if label]) + return list( + [ + label + for label in self._block.column_labels + if label and isinstance(label, str) + ] + ) def _find_indices( self, @@ -210,7 +215,7 @@ def _sql_names( tolerance: bool = False, ) -> Sequence[str]: """Retrieve sql name (column name in BQ schema) of column(s).""" - labels = columns if _is_list_like(columns) else [columns] # type:ignore + labels = columns if utils.is_list_like(columns) else [columns] # type:ignore results: Sequence[str] = [] for label in labels: col_ids = self._block.label_to_col_id.get(label, []) @@ -303,13 +308,13 @@ def _to_sql_query( index_labels = [] unnamed_index_count = 0 for index_label in self._block.index_labels: - if index_label is None: + if isinstance(index_label, str): + index_labels.append(index_label) + else: index_labels.append( indexes.INDEX_COLUMN_ID.format(unnamed_index_count), ) unnamed_index_count += 1 - else: - index_labels.append(index_label) column_labels = column_labels + typing.cast( List[Optional[str]], index_labels @@ -365,22 +370,6 @@ def query_job(self) -> Optional[bigquery.QueryJob]: def _set_internal_query_job(self, query_job: bigquery.QueryJob): self._query_job = query_job - @typing.overload - def __getitem__(self, key: bigframes.series.Series) -> DataFrame: - ... - - @typing.overload - def __getitem__(self, key: Sequence[blocks.Label]) -> DataFrame: # type:ignore - ... - - @typing.overload - def __getitem__(self, key: pandas.Index) -> DataFrame: # type:ignore - ... - - @typing.overload - def __getitem__(self, key: blocks.Label) -> bigframes.series.Series: # type:ignore - ... - def __getitem__( self, key: Union[ @@ -390,7 +379,7 @@ def __getitem__( pandas.Index, bigframes.series.Series, ], - ) -> Union[bigframes.series.Series, "DataFrame"]: + ): # No return type annotations (like pandas) as type cannot always be determined statically """Gets the specified column(s) from the DataFrame.""" # NOTE: This implements the operations described in # https://pandas.pydata.org/docs/getting_started/intro_tutorials/03_subset_data.html @@ -399,8 +388,8 @@ def __getitem__( return self._getitem_bool_series(key) sql_names = self._sql_names(key) - # Only input is a str and only find one column, returns a Series - if isinstance(key, str) and len(sql_names) == 1: + # Only input is a single key and only find one column, returns a Series + if (not utils.is_list_like(key)) and len(sql_names) == 1: return bigframes.series.Series(self._block.select_column(sql_names[0])) # Select a subset of columns or re-order columns. @@ -417,7 +406,7 @@ def __getitem__( # projection? # Select a number of columns as DF. - key = key if _is_list_like(key) else [key] # type:ignore + key = key if utils.is_list_like(key) else [key] # type:ignore selected_ids: Tuple[str, ...] = () for label in key: @@ -763,7 +752,7 @@ def drop( if index: level_id = self._resolve_levels(level or 0)[0] - if _is_list_like(index): + if utils.is_list_like(index): block, inverse_condition_id = block.apply_unary_op( level_id, ops.IsInOp(index, match_nulls=True) ) @@ -778,7 +767,7 @@ def drop( self._block.value_columns ) if columns: - if not _is_list_like(columns): + if not utils.is_list_like(columns): columns = [columns] # type:ignore columns = list(columns) @@ -796,7 +785,7 @@ def reorder_levels(self, order: LevelsType): return DataFrame(self._block.reorder_levels(resolved_level_ids)) def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: - if _is_list_like(level): + if utils.is_list_like(level): levels = list(level) else: levels = [level] @@ -804,7 +793,7 @@ def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: for level_ref in levels: if isinstance(level_ref, int): resolved_level_ids.append(self._block.index_columns[level_ref]) - elif isinstance(level_ref, str): + elif isinstance(level_ref, typing.Hashable): matching_ids = self._block.index_name_to_col_id.get(level_ref, []) if len(matching_ids) != 1: raise ValueError("level name cannot be found or is ambiguous") @@ -827,7 +816,7 @@ def rename_axis( f"rename_axis does not currently support any keyword arguments. {constants.FEEDBACK_LINK}" ) # limited implementation: the new index name is simply the 'mapper' parameter - if _is_list_like(mapper): + if utils.is_list_like(mapper): labels = mapper else: labels = [mapper] @@ -909,7 +898,7 @@ def set_index( append: bool = False, drop: bool = True, ) -> DataFrame: - if not _is_list_like(keys): + if not utils.is_list_like(keys): keys = typing.cast(typing.Sequence[blocks.Label], (keys,)) else: keys = typing.cast(typing.Sequence[blocks.Label], tuple(keys)) @@ -943,7 +932,7 @@ def sort_values( if na_position not in {"first", "last"}: raise ValueError("Param na_position must be one of 'first' or 'last'") - sort_labels = (by,) if isinstance(by, str) else tuple(by) + sort_labels = tuple(by) if utils.is_list_like(by) else (by,) sort_column_ids = self._sql_names(sort_labels) len_by = len(sort_labels) @@ -1118,7 +1107,7 @@ def nunique(self) -> bigframes.series.Series: def agg( self, func: str | typing.Sequence[str] ) -> DataFrame | bigframes.series.Series: - if _is_list_like(func): + if utils.is_list_like(func): if any( dtype not in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES for dtype in self.dtypes @@ -1410,7 +1399,7 @@ def _groupby_series( as_index: bool = True, dropna: bool = True, ): - if not isinstance(by, bigframes.series.Series) and _is_list_like(by): + if not isinstance(by, bigframes.series.Series) and utils.is_list_like(by): by = list(by) else: by = [typing.cast(typing.Union[blocks.Label, bigframes.series.Series], by)] @@ -1704,7 +1693,7 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: id_overrides = { col_id: col_label for col_id, col_label in zip(columns, column_labels) - if col_label + if col_label and isinstance(col_label, str) } if ordering_id is not None: @@ -1766,7 +1755,7 @@ def drop_duplicates( ) -> DataFrame: if subset is None: column_ids = self._block.value_columns - elif _is_list_like(subset): + elif utils.is_list_like(subset): column_ids = [ id for label in subset for id in self._block.label_to_col_id[label] ] @@ -1817,11 +1806,3 @@ def _set_block(self, block: blocks.Block): def _get_block(self) -> blocks.Block: return self._block - - -def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: - return pandas.api.types.is_list_like(obj) - - -def _is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]: - return pandas.api.types.is_dict_like(obj) diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py index 2c3c50ba754..82e2510e2af 100644 --- a/bigframes/formatting_helpers.py +++ b/bigframes/formatting_helpers.py @@ -27,7 +27,9 @@ import bigframes.constants as constants -GenericJob = Union[bigquery.LoadJob, bigquery.ExtractJob, bigquery.QueryJob] +GenericJob = Union[ + bigquery.LoadJob, bigquery.ExtractJob, bigquery.QueryJob, bigquery.CopyJob +] query_job_prop_pairs = { "Job Id": "job_id", @@ -49,7 +51,7 @@ def add_feedback_link( def repr_query_job_html(query_job: Optional[bigquery.QueryJob]): """Return query job in html format. Args: - query_job: + query_job (bigquery.QueryJob, Optional): The job representing the execution of the query on the server. Returns: Pywidget html table. @@ -114,10 +116,12 @@ def wait_for_query_job( ) -> bigquery.table.RowIterator: """Return query results. Displays a progress bar while the query is running Args: - query_job: + query_job (bigquery.QueryJob, Optional): The job representing the execution of the query on the server. - max_results: + max_results (int, Optional): The maximum number of rows the row iterator should return. + progress_bar (str, Optional): + Which progress bar to show. Returns: A row iterator over the query results. """ @@ -154,8 +158,10 @@ def wait_for_query_job( def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): """Waits for job results. Displays a progress bar while the job is running Args: - job: - The bigquery job to be executed + job (GenericJob): + The bigquery job to be executed. + progress_bar (str, Optional): + Which progress bar to show. """ loading_bar = widgets.HTML(get_base_job_loading_html(job)) if progress_bar == "auto": @@ -189,7 +195,7 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): def get_job_url(query_job: GenericJob): """Return url to the query job in cloud console. Args: - query_job: + query_job (GenericJob): The job representing the execution of the query on the server. Returns: String url. @@ -206,7 +212,7 @@ def get_job_url(query_job: GenericJob): def get_query_job_loading_html(query_job: bigquery.QueryJob): """Return progress bar html string Args: - query_job: + query_job (bigquery.QueryJob): The job representing the execution of the query on the server. Returns: Html string. @@ -217,7 +223,7 @@ def get_query_job_loading_html(query_job: bigquery.QueryJob): def get_query_job_loading_string(query_job: bigquery.QueryJob): """Return progress bar string Args: - query_job: + query_job (bigquery.QueryJob): The job representing the execution of the query on the server. Returns: String @@ -228,7 +234,7 @@ def get_query_job_loading_string(query_job: bigquery.QueryJob): def get_base_job_loading_html(job: GenericJob): """Return progress bar html string Args: - job: + job (GenericJob): The job representing the execution of the query on the server. Returns: Html string. @@ -239,7 +245,7 @@ def get_base_job_loading_html(job: GenericJob): def get_base_job_loading_string(job: GenericJob): """Return progress bar string Args: - job: + job (GenericJob): The job representing the execution of the query on the server. Returns: String @@ -250,8 +256,8 @@ def get_base_job_loading_string(job: GenericJob): def get_formatted_time(val): """Try to format time Args: - val: - Time in ms + val (Any): + Time in ms. Returns: Duration string """ diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 6f6efb9f376..9a20fd9ad74 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -73,6 +73,15 @@ def fit( ) return self + @property + def cluster_centers_(self) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError( + "A model must be fitted before calling cluster_centers_." + ) + + return self._bqml_model.centroids() + def predict( self, X: Union[bpd.DataFrame, bpd.Series], diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 9629ca0f4d5..812bb08dc38 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -159,14 +159,22 @@ def evaluate(self, input_data: Optional[bpd.DataFrame] = None): return self._session.read_gbq(sql) + def centroids(self): + assert self._model.model_type == "KMEANS" + + sql = ml_sql.ml_centroids(self.model_name) + + return self._session.read_gbq(sql) + def copy(self, new_model_name: str, replace: bool = False) -> BqmlModel: job_config = bigquery.job.CopyJobConfig() if replace: job_config.write_disposition = "WRITE_TRUNCATE" - self._session.bqclient.copy_table( + copy_job = self._session.bqclient.copy_table( self.model_name, new_model_name, job_config=job_config - ).result() + ) + self._session._start_generic_job(copy_job) new_model = self._session.bqclient.get_model(new_model_name) return BqmlModel(self._session, new_model) @@ -182,7 +190,7 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: options_sql = ml_sql.options(**{"vertex_ai_model_id": vertex_ai_model_id}) sql = ml_sql.alter_model(self.model_name, options_sql=options_sql) # Register the model and wait it to finish - self._session.bqclient.query(sql).result() + self._session._start_query(sql) self._model = self._session.bqclient.get_model(self.model_name) return self @@ -322,8 +330,7 @@ def _create_temp_model_name() -> str: def _create_bqml_model_with_sql(session: bigframes.Session, sql: str) -> BqmlModel: # fit the model, synchronously - job = session.bqclient.query(sql) - job.result() + _, job = session._start_query(sql) # real model path in the session specific hidden dataset and table prefix model_name_full = f"{job.destination.dataset_id}.{job.destination.table_id}" diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index b8d9e2c6737..80054d40e19 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -146,6 +146,13 @@ def ml_evaluate(model_name: str, source_sql: Union[str, None] = None) -> str: ({source_sql}))""" +def ml_centroids( + model_name: str, +) -> str: + """Encode ML.CENTROIDS for BQML""" + return f"""SELECT * FROM ML.CENTROIDS(MODEL `{model_name}`)""" + + def ml_predict(model_name: str, source_sql: str) -> str: """Encode ML.PREDICT for BQML""" return f"""SELECT * FROM ML.PREDICT(MODEL `{model_name}`, diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 13063af75f9..7bdd97812e9 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -134,6 +134,42 @@ def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.StringValue, x).re_search(self._pat) +class StrGetOp(UnaryOp): + def __init__(self, i: int): + self._i = i + + def _as_ibis(self, x: ibis_types.Value): + substr = typing.cast( + ibis_types.StringValue, typing.cast(ibis_types.StringValue, x)[self._i] + ) + return substr.nullif(ibis_types.literal("")) + + +class StrPadOp(UnaryOp): + def __init__( + self, length: int, fillchar: str, side: typing.Literal["both", "left", "right"] + ): + self._length = length + self._fillchar = fillchar + self._side = side + + def _as_ibis(self, x: ibis_types.Value): + str_val = typing.cast(ibis_types.StringValue, x) + + # SQL pad operations will truncate, we do not want to truncate though. + pad_length = ibis.greatest(str_val.length(), self._length) + if self._side == "left": + return str_val.lpad(pad_length, self._fillchar) + elif self._side == "right": + return str_val.rpad(pad_length, self._fillchar) + else: # side == both + # Pad more on right side if can't pad both sides equally + lpad_amount = ((pad_length - str_val.length()) // 2) + str_val.length() + return str_val.lpad(lpad_amount, self._fillchar).rpad( + pad_length, self._fillchar + ) + + class ReplaceStringOp(UnaryOp): def __init__(self, pat: str, repl: str): self._pat = pat diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index caef33919b3..361fdca0558 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -63,9 +63,9 @@ def __init__( if block: if name: - if not isinstance(name, str): - raise NotImplementedError( - f"BigQuery DataFrames only supports string series names. {constants.FEEDBACK_LINK}" + if not isinstance(name, typing.Hashable): + raise ValueError( + f"BigQuery DataFrames only supports hashable series names. {constants.FEEDBACK_LINK}" ) block = block.with_column_labels([name]) if index: diff --git a/bigframes/operations/strings.py b/bigframes/operations/strings.py index a16ecb0d328..0f1395c78f2 100644 --- a/bigframes/operations/strings.py +++ b/bigframes/operations/strings.py @@ -82,6 +82,29 @@ def repeat(self, repeats: int) -> series.Series: def capitalize(self) -> series.Series: return self._apply_unary_op(ops.capitalize_op) + def match(self, pat, case=True, flags=0) -> series.Series: + # \A anchors start of entire string rather than start of any line in multiline mode + adj_pat = rf"\A{pat}" + return self.contains(adj_pat, case=case, flags=flags) + + def fullmatch(self, pat, case=True, flags=0) -> series.Series: + # \A anchors start of entire string rather than start of any line in multiline mode + # \z likewise anchors to the end of the entire multiline string + adj_pat = rf"\A{pat}\z" + return self.contains(adj_pat, case=case, flags=flags) + + def get(self, i: int) -> series.Series: + return self._apply_unary_op(ops.StrGetOp(i)) + + def pad(self, width, side="left", fillchar=" ") -> series.Series: + return self._apply_unary_op(ops.StrPadOp(width, fillchar, side)) + + def ljust(self, width, fillchar=" ") -> series.Series: + return self._apply_unary_op(ops.StrPadOp(width, fillchar, "right")) + + def rjust(self, width, fillchar=" ") -> series.Series: + return self._apply_unary_op(ops.StrPadOp(width, fillchar, "left")) + def contains( self, pat, case: bool = True, flags: int = 0, *, regex: bool = True ) -> series.Series: @@ -95,7 +118,7 @@ def contains( else: return self._apply_unary_op(ops.ContainsStringOp(pat)) - def extract(self, pat: str, flags: int = 0): + def extract(self, pat: str, flags: int = 0) -> df.DataFrame: re2flags = _parse_flags(flags) if re2flags: pat = re2flags + pat diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index cc8b4e5cc4d..ed7a09e7b71 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections import namedtuple import inspect import threading import typing @@ -399,6 +400,9 @@ def read_gbq_function(function_name: str): Index = bigframes.core.indexes.Index Series = bigframes.series.Series +# Used by DataFrameGroupby.agg +NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"]) + # Use __all__ to let type checkers know what is part of the public API. __all___ = [ "concat", @@ -411,4 +415,5 @@ def read_gbq_function(function_name: str): "read_pandas", "remote_function", "Series", + "NamedAgg", ] diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 5924941cc54..27e2b8f7c23 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -256,34 +256,6 @@ def create_bq_connection(self): def check_bq_connection_exists(self): """Check if the BigQuery Connection exists.""" client = self._bq_connection_client - if self._bq_connection_id.count(".") == 1: - bq_location, bq_connection_id = self._bq_connection_id.split(".") - if bq_location != self._bq_location: - logger.info( - f"Reset location {self._bq_location} to match the" - + f"location in connection name: {bq_location}" - ) - self._bq_location = bq_location - self._bq_connection_id = bq_connection_id - elif self._bq_connection_id.count(".") == 2: - ( - gcp_project_id, - bq_location, - bq_connection_id, - ) = self._bq_connection_id.split(".") - if gcp_project_id != self._gcp_project_id: - raise ValueError( - "The project_id does not match BigQuery connection gcp_project_id: " - f"{self._gcp_project_id}." - ) - if bq_location != self._bq_location: - logger.info( - f"Reset location {self._bq_location} to match the" - + f"location in connection name: {bq_location}" - ) - self._gcp_project_id = gcp_project_id - self._bq_location = bq_location - self._bq_connection_id = bq_connection_id request = bigquery_connection_v1.GetConnectionRequest( name=client.connection_path( self._gcp_project_id, self._bq_location, self._bq_connection_id @@ -797,6 +769,33 @@ def remote_function( random.choices(string.ascii_lowercase + string.digits, k=8) ) + # Check connection_id with `LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID` format. + if bigquery_connection.count(".") == 1: + bq_connection_location, bq_connection_id = bigquery_connection.split(".") + if bq_connection_location.casefold() != bq_location.casefold(): + raise ValueError( + "The location does not match BigQuery connection location: " + f"{bq_location}." + ) + bigquery_connection = bq_connection_id + elif bigquery_connection.count(".") == 2: + ( + gcp_project_id, + bq_connection_location, + bq_connection_id, + ) = bigquery_connection.split(".") + if gcp_project_id.casefold() != dataset_ref.project.casefold(): + raise ValueError( + "The project_id does not match BigQuery connection gcp_project_id: " + f"{dataset_ref.project}." + ) + if bq_connection_location.casefold() != bq_location.casefold(): + raise ValueError( + "The location does not match BigQuery connection location: " + f"{bq_location}." + ) + bigquery_connection = bq_connection_id + def wrapper(f): if not callable(f): raise TypeError("f must be callable, got {}".format(f)) diff --git a/bigframes/series.py b/bigframes/series.py index 537991ed004..f8f44dc2e64 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -87,7 +87,7 @@ def iloc(self) -> bigframes.core.indexers.IlocSeriesIndexer: return bigframes.core.indexers.IlocSeriesIndexer(self) @property - def name(self) -> Optional[str]: + def name(self) -> blocks.Label: return self._name @property @@ -169,6 +169,12 @@ def rename( return Series(block) + # rename the Series name + if isinstance(index, typing.Hashable): + index = typing.cast(Optional[str], index) + block = self._block.with_column_labels([index]) + return Series(block) + raise ValueError(f"Unsupported type of parameter index: {type(index)}") def rename_axis( @@ -321,7 +327,7 @@ def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: for level_ref in levels: if isinstance(level_ref, int): resolved_level_ids.append(self._block.index_columns[level_ref]) - elif isinstance(level_ref, str): + elif isinstance(level_ref, typing.Hashable): matching_ids = self._block.index_name_to_col_id.get(level_ref, []) if len(matching_ids) != 1: raise ValueError("level name cannot be found or is ambiguous") @@ -440,6 +446,17 @@ def nsmallest(self, n: int = 5, keep: str = "first") -> Series: block = block.select_column(self._value_column) return Series(block) + def isin(self, values) -> "Series" | None: + if not _is_list_like(values): + raise TypeError( + "only list-like objects are allowed to be passed to " + f"isin(), you passed a [{type(values).__name__}]" + ) + + return self._apply_unary_op(ops.IsInOp(values, match_nulls=True)).fillna( + value=False + ) + def isna(self) -> "Series": return self._apply_unary_op(ops.isnull_op) @@ -791,6 +808,26 @@ def argmin(self) -> scalars.Scalar: scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] ) + @property + def is_monotonic_increasing(self) -> bool: + period = 1 + window = bigframes.core.WindowSpec( + preceding=period, + following=None, + ) + shifted_series = self._apply_window_op(agg_ops.ShiftOp(period), window) + return self.notna().__and__(self >= shifted_series).all() + + @property + def is_monotonic_decreasing(self) -> bool: + period = 1 + window = bigframes.core.WindowSpec( + preceding=period, + following=None, + ) + shifted_series = self._apply_window_op(agg_ops.ShiftOp(period), window) + return self.notna().__and__(self <= shifted_series).all() + def __getitem__(self, indexer): # TODO: enforce stricter alignment, should fail if indexer is missing any keys. use_iloc = ( diff --git a/bigframes/session.py b/bigframes/session.py index 28a38f93073..73fdd731063 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -358,10 +358,16 @@ def read_gbq_query( else: index_cols = list(index_col) - # Can't cluster since don't know if index_cols are clusterable data types - # TODO(tbergeron): Maybe use dryrun to determine types of index_cols to see if can cluster - _, query_job = self._start_query(query) - destination = query_job.destination + # Make sure we cluster by the index column so that subsequent + # operations are as speedy as they can be. + if index_cols: + # Since index_cols are specified, assume that we have a normal SQL + # query. DDL or DML not supported. + ibis_expr = self.ibis_client.sql(query) + destination = self._ibis_to_session_table(ibis_expr, index_cols) + else: + _, query_job = self._start_query(query) + destination = query_job.destination # If there was no destination table, that means the query must have # been DDL or DML. Return some job metadata, instead. @@ -936,15 +942,9 @@ def _create_sequential_ordering( ibis.row_number().cast(ibis_dtypes.int64).name(default_ordering_name) ) table = table.mutate(**{default_ordering_name: default_ordering_col}) - clusterable_index_cols = [ - col for col in index_cols if _can_cluster(table[col].type()) - ] - cluster_cols = (clusterable_index_cols + [default_ordering_name])[ - :_MAX_CLUSTER_COLUMNS - ] - table_ref = self._query_to_session_table( - self.ibis_client.compile(table), - cluster_cols=cluster_cols, + table_ref = self._ibis_to_session_table( + table, + cluster_cols=list(index_cols) + [default_ordering_name], ) table = self.ibis_client.sql(f"SELECT * FROM `{table_ref.table_id}`") ordering_reference = core.OrderingColumnReference(default_ordering_name) @@ -955,6 +955,17 @@ def _create_sequential_ordering( ) return table, ordering + def _ibis_to_session_table( + self, table: ibis_types.Table, cluster_cols: Iterable[str] + ) -> bigquery.TableReference: + clusterable_cols = [ + col for col in cluster_cols if _can_cluster(table[col].type()) + ][:_MAX_CLUSTER_COLUMNS] + return self._query_to_session_table( + self.ibis_client.compile(table), + cluster_cols=clusterable_cols, + ) + def _query_to_session_table( self, query_text: str, cluster_cols: Iterable[str] ) -> bigquery.TableReference: diff --git a/bigframes/version.py b/bigframes/version.py index c387f59d2c6..d386742d59f 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.1.1" +__version__ = "0.2.0" diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml new file mode 100644 index 00000000000..91abc59bc0a --- /dev/null +++ b/docs/templates/toc.yml @@ -0,0 +1,130 @@ +- items: + - href: index.md + name: Overview + - href: changelog.md + name: Changelog + - items: + - items: + - name: Options + uid: bigframes._config.Options + - name: BigQueryOptions + uid: bigframes._config.bigquery_options.BigQueryOptions + - name: DisplayOptions + uid: bigframes._config.display_options.DisplayOptions + - name: SamplingOptions + uid: bigframes._config.sampling_options.SamplingOptions + name: Options and settings + - items: + - name: Session + uid: bigframes.session.Session + name: Session + name: Core Objects + - items: + - name: DataFrame + uid: bigframes.dataframe.DataFrame + - items: + - name: DataFrameGroupBy + uid: bigframes.core.groupby.DataFrameGroupBy + - name: SeriesGroupBy + uid: bigframes.core.groupby.SeriesGroupBy + name: Groupby + - name: Indexes + uid: bigframes.core.indexes.index.Index + - name: pandas + uid: bigframes.pandas + - items: + - name: Series + uid: bigframes.series.Series + - name: DatetimeMethods + uid: bigframes.operations.datetimes.DatetimeMethods + - name: StringMethods + uid: bigframes.operations.strings.StringMethods + name: Series + - name: Window + uid: bigframes.core.window.Window + name: bigframes.pandas + - items: + - items: + - name: Overview + uid: bigframes.ml.cluster + - name: KMeans + uid: bigframes.ml.cluster.KMeans + name: cluster + - items: + - name: Overview + uid: bigframes.ml.compose + - name: ColumnTransformer + uid: bigframes.ml.compose.ColumnTransformer + name: compose + - items: + - name: Overview + uid: bigframes.ml.decomposition + - name: PCA + uid: bigframes.ml.decomposition.PCA + name: decomposition + - items: + - name: Overview + uid: bigframes.ml.ensemble + - name: RandomForestClassifier + uid: bigframes.ml.ensemble.RandomForestClassifier + - name: RandomForestRegressor + uid: bigframes.ml.ensemble.RandomForestRegressor + - name: XGBClassifier + uid: bigframes.ml.ensemble.XGBClassifier + - name: XGBRegressor + uid: bigframes.ml.ensemble.XGBRegressor + name: ensemble + - items: + - name: Overview + uid: bigframes.ml.forecasting + - name: ARIMAPlus + uid: bigframes.ml.forecasting.ARIMAPlus + name: forecasting + - items: + - name: Overview + uid: bigframes.ml.imported + - name: ONNXModel + uid: bigframes.ml.imported.ONNXModel + - name: TensorFlowModel + uid: bigframes.ml.imported.TensorFlowModel + name: imported + - items: + - name: Overview + uid: bigframes.ml.linear_model + - name: LinearRegression + uid: bigframes.ml.linear_model.LinearRegression + - name: LogisticRegression + uid: bigframes.ml.linear_model.LogisticRegression + name: linear_model + - items: + - name: Overview + uid: bigframes.ml.llm + - name: PaLM2TextGenerator + uid: bigframes.ml.llm.PaLM2TextGenerator + - name: PaLM2TextEmbeddingGenerator + uid: bigframes.ml.llm.PaLM2TextEmbeddingGenerator + name: llm + - items: + - name: metrics + uid: bigframes.ml.metrics + name: metrics + - items: + - name: model_selection + uid: bigframes.ml.model_selection + name: model_selection + - items: + - name: Overview + uid: bigframes.ml.pipeline + - name: Pipeline + uid: bigframes.ml.pipeline.Pipeline + name: pipeline + - items: + - name: Overview + uid: bigframes.ml.preprocessing + - name: OneHotEncoder + uid: bigframes.ml.preprocessing.OneHotEncoder + - name: StandardScaler + uid: bigframes.ml.preprocessing.StandardScaler + name: preprocessing + name: bigframes.ml + name: BigQuery DataFrames diff --git a/notebooks/03 - Using ML - ML fundamentals.ipynb b/notebooks/03 - Using ML - ML fundamentals.ipynb index c181b67e647..ade50bcbc2e 100644 --- a/notebooks/03 - Using ML - ML fundamentals.ipynb +++ b/notebooks/03 - Using ML - ML fundamentals.ipynb @@ -14,9 +14,65 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "590f15dab2c940b0a23af16c4b1b8fe1", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 802292ea-3f16-400e-97e2-34009774a1f0 is RUNNING. \n", " \n", " \n", - " 125\n", - " 3550.0\n", + " 241\n", + " 5300.0\n", " \n", " \n", - " 260\n", - " 5750.0\n", + " 121\n", + " 3325.0\n", " \n", " \n", - " 93\n", - " 4300.0\n", + " 209\n", + " 4075.0\n", " \n", " \n", - " 255\n", - " 4850.0\n", + " 270\n", + " 3075.0\n", " \n", " \n", - " 71\n", - " 3050.0\n", + " 187\n", + " 4600.0\n", " \n", " \n", "\n", + "

5 rows × 1 columns

\n", "[5 rows x 1 columns in total]" ], "text/plain": [ " body_mass_g\n", "penguin_id \n", - "125 3550.0\n", - "260 5750.0\n", - "93 4300.0\n", - "255 4850.0\n", - "71 3050.0\n", + "241 5300.0\n", + "121 3325.0\n", + "209 4075.0\n", + "270 3075.0\n", + "187 4600.0\n", "\n", "[5 rows x 1 columns]" ] @@ -648,6 +911,76 @@ "execution_count": 5, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "299c0c895e3d4a83a8495924a1966ce0", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job f54865db-fdb4-4022-af30-7f282a6b81c0 is DONE. 31.7 kB processed. \n", " \n", " \n", - " 23\n", - " -1.310231\n", - " 0.687253\n", - " -0.568262\n", + " 0\n", + " -1.364965\n", + " 0.629892\n", + " -1.226537\n", " \n", " \n", - " 83\n", - " 2.575955\n", - " 0.329548\n", - " -1.432429\n", + " 1\n", + " -0.771824\n", + " 0.984275\n", + " -1.226537\n", + " \n", + " \n", + " 2\n", + " -0.567932\n", + " 0.883023\n", + " -1.226537\n", + " \n", + " \n", + " 3\n", + " 0.470064\n", + " 0.376761\n", + " -0.652517\n", " \n", " \n", - " 195\n", - " 0.816172\n", - " -1.407876\n", - " 0.151877\n", + " 4\n", + " -1.235216\n", + " -0.180128\n", + " -0.652517\n", " \n", " \n", - " 264\n", - " 1.109469\n", - " -0.436962\n", - " 2.096254\n", + " 5\n", + " -0.141612\n", + " 0.680518\n", + " -0.652517\n", " \n", " \n", - " 237\n", - " 0.816172\n", - " -1.305674\n", - " 0.872017\n", + " 6\n", + " 0.544207\n", + " -0.281381\n", + " -0.652517\n", " \n", " \n", - " 261\n", - " 0.614531\n", - " -0.947969\n", - " 1.520142\n", + " 7\n", + " 1.21149\n", + " 0.629892\n", + " -0.078497\n", " \n", " \n", - " 253\n", - " 0.321234\n", - " -1.663379\n", - " 0.944031\n", + " 8\n", + " 1.026133\n", + " 0.933649\n", + " -0.078497\n", " \n", " \n", - " 164\n", - " 0.101261\n", - " -0.743566\n", - " 1.160073\n", + " 10\n", + " -0.586468\n", + " 0.883023\n", + " 0.495523\n", " \n", " \n", - " 293\n", - " -0.686975\n", - " -0.079257\n", - " -1.792499\n", + " 11\n", + " -0.92011\n", + " 0.781771\n", + " -1.154784\n", " \n", " \n", - " 176\n", - " 0.871166\n", - " -1.561178\n", - " 0.655975\n", + " 14\n", + " -1.846892\n", + " -0.028249\n", + " -1.154784\n", " \n", " \n", - " 45\n", - " 1.476091\n", - " 1.811469\n", - " 0.655975\n", + " 15\n", + " -1.290822\n", + " -0.332007\n", + " -1.154784\n", " \n", " \n", - " 185\n", - " 1.347773\n", - " -1.510077\n", - " 1.232087\n", + " 16\n", + " 0.321779\n", + " 0.073003\n", + " -0.580765\n", " \n", " \n", - " 141\n", - " 1.036145\n", - " -1.101271\n", - " 1.088059\n", + " 17\n", + " 1.230026\n", + " 1.136154\n", + " -0.580765\n", " \n", " \n", - " 147\n", - " 1.109469\n", - " -0.641365\n", - " 1.66417\n", + " 18\n", + " -0.79036\n", + " 0.376761\n", + " -0.580765\n", " \n", " \n", - " 320\n", - " -0.925279\n", - " 0.329548\n", - " -1.432429\n", + " 20\n", + " 0.4886\n", + " 0.326134\n", + " -0.580765\n", " \n", " \n", - " 285\n", - " -0.192036\n", - " -2.072185\n", - " 1.016045\n", + " 21\n", + " 1.359775\n", + " 1.034902\n", + " -0.580765\n", " \n", " \n", - " 305\n", - " 0.027937\n", - " 0.43175\n", - " 0.655975\n", + " 24\n", + " 1.044669\n", + " 0.528639\n", + " -0.580765\n", " \n", " \n", - " 37\n", - " 1.421098\n", - " 1.607066\n", - " -0.496248\n", + " 26\n", + " -0.456718\n", + " 0.680518\n", + " -0.006745\n", " \n", " \n", - " 98\n", - " -0.521995\n", - " 0.176246\n", - " -0.784304\n", + " 27\n", + " 1.21149\n", + " 1.237407\n", + " -0.006745\n", " \n", " \n", - " 10\n", - " -0.576989\n", - " 0.891656\n", - " 0.511947\n", + " 29\n", + " 1.378311\n", + " 0.933649\n", + " -0.006745\n", + " \n", + " \n", + " 31\n", + " -2.013713\n", + " -0.534512\n", + " -1.657052\n", + " \n", + " \n", + " 32\n", + " -1.253751\n", + " 0.478013\n", + " -1.657052\n", + " \n", + " \n", + " 33\n", + " -0.827431\n", + " -0.230754\n", + " -1.657052\n", " \n", " \n", "\n", + "

25 rows × 3 columns

\n", "[267 rows x 3 columns in total]" ], "text/plain": [ " scaled_culmen_length_mm scaled_culmen_depth_mm \\\n", "penguin_id \n", - "23 -1.310231 0.687253 \n", - "83 2.575955 0.329548 \n", - "195 0.816172 -1.407876 \n", - "264 1.109469 -0.436962 \n", - "237 0.816172 -1.305674 \n", - "261 0.614531 -0.947969 \n", - "253 0.321234 -1.663379 \n", - "164 0.101261 -0.743566 \n", - "293 -0.686975 -0.079257 \n", - "176 0.871166 -1.561178 \n", - "45 1.476091 1.811469 \n", - "185 1.347773 -1.510077 \n", - "141 1.036145 -1.101271 \n", - "147 1.109469 -0.641365 \n", - "320 -0.925279 0.329548 \n", - "285 -0.192036 -2.072185 \n", - "305 0.027937 0.43175 \n", - "37 1.421098 1.607066 \n", - "98 -0.521995 0.176246 \n", - "10 -0.576989 0.891656 \n", - "307 -1.420218 0.022944 \n", - "337 -0.393678 1.14716 \n", - "26 -0.448671 0.687253 \n", - "97 -1.45688 0.380649 \n", - "27 1.201125 1.249361 \n", + "0 -1.364965 0.629892 \n", + "1 -0.771824 0.984275 \n", + "2 -0.567932 0.883023 \n", + "3 0.470064 0.376761 \n", + "4 -1.235216 -0.180128 \n", + "5 -0.141612 0.680518 \n", + "6 0.544207 -0.281381 \n", + "7 1.21149 0.629892 \n", + "8 1.026133 0.933649 \n", + "10 -0.586468 0.883023 \n", + "11 -0.92011 0.781771 \n", + "14 -1.846892 -0.028249 \n", + "15 -1.290822 -0.332007 \n", + "16 0.321779 0.073003 \n", + "17 1.230026 1.136154 \n", + "18 -0.79036 0.376761 \n", + "20 0.4886 0.326134 \n", + "21 1.359775 1.034902 \n", + "24 1.044669 0.528639 \n", + "26 -0.456718 0.680518 \n", + "27 1.21149 1.237407 \n", + "29 1.378311 0.933649 \n", + "31 -2.013713 -0.534512 \n", + "32 -1.253751 0.478013 \n", + "33 -0.827431 -0.230754 \n", "\n", " scaled_flipper_length_mm \n", "penguin_id \n", - "23 -0.568262 \n", - "83 -1.432429 \n", - "195 0.151877 \n", - "264 2.096254 \n", - "237 0.872017 \n", - "261 1.520142 \n", - "253 0.944031 \n", - "164 1.160073 \n", - "293 -1.792499 \n", - "176 0.655975 \n", - "45 0.655975 \n", - "185 1.232087 \n", - "141 1.088059 \n", - "147 1.66417 \n", - "320 -1.432429 \n", - "285 1.016045 \n", - "305 0.655975 \n", - "37 -0.496248 \n", - "98 -0.784304 \n", - "10 0.511947 \n", - "307 -1.000346 \n", - "337 -0.208192 \n", - "26 0.00785 \n", - "97 -0.784304 \n", - "27 0.00785 \n", + "0 -1.226537 \n", + "1 -1.226537 \n", + "2 -1.226537 \n", + "3 -0.652517 \n", + "4 -0.652517 \n", + "5 -0.652517 \n", + "6 -0.652517 \n", + "7 -0.078497 \n", + "8 -0.078497 \n", + "10 0.495523 \n", + "11 -1.154784 \n", + "14 -1.154784 \n", + "15 -1.154784 \n", + "16 -0.580765 \n", + "17 -0.580765 \n", + "18 -0.580765 \n", + "20 -0.580765 \n", + "21 -0.580765 \n", + "24 -0.580765 \n", + "26 -0.006745 \n", + "27 -0.006745 \n", + "29 -0.006745 \n", + "31 -1.657052 \n", + "32 -1.657052 \n", + "33 -1.657052 \n", "...\n", "\n", "[267 rows x 3 columns]" @@ -890,6 +1254,76 @@ "execution_count": 6, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "6ec4ab3e60844e989dbebad89e7665ca", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job a8029b51-2ef1-4acd-9759-d808db954298 is DONE. 31.7 kB processed. \n", " \n", " \n", - " 125\n", - " -0.77863\n", - " 0.891656\n", - " -1.216387\n", + " 9\n", + " -0.697682\n", + " 1.490538\n", + " -0.078497\n", " \n", " \n", - " 260\n", - " 1.182794\n", - " -0.743566\n", - " 1.520142\n", + " 12\n", + " -1.290822\n", + " -0.129502\n", + " -1.154784\n", " \n", " \n", - " 93\n", - " -0.521995\n", - " 0.482851\n", - " 0.295905\n", + " 13\n", + " 0.562742\n", + " 0.073003\n", + " -1.154784\n", " \n", " \n", - " 255\n", - " 0.101261\n", - " -1.254574\n", - " 0.944031\n", + " 19\n", + " -1.142537\n", + " 0.478013\n", + " -0.580765\n", " \n", " \n", - " 71\n", - " -2.171791\n", - " -0.845768\n", - " -0.928332\n", + " 22\n", + " -0.697682\n", + " -0.028249\n", + " -0.580765\n", " \n", " \n", - " 327\n", - " -0.521995\n", - " 0.227347\n", - " -1.360415\n", + " 23\n", + " -1.327894\n", + " 0.680518\n", + " -0.580765\n", " \n", " \n", - " 338\n", - " -0.980272\n", - " 2.066973\n", - " -0.71229\n", + " 25\n", + " 1.489525\n", + " 0.478013\n", + " -0.006745\n", " \n", " \n", - " 68\n", - " 1.29278\n", - " 0.840556\n", - " 0.151877\n", + " 28\n", + " 1.897309\n", + " 1.844922\n", + " -0.006745\n", " \n", " \n", - " 17\n", - " 1.219456\n", - " 1.14716\n", - " -0.568262\n", + " 30\n", + " 1.267097\n", + " 0.680518\n", + " -0.006745\n", " \n", " \n", - " 116\n", - " -0.302023\n", - " 2.066973\n", - " -0.71229\n", + " 38\n", + " -0.49379\n", + " 1.59179\n", + " -0.509012\n", " \n", " \n", - " 106\n", - " -0.94361\n", - " 1.453764\n", - " -0.784304\n", + " 41\n", + " -1.531786\n", + " 0.427387\n", + " 0.065008\n", " \n", " \n", - " 227\n", - " 1.512753\n", - " -0.028156\n", - " 1.952226\n", + " 45\n", + " 1.489525\n", + " 1.794295\n", + " 0.639028\n", " \n", " \n", - " 314\n", - " -0.210367\n", - " 0.687253\n", - " -0.424234\n", + " 49\n", + " -0.567932\n", + " -0.281381\n", + " -1.011279\n", " \n", " \n", - " 246\n", - " -1.126921\n", - " 0.585052\n", - " -1.936527\n", + " 50\n", + " -0.271362\n", + " -0.230754\n", + " -1.011279\n", " \n", " \n", - " 138\n", - " 1.201125\n", - " -0.99907\n", - " 1.088059\n", + " 62\n", + " -1.402036\n", + " -0.078876\n", + " -0.43726\n", " \n", " \n", - " 126\n", - " -1.383556\n", - " -0.028156\n", - " -1.216387\n", + " 65\n", + " -0.586468\n", + " 0.629892\n", + " -0.43726\n", " \n", " \n", - " 121\n", - " 0.761179\n", - " -0.385861\n", - " -0.136178\n", + " 66\n", + " 0.989062\n", + " 1.389285\n", + " 0.13676\n", " \n", " \n", - " 231\n", - " -0.320354\n", - " 1.198261\n", - " -0.280206\n", + " 83\n", + " 2.601663\n", + " 0.326134\n", + " -1.441794\n", " \n", " \n", - " 40\n", - " 1.146132\n", - " 0.840556\n", - " 0.079863\n", + " 87\n", + " 1.526596\n", + " 0.832397\n", + " -0.293755\n", " \n", " \n", - " 60\n", - " 0.504544\n", - " 0.380649\n", - " -0.424234\n", + " 92\n", + " 1.767559\n", + " 1.389285\n", + " 0.280265\n", " \n", - " \n", - "\n", - "[67 rows x 3 columns in total]" - ], + " \n", + " 98\n", + " -0.530861\n", + " 0.174256\n", + " -0.796022\n", + " \n", + " \n", + " 104\n", + " 0.35885\n", + " -0.028249\n", + " -0.796022\n", + " \n", + " \n", + " 114\n", + " 0.229101\n", + " -0.281381\n", + " -0.72427\n", + " \n", + " \n", + " 115\n", + " 0.840777\n", + " 0.174256\n", + " -0.72427\n", + " \n", + " \n", + " 118\n", + " -1.550322\n", + " 0.174256\n", + " -0.72427\n", + " \n", + " \n", + "\n", + "

25 rows × 3 columns

\n", + "[67 rows x 3 columns in total]" + ], "text/plain": [ " scaled_culmen_length_mm scaled_culmen_depth_mm \\\n", "penguin_id \n", - "125 -0.77863 0.891656 \n", - "260 1.182794 -0.743566 \n", - "93 -0.521995 0.482851 \n", - "255 0.101261 -1.254574 \n", - "71 -2.171791 -0.845768 \n", - "327 -0.521995 0.227347 \n", - "338 -0.980272 2.066973 \n", - "68 1.29278 0.840556 \n", - "17 1.219456 1.14716 \n", - "116 -0.302023 2.066973 \n", - "106 -0.94361 1.453764 \n", - "227 1.512753 -0.028156 \n", - "314 -0.210367 0.687253 \n", - "246 -1.126921 0.585052 \n", - "138 1.201125 -0.99907 \n", - "126 -1.383556 -0.028156 \n", - "121 0.761179 -0.385861 \n", - "231 -0.320354 1.198261 \n", - "40 1.146132 0.840556 \n", - "60 0.504544 0.380649 \n", - "331 -1.548535 0.176246 \n", - "218 0.192916 -0.896869 \n", - "154 -1.10859 0.738354 \n", - "130 -0.357016 1.198261 \n", - "89 1.604408 1.351563 \n", + "9 -0.697682 1.490538 \n", + "12 -1.290822 -0.129502 \n", + "13 0.562742 0.073003 \n", + "19 -1.142537 0.478013 \n", + "22 -0.697682 -0.028249 \n", + "23 -1.327894 0.680518 \n", + "25 1.489525 0.478013 \n", + "28 1.897309 1.844922 \n", + "30 1.267097 0.680518 \n", + "38 -0.49379 1.59179 \n", + "41 -1.531786 0.427387 \n", + "45 1.489525 1.794295 \n", + "49 -0.567932 -0.281381 \n", + "50 -0.271362 -0.230754 \n", + "62 -1.402036 -0.078876 \n", + "65 -0.586468 0.629892 \n", + "66 0.989062 1.389285 \n", + "83 2.601663 0.326134 \n", + "87 1.526596 0.832397 \n", + "92 1.767559 1.389285 \n", + "98 -0.530861 0.174256 \n", + "104 0.35885 -0.028249 \n", + "114 0.229101 -0.281381 \n", + "115 0.840777 0.174256 \n", + "118 -1.550322 0.174256 \n", "\n", " scaled_flipper_length_mm \n", "penguin_id \n", - "125 -1.216387 \n", - "260 1.520142 \n", - "93 0.295905 \n", - "255 0.944031 \n", - "71 -0.928332 \n", - "327 -1.360415 \n", - "338 -0.71229 \n", - "68 0.151877 \n", - "17 -0.568262 \n", - "116 -0.71229 \n", - "106 -0.784304 \n", - "227 1.952226 \n", - "314 -0.424234 \n", - "246 -1.936527 \n", - "138 1.088059 \n", - "126 -1.216387 \n", - "121 -0.136178 \n", - "231 -0.280206 \n", - "40 0.079863 \n", - "60 -0.424234 \n", - "331 -0.784304 \n", - "218 1.376114 \n", - "154 -0.568262 \n", - "130 -0.064164 \n", - "89 -0.280206 \n", + "9 -0.078497 \n", + "12 -1.154784 \n", + "13 -1.154784 \n", + "19 -0.580765 \n", + "22 -0.580765 \n", + "23 -0.580765 \n", + "25 -0.006745 \n", + "28 -0.006745 \n", + "30 -0.006745 \n", + "38 -0.509012 \n", + "41 0.065008 \n", + "45 0.639028 \n", + "49 -1.011279 \n", + "50 -1.011279 \n", + "62 -0.43726 \n", + "65 -0.43726 \n", + "66 0.13676 \n", + "83 -1.441794 \n", + "87 -0.293755 \n", + "92 0.280265 \n", + "98 -0.796022 \n", + "104 -0.796022 \n", + "114 -0.72427 \n", + "115 -0.72427 \n", + "118 -0.72427 \n", "...\n", "\n", "[67 rows x 3 columns]" @@ -1133,6 +1598,104 @@ "execution_count": 7, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "94efb3f362234593906e44c16b1842e4", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 792a19f4-e3c2-4344-8555-7f5a781c3c7d is DONE. 32.0 kB processed. \n", " \n", " \n", - " 23\n", + " 0\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.364965\n", + " 0.629892\n", + " -1.226537\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.310231\n", - " 0.687253\n", - " -0.568262\n", " [{'index': 1, 'value': 1.0}]\n", + " \n", + " \n", + " 1\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.771824\n", + " 0.984275\n", + " -1.226537\n", + " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 83\n", + " 2\n", " [{'index': 2, 'value': 1.0}]\n", - " 2.575955\n", - " 0.329548\n", - " -1.432429\n", + " -0.567932\n", + " 0.883023\n", + " -1.226537\n", + " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", + " \n", + " \n", + " 3\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.470064\n", + " 0.376761\n", + " -0.652517\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 195\n", - " [{'index': 1, 'value': 1.0}]\n", - " 0.816172\n", - " -1.407876\n", - " 0.151877\n", + " 4\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.235216\n", + " -0.180128\n", + " -0.652517\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", - " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 264\n", - " [{'index': 1, 'value': 1.0}]\n", - " 1.109469\n", - " -0.436962\n", - " 2.096254\n", + " 5\n", " [{'index': 2, 'value': 1.0}]\n", + " -0.141612\n", + " 0.680518\n", + " -0.652517\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 237\n", - " [{'index': 1, 'value': 1.0}]\n", - " 0.816172\n", - " -1.305674\n", - " 0.872017\n", + " 6\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.544207\n", + " -0.281381\n", + " -0.652517\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 261\n", - " [{'index': 1, 'value': 1.0}]\n", - " 0.614531\n", - " -0.947969\n", - " 1.520142\n", + " 7\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.21149\n", + " 0.629892\n", + " -0.078497\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 253\n", - " [{'index': 1, 'value': 1.0}]\n", - " 0.321234\n", - " -1.663379\n", - " 0.944031\n", - " [{'index': 1, 'value': 1.0}]\n", + " 8\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.026133\n", + " 0.933649\n", + " -0.078497\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 164\n", - " [{'index': 1, 'value': 1.0}]\n", - " 0.101261\n", - " -0.743566\n", - " 1.160073\n", - " [{'index': 0, 'value': 1.0}]\n", + " 10\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.586468\n", + " 0.883023\n", + " 0.495523\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 293\n", + " 11\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.92011\n", + " 0.781771\n", + " -1.154784\n", " [{'index': 3, 'value': 1.0}]\n", - " -0.686975\n", - " -0.079257\n", - " -1.792499\n", - " [{'index': 1, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 176\n", + " 14\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.846892\n", + " -0.028249\n", + " -1.154784\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", - " 0.871166\n", - " -1.561178\n", - " 0.655975\n", + " \n", + " \n", + " 15\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.290822\n", + " -0.332007\n", + " -1.154784\n", + " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", - " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 45\n", + " 16\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.476091\n", - " 1.811469\n", - " 0.655975\n", + " 0.321779\n", + " 0.073003\n", + " -0.580765\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 185\n", - " [{'index': 1, 'value': 1.0}]\n", - " 1.347773\n", - " -1.510077\n", - " 1.232087\n", + " 17\n", " [{'index': 2, 'value': 1.0}]\n", + " 1.230026\n", + " 1.136154\n", + " -0.580765\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 141\n", + " 18\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.79036\n", + " 0.376761\n", + " -0.580765\n", + " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", - " 1.036145\n", - " -1.101271\n", - " 1.088059\n", + " \n", + " \n", + " 20\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.4886\n", + " 0.326134\n", + " -0.580765\n", " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", + " \n", + " \n", + " 21\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.359775\n", + " 1.034902\n", + " -0.580765\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 147\n", - " [{'index': 1, 'value': 1.0}]\n", - " 1.109469\n", - " -0.641365\n", - " 1.66417\n", + " 24\n", " [{'index': 2, 'value': 1.0}]\n", + " 1.044669\n", + " 0.528639\n", + " -0.580765\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 320\n", + " 26\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.456718\n", + " 0.680518\n", + " -0.006745\n", " [{'index': 3, 'value': 1.0}]\n", - " -0.925279\n", - " 0.329548\n", - " -1.432429\n", - " [{'index': 1, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 285\n", - " [{'index': 1, 'value': 1.0}]\n", - " -0.192036\n", - " -2.072185\n", - " 1.016045\n", - " [{'index': 1, 'value': 1.0}]\n", + " 27\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.21149\n", + " 1.237407\n", + " -0.006745\n", " [{'index': 3, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 305\n", + " 29\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.378311\n", + " 0.933649\n", + " -0.006745\n", " [{'index': 3, 'value': 1.0}]\n", - " 0.027937\n", - " 0.43175\n", - " 0.655975\n", " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 37\n", - " [{'index': 2, 'value': 1.0}]\n", - " 1.421098\n", - " 1.607066\n", - " -0.496248\n", + " 31\n", " [{'index': 2, 'value': 1.0}]\n", + " -2.013713\n", + " -0.534512\n", + " -1.657052\n", " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 98\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.521995\n", - " 0.176246\n", - " -0.784304\n", + " 32\n", " [{'index': 2, 'value': 1.0}]\n", + " -1.253751\n", + " 0.478013\n", + " -1.657052\n", + " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 10\n", + " 33\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.576989\n", - " 0.891656\n", - " 0.511947\n", + " -0.827431\n", + " -0.230754\n", + " -1.657052\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", "\n", + "

25 rows × 6 columns

\n", "[267 rows x 6 columns in total]" ], "text/plain": [ " onehotencoded_island scaled_culmen_length_mm \\\n", "penguin_id \n", - "23 [{'index': 2, 'value': 1.0}] -1.310231 \n", - "83 [{'index': 2, 'value': 1.0}] 2.575955 \n", - "195 [{'index': 1, 'value': 1.0}] 0.816172 \n", - "264 [{'index': 1, 'value': 1.0}] 1.109469 \n", - "237 [{'index': 1, 'value': 1.0}] 0.816172 \n", - "261 [{'index': 1, 'value': 1.0}] 0.614531 \n", - "253 [{'index': 1, 'value': 1.0}] 0.321234 \n", - "164 [{'index': 1, 'value': 1.0}] 0.101261 \n", - "293 [{'index': 3, 'value': 1.0}] -0.686975 \n", - "176 [{'index': 1, 'value': 1.0}] 0.871166 \n", - "45 [{'index': 2, 'value': 1.0}] 1.476091 \n", - "185 [{'index': 1, 'value': 1.0}] 1.347773 \n", - "141 [{'index': 1, 'value': 1.0}] 1.036145 \n", - "147 [{'index': 1, 'value': 1.0}] 1.109469 \n", - "320 [{'index': 3, 'value': 1.0}] -0.925279 \n", - "285 [{'index': 1, 'value': 1.0}] -0.192036 \n", - "305 [{'index': 3, 'value': 1.0}] 0.027937 \n", - "37 [{'index': 2, 'value': 1.0}] 1.421098 \n", - "98 [{'index': 2, 'value': 1.0}] -0.521995 \n", - "10 [{'index': 2, 'value': 1.0}] -0.576989 \n", - "307 [{'index': 3, 'value': 1.0}] -1.420218 \n", - "337 [{'index': 3, 'value': 1.0}] -0.393678 \n", - "26 [{'index': 2, 'value': 1.0}] -0.448671 \n", - "97 [{'index': 2, 'value': 1.0}] -1.45688 \n", - "27 [{'index': 2, 'value': 1.0}] 1.201125 \n", + "0 [{'index': 2, 'value': 1.0}] -1.364965 \n", + "1 [{'index': 2, 'value': 1.0}] -0.771824 \n", + "2 [{'index': 2, 'value': 1.0}] -0.567932 \n", + "3 [{'index': 2, 'value': 1.0}] 0.470064 \n", + "4 [{'index': 2, 'value': 1.0}] -1.235216 \n", + "5 [{'index': 2, 'value': 1.0}] -0.141612 \n", + "6 [{'index': 2, 'value': 1.0}] 0.544207 \n", + "7 [{'index': 2, 'value': 1.0}] 1.21149 \n", + "8 [{'index': 2, 'value': 1.0}] 1.026133 \n", + "10 [{'index': 2, 'value': 1.0}] -0.586468 \n", + "11 [{'index': 2, 'value': 1.0}] -0.92011 \n", + "14 [{'index': 2, 'value': 1.0}] -1.846892 \n", + "15 [{'index': 2, 'value': 1.0}] -1.290822 \n", + "16 [{'index': 2, 'value': 1.0}] 0.321779 \n", + "17 [{'index': 2, 'value': 1.0}] 1.230026 \n", + "18 [{'index': 2, 'value': 1.0}] -0.79036 \n", + "20 [{'index': 2, 'value': 1.0}] 0.4886 \n", + "21 [{'index': 2, 'value': 1.0}] 1.359775 \n", + "24 [{'index': 2, 'value': 1.0}] 1.044669 \n", + "26 [{'index': 2, 'value': 1.0}] -0.456718 \n", + "27 [{'index': 2, 'value': 1.0}] 1.21149 \n", + "29 [{'index': 2, 'value': 1.0}] 1.378311 \n", + "31 [{'index': 2, 'value': 1.0}] -2.013713 \n", + "32 [{'index': 2, 'value': 1.0}] -1.253751 \n", + "33 [{'index': 2, 'value': 1.0}] -0.827431 \n", "\n", " scaled_culmen_depth_mm scaled_flipper_length_mm \\\n", "penguin_id \n", - "23 0.687253 -0.568262 \n", - "83 0.329548 -1.432429 \n", - "195 -1.407876 0.151877 \n", - "264 -0.436962 2.096254 \n", - "237 -1.305674 0.872017 \n", - "261 -0.947969 1.520142 \n", - "253 -1.663379 0.944031 \n", - "164 -0.743566 1.160073 \n", - "293 -0.079257 -1.792499 \n", - "176 -1.561178 0.655975 \n", - "45 1.811469 0.655975 \n", - "185 -1.510077 1.232087 \n", - "141 -1.101271 1.088059 \n", - "147 -0.641365 1.66417 \n", - "320 0.329548 -1.432429 \n", - "285 -2.072185 1.016045 \n", - "305 0.43175 0.655975 \n", - "37 1.607066 -0.496248 \n", - "98 0.176246 -0.784304 \n", - "10 0.891656 0.511947 \n", - "307 0.022944 -1.000346 \n", - "337 1.14716 -0.208192 \n", - "26 0.687253 0.00785 \n", - "97 0.380649 -0.784304 \n", - "27 1.249361 0.00785 \n", + "0 0.629892 -1.226537 \n", + "1 0.984275 -1.226537 \n", + "2 0.883023 -1.226537 \n", + "3 0.376761 -0.652517 \n", + "4 -0.180128 -0.652517 \n", + "5 0.680518 -0.652517 \n", + "6 -0.281381 -0.652517 \n", + "7 0.629892 -0.078497 \n", + "8 0.933649 -0.078497 \n", + "10 0.883023 0.495523 \n", + "11 0.781771 -1.154784 \n", + "14 -0.028249 -1.154784 \n", + "15 -0.332007 -1.154784 \n", + "16 0.073003 -0.580765 \n", + "17 1.136154 -0.580765 \n", + "18 0.376761 -0.580765 \n", + "20 0.326134 -0.580765 \n", + "21 1.034902 -0.580765 \n", + "24 0.528639 -0.580765 \n", + "26 0.680518 -0.006745 \n", + "27 1.237407 -0.006745 \n", + "29 0.933649 -0.006745 \n", + "31 -0.534512 -1.657052 \n", + "32 0.478013 -1.657052 \n", + "33 -0.230754 -1.657052 \n", "\n", " onehotencoded_sex onehotencoded_species \n", "penguin_id \n", - "23 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "83 [{'index': 1, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", - "195 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "264 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "237 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "261 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "253 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "164 [{'index': 0, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "293 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "176 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "45 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", - "185 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "141 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "147 [{'index': 2, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "320 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "285 [{'index': 1, 'value': 1.0}] [{'index': 3, 'value': 1.0}] \n", - "305 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "37 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", - "98 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "10 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "307 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "337 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "26 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "97 [{'index': 1, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", - "27 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "0 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "1 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "2 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "3 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "4 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "5 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "6 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "7 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "8 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "10 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "11 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "14 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "15 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "16 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "17 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "18 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "20 [{'index': 2, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "21 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "24 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "26 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "27 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "29 [{'index': 3, 'value': 1.0}] [{'index': 2, 'value': 1.0}] \n", + "31 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "32 [{'index': 3, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", + "33 [{'index': 2, 'value': 1.0}] [{'index': 1, 'value': 1.0}] \n", "...\n", "\n", "[267 rows x 6 columns]" @@ -1490,6 +2099,76 @@ "execution_count": 8, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "b7c22c0858914b46951adde174b43e25", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job b8c690b1-fa1b-4c4c-b802-c8ab37c930c9 is DONE. 359 Bytes processed. \n", " \n", " \n", - " 255\n", - " 4789.541771\n", + " 9\n", + " 4295.335461\n", " \n", " \n", - " 158\n", - " 4642.382347\n", + " 12\n", + " 3338.44131\n", " \n", " \n", - " 230\n", - " 3562.484258\n", + " 13\n", + " 3201.820204\n", " \n", " \n", - " 218\n", - " 5321.625743\n", + " 19\n", + " 3982.814079\n", " \n", " \n", - " 297\n", - " 3926.179647\n", + " 22\n", + " 3538.610664\n", " \n", " \n", - " 93\n", - " 4158.344927\n", + " 23\n", + " 3613.50305\n", " \n", " \n", - " 338\n", - " 4093.191646\n", + " 25\n", + " 4009.759444\n", " \n", " \n", - " 267\n", - " 5590.390771\n", + " 28\n", + " 4240.515635\n", " \n", " \n", - " 275\n", - " 3623.482531\n", + " 30\n", + " 4028.904195\n", " \n", " \n", - " 200\n", - " 5494.402339\n", + " 38\n", + " 4206.810346\n", " \n", " \n", - " 107\n", - " 3450.730176\n", + " 41\n", + " 3736.225488\n", " \n", " \n", - " 227\n", - " 5707.948707\n", + " 45\n", + " 4383.368544\n", " \n", " \n", - " 199\n", - " 4855.771085\n", + " 49\n", + " 3125.772789\n", " \n", " \n", - " 260\n", - " 5477.219436\n", + " 50\n", + " 3149.28765\n", " \n", " \n", - " 201\n", - " 4965.59121\n", + " 62\n", + " 3531.69488\n", " \n", " \n", - " 7\n", - " 3670.281516\n", + " 65\n", + " 4073.900616\n", " \n", " \n", - " 8\n", - " 4072.678667\n", + " 66\n", + " 4160.810162\n", " \n", " \n", - " 126\n", - " 3339.624869\n", + " 83\n", + " 3272.396279\n", " \n", " \n", - " 39\n", - " 3556.890093\n", + " 87\n", + " 3989.364493\n", " \n", " \n", - " 12\n", - " 3302.203276\n", + " 92\n", + " 4240.495294\n", + " \n", + " \n", + " 98\n", + " 3911.455384\n", + " \n", + " \n", + " 104\n", + " 3271.202866\n", + " \n", + " \n", + " 114\n", + " 3244.728549\n", + " \n", + " \n", + " 115\n", + " 3737.374636\n", + " \n", + " \n", + " 118\n", + " 3485.95604\n", " \n", " \n", "\n", + "

25 rows × 1 columns

\n", "[67 rows x 1 columns in total]" ], "text/plain": [ " predicted_body_mass_g\n", "penguin_id \n", - "255 4789.541771\n", - "158 4642.382347\n", - "230 3562.484258\n", - "218 5321.625743\n", - "297 3926.179647\n", - "93 4158.344927\n", - "338 4093.191646\n", - "267 5590.390771\n", - "275 3623.482531\n", - "200 5494.402339\n", - "107 3450.730176\n", - "227 5707.948707\n", - "199 4855.771085\n", - "260 5477.219436\n", - "201 4965.59121\n", - "7 3670.281516\n", - "8 4072.678667\n", - "126 3339.624869\n", - "39 3556.890093\n", - "12 3302.203276\n", - "34 3227.822355\n", - "206 3864.010614\n", - "53 3329.248666\n", - "282 5297.303727\n", - "225 5770.689468\n", + "9 4295.335461\n", + "12 3338.44131\n", + "13 3201.820204\n", + "19 3982.814079\n", + "22 3538.610664\n", + "23 3613.50305\n", + "25 4009.759444\n", + "28 4240.515635\n", + "30 4028.904195\n", + "38 4206.810346\n", + "41 3736.225488\n", + "45 4383.368544\n", + "49 3125.772789\n", + "50 3149.28765\n", + "62 3531.69488\n", + "65 4073.900616\n", + "66 4160.810162\n", + "83 3272.396279\n", + "87 3989.364493\n", + "92 4240.495294\n", + "98 3911.455384\n", + "104 3271.202866\n", + "114 3244.728549\n", + "115 3737.374636\n", + "118 3485.95604\n", "...\n", "\n", "[67 rows x 1 columns]" @@ -1670,6 +2370,76 @@ "execution_count": 9, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "fa6bd24b44cb42ec946e262ac2f25d09", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 8e88f903-5b13-4305-b259-d17194be16da is DONE. 809 Bytes processed. \n", " \n", " \n", - " 17\n", - " 1\n", + " 9\n", + " 4\n", " \n", " \n", - " 230\n", + " 12\n", + " 4\n", + " \n", + " \n", + " 13\n", " 2\n", " \n", " \n", - " 291\n", + " 19\n", + " 4\n", + " \n", + " \n", + " 22\n", " 4\n", " \n", " \n", - " 177\n", - " 3\n", + " 23\n", + " 4\n", " \n", " \n", - " 206\n", + " 25\n", " 2\n", " \n", " \n", - " 331\n", + " 28\n", " 2\n", " \n", " \n", - " 18\n", + " 30\n", " 2\n", " \n", " \n", - " 218\n", + " 38\n", " 4\n", " \n", " \n", - " 116\n", - " 2\n", + " 41\n", + " 4\n", " \n", " \n", - " 271\n", + " 45\n", " 2\n", " \n", " \n", - " 165\n", - " 3\n", + " 49\n", + " 4\n", " \n", " \n", - " 194\n", + " 50\n", " 2\n", " \n", " \n", - " 39\n", - " 1\n", + " 62\n", + " 4\n", " \n", " \n", - " 255\n", - " 3\n", + " 65\n", + " 4\n", " \n", " \n", - " 225\n", - " 4\n", + " 66\n", + " 2\n", " \n", " \n", - " 40\n", - " 1\n", + " 83\n", + " 2\n", " \n", " \n", - " 8\n", - " 1\n", + " 87\n", + " 2\n", " \n", " \n", - " 155\n", - " 4\n", + " 92\n", + " 2\n", " \n", " \n", - " 239\n", + " 98\n", " 4\n", " \n", " \n", - " 138\n", - " 3\n", + " 104\n", + " 2\n", + " \n", + " \n", + " 114\n", + " 2\n", + " \n", + " \n", + " 115\n", + " 2\n", + " \n", + " \n", + " 118\n", + " 4\n", " \n", " \n", "\n", + "

25 rows × 1 columns

\n", "[67 rows x 1 columns in total]" ], "text/plain": [ " CENTROID_ID\n", "penguin_id \n", - "17 1\n", - "230 2\n", - "291 4\n", - "177 3\n", - "206 2\n", - "331 2\n", - "18 2\n", - "218 4\n", - "116 2\n", - "271 2\n", - "165 3\n", - "194 2\n", - "39 1\n", - "255 3\n", - "225 4\n", - "40 1\n", - "8 1\n", - "155 4\n", - "239 4\n", - "138 3\n", - "89 1\n", - "34 1\n", - "106 2\n", - "12 2\n", - "158 3\n", + "9 4\n", + "12 4\n", + "13 2\n", + "19 4\n", + "22 4\n", + "23 4\n", + "25 2\n", + "28 2\n", + "30 2\n", + "38 4\n", + "41 4\n", + "45 2\n", + "49 4\n", + "50 2\n", + "62 4\n", + "65 4\n", + "66 2\n", + "83 2\n", + "87 2\n", + "92 2\n", + "98 4\n", + "104 2\n", + "114 2\n", + "115 2\n", + "118 4\n", "...\n", "\n", "[67 rows x 1 columns]" @@ -1890,6 +2681,76 @@ "execution_count": 11, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "f0465986682742af92759c3f5fce96e0", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job c7d094cb-cc51-4f11-8887-b169c23aceb2 is DONE. 32.3 kB processed. \n", " \n", " \n", - " 125\n", - " 3911.489556\n", + " 9\n", + " 4295.328991\n", " \n", " \n", - " 260\n", - " 5477.221344\n", + " 12\n", + " 3338.434943\n", " \n", " \n", - " 93\n", - " 4158.347237\n", + " 13\n", + " 3201.813783\n", " \n", " \n", - " 255\n", - " 4789.543554\n", + " 19\n", + " 3982.807707\n", " \n", " \n", - " 71\n", - " 3153.676672\n", + " 22\n", + " 3538.604385\n", " \n", " \n", - " 327\n", - " 3358.229939\n", + " 23\n", + " 3613.496641\n", " \n", " \n", - " 338\n", - " 4093.194446\n", + " 25\n", + " 4009.753161\n", " \n", " \n", - " 68\n", - " 4133.390538\n", + " 28\n", + " 4240.509087\n", " \n", " \n", - " 17\n", - " 4013.257221\n", + " 30\n", + " 4028.897875\n", " \n", " \n", - " 116\n", - " 4189.275197\n", + " 38\n", + " 4206.80377\n", " \n", " \n", - " 106\n", - " 4019.304549\n", + " 41\n", + " 3736.219256\n", " \n", " \n", - " 227\n", - " 5707.950809\n", + " 45\n", + " 4383.362136\n", " \n", " \n", - " 314\n", - " 4037.311435\n", + " 49\n", + " 3125.766474\n", " \n", " \n", - " 246\n", - " 3297.59755\n", + " 50\n", + " 3149.281322\n", " \n", " \n", - " 138\n", - " 4971.171877\n", + " 62\n", + " 3531.688645\n", " \n", " \n", - " 126\n", - " 3339.627076\n", + " 65\n", + " 4073.894238\n", " \n", " \n", - " 121\n", - " 3462.463115\n", + " 66\n", + " 4160.803738\n", + " \n", + " \n", + " 83\n", + " 3272.389735\n", + " \n", + " \n", + " 87\n", + " 3989.358086\n", + " \n", + " \n", + " 92\n", + " 4240.488891\n", + " \n", + " \n", + " 98\n", + " 3911.449023\n", " \n", " \n", - " 231\n", - " 4206.392245\n", + " 104\n", + " 3271.196535\n", " \n", " \n", - " 40\n", - " 4102.581912\n", + " 114\n", + " 3244.722283\n", " \n", " \n", - " 60\n", - " 3483.351408\n", + " 115\n", + " 3737.368277\n", + " \n", + " \n", + " 118\n", + " 3485.949702\n", " \n", " \n", "\n", + "

25 rows × 1 columns

\n", "[67 rows x 1 columns in total]" ], "text/plain": [ " predicted_body_mass_g\n", "penguin_id \n", - "125 3911.489556\n", - "260 5477.221344\n", - "93 4158.347237\n", - "255 4789.543554\n", - "71 3153.676672\n", - "327 3358.229939\n", - "338 4093.194446\n", - "68 4133.390538\n", - "17 4013.257221\n", - "116 4189.275197\n", - "106 4019.304549\n", - "227 5707.950809\n", - "314 4037.311435\n", - "246 3297.59755\n", - "138 4971.171877\n", - "126 3339.627076\n", - "121 3462.463115\n", - "231 4206.392245\n", - "40 4102.581912\n", - "60 3483.351408\n", - "331 3370.089387\n", - "218 5321.627618\n", - "154 3618.98119\n", - "130 4249.575399\n", - "89 4145.000048\n", + "9 4295.328991\n", + "12 3338.434943\n", + "13 3201.813783\n", + "19 3982.807707\n", + "22 3538.604385\n", + "23 3613.496641\n", + "25 4009.753161\n", + "28 4240.509087\n", + "30 4028.897875\n", + "38 4206.80377\n", + "41 3736.219256\n", + "45 4383.362136\n", + "49 3125.766474\n", + "50 3149.281322\n", + "62 3531.688645\n", + "65 4073.894238\n", + "66 4160.803738\n", + "83 3272.389735\n", + "87 3989.358086\n", + "92 4240.488891\n", + "98 3911.449023\n", + "104 3271.196535\n", + "114 3244.722283\n", + "115 3737.368277\n", + "118 3485.949702\n", "...\n", "\n", "[67 rows x 1 columns]" @@ -2071,6 +2953,76 @@ "execution_count": 12, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "e6409bce719940f4ae326a8b18871d9b", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job a427bad2-9875-453f-ad2a-1eefaf085657 is DONE. 32.3 kB processed. \n", " \n", " 0\n", - " 223.862457\n", - " 73619.165369\n", - " 0.004765\n", - " 209.607289\n", - " 0.888368\n", - " 0.901838\n", + " 241.640738\n", + " 90117.84266\n", + " 0.005652\n", + " 200.718678\n", + " 0.8727\n", + " 0.878359\n", " \n", " \n", "\n", + "

1 rows × 6 columns

\n", "[1 rows x 6 columns in total]" ], "text/plain": [ " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", - "0 223.862457 73619.165369 0.004765 \n", + "0 241.640738 90117.84266 0.005652 \n", "\n", " median_absolute_error r2_score explained_variance \n", - "0 209.607289 0.888368 0.901838 \n", + "0 200.718678 0.8727 0.878359 \n", "\n", "[1 rows x 6 columns]" ] @@ -2147,10 +3100,52 @@ "execution_count": 13, "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "e56e2cc197894ee9acc448a8c12e8a30", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 929c826c-1051-47fc-9256-546f4ef11c32 is DONE. 31.7 kB processed. str: + """Provides a pretrained model as a test fixture that is cached across test runs. + This lets us run system tests without having to wait for a model.fit(...)""" + sql = f""" +CREATE OR REPLACE MODEL `$model_name` +OPTIONS ( + model_type='kmeans', + num_clusters=3 +) AS SELECT + culmen_length_mm, + culmen_depth_mm, + flipper_length_mm, + sex +FROM `{penguins_table_id}`""" + # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited + model_name = f"{dataset_id_permanent}.penguins_logistic_reg_{hashlib.md5(sql.encode()).hexdigest()}" + sql = sql.replace("$model_name", model_name) + + try: + session.bqclient.get_model(model_name) + except google.cloud.exceptions.NotFound: + logging.info( + "penguins_logistic_model fixture was not found in the permanent dataset, regenerating it..." + ) + session.bqclient.query(sql).result() + finally: + return model_name + + @pytest.fixture(scope="session") def penguins_xgbregressor_model_name( session: bigframes.Session, dataset_id_permanent, penguins_table_id diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 8325cec50aa..87ea46f969c 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -48,13 +48,19 @@ def penguins_bqml_linear_model(session, penguins_linear_model_name) -> core.Bqml @pytest.fixture(scope="function") def ephemera_penguins_bqml_linear_model( penguins_bqml_linear_model, -) -> linear_model.LinearRegression: +) -> core.BqmlModel: model = penguins_bqml_linear_model return model.copy( f"{model._model.project}.{model._model.dataset_id}.{uuid.uuid4().hex}" ) +@pytest.fixture(scope="session") +def penguins_bqml_kmeans_model(session, penguins_kmeans_model_name) -> core.BqmlModel: + model = session.bqclient.get_model(penguins_kmeans_model_name) + return core.BqmlModel(session, model) + + @pytest.fixture(scope="session") def penguins_linear_model( session, penguins_linear_model_name: str @@ -125,34 +131,11 @@ def penguins_randomforest_classifier_model( @pytest.fixture(scope="session") -def penguins_kmeans_model( - session: bigframes.Session, dataset_id_permanent, penguins_table_id -) -> cluster.KMeans: - """Provides a pretrained model as a test fixture that is cached across test runs. - This lets us run system tests without having to wait for a model.fit(...)""" - sql = f""" -CREATE OR REPLACE MODEL `$model_name` -OPTIONS ( - model_type='kmeans', - num_clusters=3 -) AS SELECT - culmen_length_mm, - culmen_depth_mm, - flipper_length_mm, - sex -FROM `{penguins_table_id}`""" - # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited - model_name = f"{dataset_id_permanent}.penguins_cluster_{hashlib.md5(sql.encode()).hexdigest()}" - sql = sql.replace("$model_name", model_name) - - try: - return session.read_gbq_model(model_name) - except google.cloud.exceptions.NotFound: - logging.info( - "penguins_kmeans_model fixture was not found in the permanent dataset, regenerating it..." - ) - session.bqclient.query(sql).result() - return session.read_gbq_model(model_name) +def penguins_kmeans_model(session, penguins_kmeans_model_name: str) -> cluster.KMeans: + return cast( + cluster.KMeans, + session.read_gbq_model(penguins_kmeans_model_name), + ) @pytest.fixture(scope="session") diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py index a003cd1ec1b..d95a1e1bc22 100644 --- a/tests/system/small/ml/test_cluster.py +++ b/tests/system/small/ml/test_cluster.py @@ -88,5 +88,71 @@ def test_kmeans_score(session, penguins_kmeans_model: cluster.KMeans): ) +def test_kmeans_cluster_centers(penguins_kmeans_model: cluster.KMeans): + result = penguins_kmeans_model.cluster_centers_.to_pandas() + expected = pd.DataFrame( + { + "centroid_id": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3], + "feature": [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + * 3, + "numerical_value": [ + 47.509677, + 14.993548, + 217.040123, + pd.NA, + 38.207813, + 18.03125, + 187.992188, + pd.NA, + 47.036346, + 18.834808, + 197.1612, + pd.NA, + ], + "categorical_value": [ + [], + [], + [], + [ + {"category": ".", "value": 0.008064516129032258}, + {"category": "MALE", "value": 0.49193548387096775}, + {"category": "FEMALE", "value": 0.47580645161290325}, + {"category": "_null_filler", "value": 0.024193548387096774}, + ], + [], + [], + [], + [ + {"category": "MALE", "value": 0.34375}, + {"category": "FEMALE", "value": 0.625}, + {"category": "_null_filler", "value": 0.03125}, + ], + [], + [], + [], + [ + {"category": "MALE", "value": 0.6847826086956522}, + {"category": "FEMALE", "value": 0.2826086956521739}, + {"category": "_null_filler", "value": 0.03260869565217391}, + ], + ], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + check_dtype=False, + ) + + def test_loaded_config(penguins_kmeans_model): assert penguins_kmeans_model.n_clusters == 3 diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index 8b864d9b55c..4b184b0d4c8 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -21,7 +21,7 @@ import pytz import bigframes -import bigframes.ml.core +from bigframes.ml import core def test_model_eval( @@ -74,9 +74,73 @@ def test_model_eval_with_data(penguins_bqml_linear_model, penguins_df_default_in ) -def test_model_predict( - penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, new_penguins_df -): +def test_model_centroids(penguins_bqml_kmeans_model: core.BqmlModel): + result = penguins_bqml_kmeans_model.centroids().to_pandas() + expected = pd.DataFrame( + { + "centroid_id": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3], + "feature": [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + * 3, + "numerical_value": [ + 47.509677, + 14.993548, + 217.040123, + pd.NA, + 38.207813, + 18.03125, + 187.992188, + pd.NA, + 47.036346, + 18.834808, + 197.1612, + pd.NA, + ], + "categorical_value": [ + [], + [], + [], + [ + {"category": ".", "value": 0.008064516129032258}, + {"category": "MALE", "value": 0.49193548387096775}, + {"category": "FEMALE", "value": 0.47580645161290325}, + {"category": "_null_filler", "value": 0.024193548387096774}, + ], + [], + [], + [], + [ + {"category": "MALE", "value": 0.34375}, + {"category": "FEMALE", "value": 0.625}, + {"category": "_null_filler", "value": 0.03125}, + ], + [], + [], + [], + [ + {"category": "MALE", "value": 0.6847826086956522}, + {"category": "FEMALE", "value": 0.2826086956521739}, + {"category": "_null_filler", "value": 0.03260869565217391}, + ], + ], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + check_dtype=False, + ) + + +def test_model_predict(penguins_bqml_linear_model: core.BqmlModel, new_penguins_df): predictions = penguins_bqml_linear_model.predict(new_penguins_df).to_pandas() expected = pd.DataFrame( {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]}, @@ -92,7 +156,7 @@ def test_model_predict( def test_model_predict_with_unnamed_index( - penguins_bqml_linear_model: bigframes.ml.core.BqmlModel, new_penguins_df + penguins_bqml_linear_model: core.BqmlModel, new_penguins_df ): # This will result in an index that lacks a name, which the ML library will @@ -121,7 +185,7 @@ def test_model_predict_with_unnamed_index( def test_model_generate_text( - bqml_palm2_text_generator_model: bigframes.ml.core.BqmlModel, llm_text_df + bqml_palm2_text_generator_model: core.BqmlModel, llm_text_df ): options = { "temperature": 0.5, @@ -148,7 +212,7 @@ def test_model_generate_text( assert all(series.str.len() > 20) -def test_model_forecast(time_series_bqml_arima_plus_model: bigframes.ml.core.BqmlModel): +def test_model_forecast(time_series_bqml_arima_plus_model: core.BqmlModel): utc = pytz.utc forecast = time_series_bqml_arima_plus_model.forecast().to_pandas()[ ["forecast_timestamp", "forecast_value"] diff --git a/tests/system/small/operations/test_strings.py b/tests/system/small/operations/test_strings.py index d7bf3312f09..31b64f4314d 100644 --- a/tests/system/small/operations/test_strings.py +++ b/tests/system/small/operations/test_strings.py @@ -357,3 +357,86 @@ def test_cat_with_series(scalars_dfs): pd_result, bf_result, ) + + +def test_str_match(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + pattern = "[A-Z].*" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.match(pattern).to_pandas() + pd_result = scalars_pandas_df[col_name].str.match(pattern) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_str_fullmatch(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + pattern = "[A-Z].*!" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.fullmatch(pattern).to_pandas() + pd_result = scalars_pandas_df[col_name].str.fullmatch(pattern) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_str_get(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.get(8).to_pandas() + pd_result = scalars_pandas_df[col_name].str.get(8) + + print(pd_result) + print(bf_result) + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_str_pad(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.pad(8, side="both", fillchar="%").to_pandas() + pd_result = scalars_pandas_df[col_name].str.pad(8, side="both", fillchar="%") + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_str_ljust(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.ljust(7, fillchar="%").to_pandas() + pd_result = scalars_pandas_df[col_name].str.ljust(7, fillchar="%") + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) + + +def test_str_rjust(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_series: bigframes.series.Series = scalars_df[col_name] + bf_result = bf_series.str.rjust(9, fillchar="%").to_pandas() + pd_result = scalars_pandas_df[col_name].str.rjust(9, fillchar="%") + + assert_series_equal_ignoring_order( + pd_result, + bf_result, + ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index c7e17f5a2dd..2c44dd80670 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -86,6 +86,14 @@ def test_get_column(scalars_dfs): assert_series_equal_ignoring_order(bf_result, pd_result) +def test_get_column_nonstring(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + series = scalars_df.rename(columns={"int64_col": 123.1})[123.1] + bf_result = series.to_pandas() + pd_result = scalars_pandas_df.rename(columns={"int64_col": 123.1})[123.1] + assert_series_equal_ignoring_order(bf_result, pd_result) + + def test_hasattr(scalars_dfs): scalars_df, _ = scalars_dfs assert hasattr(scalars_df, "int64_col") @@ -215,7 +223,7 @@ def test_drop_index_and_columns(scalars_dfs): def test_rename(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs - col_name_dict = {"bool_col": "boolean_col"} + col_name_dict = {"bool_col": 1.2345} df_pandas = scalars_df.rename(columns=col_name_dict).to_pandas() pd.testing.assert_index_equal( df_pandas.columns, scalars_pandas_df.rename(columns=col_name_dict).columns @@ -1751,6 +1759,16 @@ def test_rename_axis(scalars_df_index, scalars_pandas_df_index): ) +def test_rename_axis_nonstring(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.rename_axis((4,)) + pd_result = scalars_pandas_df_index.rename_axis((4,)) + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): pd_string_series = scalars_pandas_df_index.string_col.iloc[[0, 5, 1, 1, 5]] bf_string_series = scalars_df_index.string_col.iloc[[0, 5, 1, 1, 5]] diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index 7ad753e1bc0..1f5aa906c8c 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -28,6 +28,8 @@ except ImportError: pandas_gbq = None +import typing + import bigframes import bigframes.dataframe @@ -80,7 +82,7 @@ def test_to_csv_index( if scalars_df.index.name is not None: path = gcs_folder + f"test_index_df_to_csv_index_{index}*.csv" if index: - index_col = scalars_df.index.name + index_col = typing.cast(str, scalars_df.index.name) else: path = gcs_folder + f"test_default_index_df_to_csv_index_{index}*.csv" @@ -118,7 +120,7 @@ def test_to_csv_tabs( pytest.skip("date_format parameter not supported in pandas 1.x.") """Test the `to_csv` API with the `sep` parameter.""" scalars_df, scalars_pandas_df = scalars_dfs - index_col = scalars_df.index.name + index_col = typing.cast(str, scalars_df.index.name) path = gcs_folder + "test_to_csv_tabs*.csv" # TODO(swast): Support "date_format" parameter and make sure our diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index e72d75729be..d5dd4e357bc 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -15,6 +15,8 @@ import pandas as pd import pytest +import bigframes.pandas as bpd + @pytest.mark.parametrize( ("operator"), @@ -86,6 +88,73 @@ def test_dataframe_groupby_aggregate( pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) +def test_dataframe_groupby_agg_string(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = scalars_df_index[col_names].groupby("string_col").agg("count") + pd_result = scalars_pandas_df_index[col_names].groupby("string_col").agg("count") + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_frame_equal( + pd_result, + bf_result_computed, + check_dtype=False, + ) + + +def test_dataframe_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = scalars_df_index[col_names].groupby("string_col").agg(["count", "min"]) + pd_result = ( + scalars_pandas_df_index[col_names].groupby("string_col").agg(["count", "min"]) + ) + bf_result_computed = bf_result.to_pandas() + + # Pandas produces multi-index which isn't supported in bq df yet + pd_result = pd_result.set_axis(bf_result.columns, axis=1) + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + +def test_dataframe_groupby_agg_dict(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names] + .groupby("string_col") + .agg({"int64_too": ["mean", "max"], "string_col": "count"}) + ) + pd_result = ( + scalars_pandas_df_index[col_names] + .groupby("string_col") + .agg({"int64_too": ["mean", "max"], "string_col": "count"}) + ) + bf_result_computed = bf_result.to_pandas() + + # Pandas produces multi-index which isn't supported in bq df yet + pd_result = pd_result.set_axis(bf_result.columns, axis=1) + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + +def test_dataframe_groupby_agg_named(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names] + .groupby("string_col") + .agg( + agg1=bpd.NamedAgg("int64_too", "sum"), + agg2=bpd.NamedAgg("float64_col", "max"), + ) + ) + pd_result = ( + scalars_pandas_df_index[col_names] + .groupby("string_col") + .agg( + agg1=pd.NamedAgg("int64_too", "sum"), agg2=pd.NamedAgg("float64_col", "max") + ) + ) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + @pytest.mark.parametrize( ("as_index"), [ @@ -173,3 +242,39 @@ def test_dataframe_groupby_getitem_list( ) pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) + + +def test_series_groupby_agg_string(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_col"] + .groupby(scalars_df_index["string_col"]) + .agg("count") + ) + pd_result = ( + scalars_pandas_df_index["int64_col"] + .groupby(scalars_pandas_df_index["string_col"]) + .agg("count") + ) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_series_equal( + pd_result, bf_result_computed, check_dtype=False, check_names=False + ) + + +def test_series_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_col"] + .groupby(scalars_df_index["string_col"]) + .agg(["sum", "mean"]) + ) + pd_result = ( + scalars_pandas_df_index["int64_col"] + .groupby(scalars_pandas_df_index["string_col"]) + .agg(["sum", "mean"]) + ) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_frame_equal( + pd_result, bf_result_computed, check_dtype=False, check_names=False + ) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index fe4b1c5a97a..47d758763ba 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -37,11 +37,11 @@ def bq_cf_connection_location() -> str: @pytest.fixture(scope="module") -def bq_cf_connection_location_mistached() -> str: +def bq_cf_connection_location_mismatched() -> str: """Pre-created BQ connection to invoke cloud function for bigframes-dev - $ bq show --connection --location=us-east1 --project_id=bigframes-dev bigframes-rf-conn + $ bq show --connection --location=eu --project_id=bigframes-dev bigframes-rf-conn """ - return "us-east1.bigframes-rf-conn" + return "eu.bigframes-rf-conn" @pytest.fixture(scope="module") @@ -53,11 +53,11 @@ def bq_cf_connection_location_project() -> str: @pytest.fixture(scope="module") -def bq_cf_connection_location_project_mistached() -> str: +def bq_cf_connection_location_project_mismatched() -> str: """Pre-created BQ connection to invoke cloud function for bigframes-dev - $ bq show --connection --location=us-east1 --project_id=bigframes-metrics bigframes-rf-conn + $ bq show --connection --location=eu --project_id=bigframes-metrics bigframes-rf-conn """ - return "bigframes-metrics.us-east1.bigframes-rf-conn" + return "bigframes-metrics.eu.bigframes-rf-conn" @pytest.fixture(scope="module") @@ -196,50 +196,28 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_direct_no_session_param_location_mistached( +def test_remote_function_direct_no_session_param_location_mismatched( bigquery_client, bigqueryconnection_client, cloudfunctions_client, - scalars_dfs, dataset_id_permanent, - bq_cf_connection_location_mistached, + bq_cf_connection_location_mismatched, ): - @remote_function( - [int], - int, - bigquery_client=bigquery_client, - bigquery_connection_client=bigqueryconnection_client, - cloud_functions_client=cloudfunctions_client, - dataset=dataset_id_permanent, - bigquery_connection=bq_cf_connection_location_mistached, - # See e2e tests for tests that actually deploy the Cloud Function. - reuse=True, - ) - def square(x): - return x * x - - scalars_df, scalars_pandas_df = scalars_dfs - - bf_int64_col = scalars_df["int64_col"] - bf_int64_col_filter = bf_int64_col.notnull() - bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] - bf_result_col = bf_int64_col_filtered.apply(square) - bf_result = ( - bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() - ) - - pd_int64_col = scalars_pandas_df["int64_col"] - pd_int64_col_filter = pd_int64_col.notnull() - pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] - pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) - # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. - # pd_int64_col_filtered.dtype is Int64Dtype() - # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. - # For this test let's force the pandas dtype to be same as bigframes' dtype. - pd_result_col = pd_result_col.astype(pd.Int64Dtype()) - pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + with pytest.raises(ValueError): - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + @remote_function( + [int], + int, + bigquery_client=bigquery_client, + bigquery_connection_client=bigqueryconnection_client, + cloud_functions_client=cloudfunctions_client, + dataset=dataset_id_permanent, + bigquery_connection=bq_cf_connection_location_mismatched, + # See e2e tests for tests that actually deploy the Cloud Function. + reuse=True, + ) + def square(x): + return x * x @pytest.mark.flaky(retries=2, delay=120) @@ -295,7 +273,7 @@ def test_remote_function_direct_no_session_param_project_mismatched( bigqueryconnection_client, cloudfunctions_client, dataset_id_permanent, - bq_cf_connection_location_project_mistached, + bq_cf_connection_location_project_mismatched, ): with pytest.raises(ValueError): @@ -306,7 +284,7 @@ def test_remote_function_direct_no_session_param_project_mismatched( bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, dataset=dataset_id_permanent, - bigquery_connection=bq_cf_connection_location_project_mistached, + bigquery_connection=bq_cf_connection_location_project_mismatched, # See e2e tests for tests that actually deploy the Cloud Function. reuse=True, ) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 1c1e7b035bd..70c56e5e137 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -675,6 +675,51 @@ def test_copy(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_series_equal(bf_copy.to_pandas(), pd_copy) +def test_isin_raise_error(scalars_df_index, scalars_pandas_df_index): + col_name = "int64_too" + with pytest.raises(TypeError): + scalars_df_index[col_name].isin("whatever").to_pandas() + + +@pytest.mark.parametrize( + ( + "col_name", + "test_set", + ), + [ + ( + "int64_col", + [314159, 2.0, 3, pd.NA], + ), + ( + "int64_col", + [2, 55555, 4], + ), + ( + "float64_col", + [-123.456, 1.25, pd.NA], + ), + ( + "int64_too", + [1, 2, pd.NA], + ), + ( + "string_col", + ["Hello, World!", "Hi", "こんにちは"], + ), + ], +) +def test_isin(scalars_dfs, col_name, test_set): + scalars_df, scalars_pandas_df = scalars_dfs + print(type(scalars_pandas_df["datetime_col"].iloc[0])) + bf_result = scalars_df[col_name].isin(test_set).to_pandas() + pd_result = scalars_pandas_df[col_name].isin(test_set).astype("boolean") + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + def test_isnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "float64_col" @@ -2197,6 +2242,16 @@ def test_rename(scalars_df_index, scalars_pandas_df_index): ) +def test_rename_nonstring(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.rename((4, 2)) + pd_result = scalars_pandas_df_index.string_col.rename((4, 2)) + + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + def test_rename_dict_same_type(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.string_col.rename({1: 100, 2: 200}) pd_result = scalars_pandas_df_index.string_col.rename({1: 100, 2: 200}) @@ -2356,3 +2411,41 @@ def test_query_job_setters(scalars_dfs): series.to_pandas() job_ids.add(series.query_job.job_id) assert len(job_ids) == 2 + + +@pytest.mark.parametrize( + ("series_input",), + [ + ([1, 2, 3, 4, 5],), + ([1, 1, 3, 5, 5],), + ([1, pd.NA, 4, 5, 5],), + ([1, 3, 2, 5, 4],), + ([pd.NA, pd.NA],), + ([1, 1, 1, 1, 1],), + ], +) +def test_is_monotonic_increasing(series_input): + scalars_df = series.Series(series_input) + scalars_pandas_df = pd.Series(series_input) + assert ( + scalars_df.is_monotonic_increasing == scalars_pandas_df.is_monotonic_increasing + ) + + +@pytest.mark.parametrize( + ("series_input",), + [ + ([1],), + ([5, 4, 3, 2, 1],), + ([5, 5, 3, 1, 1],), + ([1, pd.NA, 4, 5, 5],), + ([5, pd.NA, 4, 2, 1],), + ([1, 1, 1, 1, 1],), + ], +) +def test_is_monotonic_decreasing(series_input): + scalars_df = series.Series(series_input) + scalars_pandas_df = pd.Series(series_input) + assert ( + scalars_df.is_monotonic_decreasing == scalars_pandas_df.is_monotonic_decreasing + ) diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 599b8aabbc9..2fc34f9bae9 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -11,8 +11,10 @@ # 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 random import tempfile +import textwrap import typing from typing import List @@ -160,6 +162,25 @@ def test_read_gbq_w_index_col( assert bf_shape == result.shape +def test_read_gbq_w_anonymous_query_results_table(session: bigframes.Session): + """Ensure BigQuery DataFrames can be used to inspect the results of a query job.""" + query = textwrap.dedent( + """ + SELECT SUM(`number`) AS total_people, name + FROM `bigquery-public-data.usa_names.usa_1910_2013` + GROUP BY name + HAVING name < "B" + """ + ) + job = session.bqclient.query(query) + expected = job.to_dataframe().set_index("name").sort_index() + destination = f"{job.destination.project}.{job.destination.dataset_id}.{job.destination.table_id}" + df = session.read_gbq(destination, index_col="name") + result = df.to_pandas() + expected.index = expected.index.astype(result.index.dtype) + pd.testing.assert_frame_equal(result, expected, check_dtype=False) + + @pytest.mark.parametrize( ("query_or_table", "max_results"), [ diff --git a/tests/unit/core/test_io.py b/tests/unit/core/test_io.py new file mode 100644 index 00000000000..c5074f80c2c --- /dev/null +++ b/tests/unit/core/test_io.py @@ -0,0 +1,49 @@ +# 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 datetime + +import google.cloud.bigquery as bigquery + +import bigframes.core.io + + +def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): + table_ref = bigquery.TableReference.from_string( + "my-test-project._e8166e0cdb.anonbb92cd" + ) + + sql = bigframes.core.io.create_snapshot_sql( + table_ref, datetime.datetime.now(datetime.timezone.utc) + ) + + # Anonymous query results tables don't support time travel. + assert "SYSTEM_TIME" not in sql + + # Need fully-qualified table name. + assert "`my-test-project`.`_e8166e0cdb`.`anonbb92cd`" in sql + + +def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): + table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") + + sql = bigframes.core.io.create_snapshot_sql( + table_ref, datetime.datetime.now(datetime.timezone.utc) + ) + + # We aren't modifying _SESSION tables, so don't use time travel. + assert "SYSTEM_TIME" not in sql + + # Don't need the project ID for _SESSION tables. + assert "my-test-project" not in sql diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 6f5c96da49d..d8c8a2d108e 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -146,6 +146,11 @@ def test_ml_evaluate_no_source_produces_correct_sql(): assert sql == """SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`)""" +def test_ml_centroids_produces_correct_sql(): + sql = ml_sql.ml_centroids(model_name="my_dataset.my_model") + assert sql == """SELECT * FROM ML.CENTROIDS(MODEL `my_dataset.my_model`)""" + + def test_ml_generate_text_produces_correct_sql(): sql = ml_sql.ml_generate_text( model_name="my_dataset.my_model", diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index 6975e6edf2b..95822718c3a 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -302,8 +302,49 @@ def expanding(self, *args, **kwargs): class SeriesGroupBy(GroupBy): - pass + def agg(self, func): + """ + Aggregate using one or more operations. + + Args: + func : function, str, list, dict or None + Function to use for aggregating the data. + + Accepted combinations are: + + - string function name + - list of function names, e.g. ``['sum', 'mean']`` + + Returns: + Series or DataFrame + """ + raise NotImplementedError("abstract property") class DataFrameGroupBy(GroupBy): - pass + def agg(self, func, **kwargs): + """ + Aggregate using one or more operations. + + Args: + func (function, str, list, dict or None): + Function to use for aggregating the data. + + Accepted combinations are: + + - string function name + - list of function names, e.g. ``['sum', 'mean']`` + - dict of axis labels -> function names or list of such. + - None, in which case ``**kwargs`` are used with Named Aggregation. Here the + output has one column for each element in ``**kwargs``. The name of the + column is keyword, whereas the value determines the aggregation used to compute + the values in the column. + + kwargs + If ``func`` is None, ``**kwargs`` are used to define the output names and + aggregations via Named Aggregation. See ``func`` entry. + + Returns: + DataFrame + """ + raise NotImplementedError("abstract property") diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 215f7ec4e08..039dc1eae0f 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1572,3 +1572,48 @@ def str(self): after Python’s string methods, with some inspiration from R’s stringr package. """ raise NotImplementedError("abstract property") + + def isin(self, values): + """ + Whether elements in Series are contained in values. + + Return a boolean Series showing whether each element in the Series matches an + element in the passed sequence of values exactly. + + .. note:: + This function treats all NaN-like values(e.g., pd.NA, numpy.nan, None) as + the same. That is, if any form of NaN is present in values, all forms + of NaN in the series will be considered a match. (though pandas may not) + + Args: + values (list-like): + The sequence of values to test. Passing in a single string will raise a + TypeError. Instead, turn a single string into a list of one element. + + Returns: + bigframes.series.Series: Series of booleans indicating if each element is in values. + + Raises: + TypeError: If input is not list-like. + """ + raise NotImplementedError("abstract method") + + @property + def is_monotonic_increasing(self) -> bool: + """ + Return boolean if values in the object are monotonically increasing. + + Returns: + bool + """ + raise NotImplementedError("abstract property") + + @property + def is_monotonic_decreasing(self) -> bool: + """ + Return boolean if values in the object are monotonically decreasing. + + Returns: + bool + """ + raise NotImplementedError("abstract property") diff --git a/third_party/bigframes_vendored/pandas/core/strings/accessor.py b/third_party/bigframes_vendored/pandas/core/strings/accessor.py index e464843c779..a27093b552c 100644 --- a/third_party/bigframes_vendored/pandas/core/strings/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/strings/accessor.py @@ -314,3 +314,116 @@ def endswith( pattern matches the end of each string element. """ raise NotImplementedError("abstract method") + + def match(self, pat: str, case: bool = True, flags: int = 0): + """ + Determine if each string starts with a match of a regular expression. + + Args: + pat (str): + Character sequence or regular expression. + case (bool): + If True, case sensitive. + flags (int, default 0): + Regex module flags, e.g. re.IGNORECASE. + + Returns: + bigframes.series.Series: Series of boolean values + """ + raise NotImplementedError("abstract method") + + def fullmatch(self, pat: str, case: bool = True, flags: int = 0): + """ + Determine if each string entirely matches a regular expression. + + Args: + pat (str): + Character sequence or regular expression. + case (bool): + If True, case sensitive. + flags (int, default 0): + Regex module flags, e.g. re.IGNORECASE. + + Returns: + bigframes.series.Series: Series of boolean values + """ + raise NotImplementedError("abstract method") + + def get(self, i: int): + """ + Extract element from each component at specified position or with specified key. + + Extract element from lists, tuples, dict, or strings in each element in the + Series/Index. + + Args: + i (int): + Position or key of element to extract. + + Returns: + bigframes.series.Series: Series + """ + raise NotImplementedError("abstract method") + + def pad( + self, + width: int, + side: typing.Literal["left", "right", "both"] = "left", + fillchar: str = " ", + ): + """ + Pad strings in the Series/Index up to width. + + Args: + width (int): + Minimum width of resulting string; additional characters will be filled + with character defined in `fillchar`. + side ({'left', 'right', 'both'}, default 'left'): + Side from which to fill resulting string. + fillchar (str, default ' '): + Additional character for filling, default is whitespace. + + Returns: + bigframes.series.Series: Returns Series or Index with minimum number of char in object. + """ + raise NotImplementedError("abstract method") + + def ljust( + self, + width: int, + fillchar: str = " ", + ): + """ + Pad right side of strings in the Series/Index up to width. + + Args: + width (int): + Minimum width of resulting string; additional characters will be filled + with character defined in `fillchar`. + fillchar (str, default ' '): + Additional character for filling, default is whitespace. + + Returns: + bigframes.series.Series: Returns Series or Index with minimum number of char in object. + """ + raise NotImplementedError("abstract method") + + def rjust( + self, + width: int, + fillchar: str = " ", + ): + """ + Pad left side of strings in the Series/Index up to width. + + Args: + width (int): + Minimum width of resulting string; additional characters will be filled + with character defined in `fillchar`. + fillchar (str, default ' '): + Additional character for filling, default is whitespace. + + Returns: + bigframes.series.Series: Returns Series or Index with minimum number of char in object. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index bddb82c7ba0..068aa4d2908 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -101,3 +101,24 @@ def score( bigframes.dataframe.DataFrame: DataFrame of the metrics. """ raise NotImplementedError("abstract method") + + @property + def cluster_centers_(self): + """Information of cluster centers. + + Returns: + bigframes.dataframe.DataFrame: DataFrame of cluster centers, containing following columns: + centroid_id: An integer that identifies the centroid. + + feature: The column name that contains the feature. + + numerical_value: If feature is numeric, the value of feature for the centroid that centroid_id identifies. If feature is not numeric, the value is NULL. + + categorical_value: An list of mappings containing information about categorical features. Each mapping contains the following fields: + categorical_value.category: The name of each category. + + categorical_value.value: The value of categorical_value.category for the centroid that centroid_id identifies. + + The output contains one row per feature per centroid. + """ + raise NotImplementedError("abstract method") From a32b74751785c8e8aec40ce01df639dd7c4fbb77 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 24 Aug 2023 11:04:40 -0500 Subject: [PATCH 06/80] chore: sync latest changes from internal repo (#7) docs: highlight bigframes is open-source docs: correct the return types of Dataframe and Series docs: create subfolders for notebooks feat: add `bigframes.get_global_session()` and `bigframes.reset_session()` aliases chore: mark ml.llm tests flaky chore: make kokoro/build.sh executable feat: add `Series.str` methods `isalpha`, `isdigit`, `isdecimal`, `isalnum`, `isspace`, `islower`, `isupper`, `zfill`, `center` chore: pin max pytest-retry plugin version in tests docs: sample ML Drug Name Generation notebook docs: add samples and best practices to `read_gbq` docs chore: fix Python download path in docs-presubmit tests perf: add local cache for `__repr_*__` methods feat: support `DataFrame.pivot` fix: don't use query cache for Session construction feat: add `bigframes.pandas.read_pickle` function feat: support MultiIndex for DataFrame columns chore: change the docs kokoro setup to Gerrit path docs: transform remote function user guide into sample code fix: raise exception for invalid function in `read_gbq_function` docs: add release status to table of contents feat: add `fit_transform` to `bigquery.ml` transformers feat: use `pandas.Index` for column labels docs: add ML section under Overview fix: check that types are specified in `read_gbq_function` fix: add error message to `set_index` --- .kokoro/build.sh | 0 .kokoro/docker/docs/Dockerfile | 13 +- .kokoro/docs/common.cfg | 4 +- .kokoro/docs/docs-presubmit.cfg | 2 +- README.rst | 126 +- bigframes/__init__.py | 5 +- bigframes/core/blocks.py | 239 +- bigframes/core/global_session.py | 65 + bigframes/core/groupby/__init__.py | 149 +- bigframes/core/indexers.py | 2 +- bigframes/dataframe.py | 120 +- bigframes/ml/base.py | 24 +- bigframes/ml/compose.py | 2 +- bigframes/ml/preprocessing.py | 4 +- bigframes/operations/__init__.py | 96 +- bigframes/operations/aggregations.py | 32 +- bigframes/operations/strings.py | 41 + bigframes/pandas/__init__.py | 138 +- bigframes/remote_function.py | 83 +- bigframes/series.py | 4 +- bigframes/session.py | 42 +- docs/index.rst | 1 - docs/templates/toc.yml | 1 + .../bigframes.pandas/remote_functions.rst | 134 -- docs/user_guide/index.rst | 9 - notebooks/00 - Summary.ipynb | 2060 ----------------- notebooks/01 - Getting Started.ipynb | 1190 ---------- .../dataframe.ipynb} | 4 +- .../longer_ml_demo.ipynb} | 0 ...q_dataframes_ml_drug_name_generation.ipynb | 980 ++++++++ .../large_language_models.ipynb} | 0 .../bq_dataframes_llm_code_generation.ipynb | 891 +++++++ .../bq_dataframes_ml_linear_regression.ipynb | 743 ++++++ .../getting_started_bq_dataframes.ipynb | 971 ++++++++ .../ml_fundamentals.ipynb} | 0 .../regionalized.ipynb} | 0 .../easy_linear_regression.ipynb} | 0 .../sklearn_linear_regression.ipynb} | 0 .../remote_function.ipynb} | 0 noxfile.py | 50 +- pytest.ini | 1 + samples/snippets/remote_function.py | 147 ++ samples/snippets/remote_function_test.py | 32 + tests/data/hockey_players.json | 37 + tests/data/hockey_players.jsonl | 10 + tests/system/conftest.py | 79 + tests/system/large/ml/test_compose.py | 53 +- tests/system/small/ml/test_llm.py | 8 + tests/system/small/ml/test_preprocessing.py | 68 +- tests/system/small/operations/test_strings.py | 125 +- tests/system/small/test_dataframe.py | 44 + tests/system/small/test_groupby.py | 4 - tests/system/small/test_ipython.py | 28 + tests/system/small/test_multiindex.py | 177 ++ tests/system/small/test_pandas_options.py | 5 +- tests/system/small/test_remote_function.py | 151 +- tests/system/small/test_session.py | 38 + tests/unit/test_dtypes.py | 7 +- .../bigframes_vendored/pandas/core/frame.py | 115 +- .../bigframes_vendored/pandas/core/series.py | 91 +- .../pandas/core/strings/accessor.py | 140 ++ .../bigframes_vendored/pandas/io/gbq.py | 51 +- .../bigframes_vendored/pandas/io/pickle.py | 55 + .../bigframes_vendored/sklearn/base.py | 23 +- .../sklearn/cluster/_kmeans.py | 2 - 65 files changed, 5909 insertions(+), 3807 deletions(-) mode change 100644 => 100755 .kokoro/build.sh create mode 100644 bigframes/core/global_session.py delete mode 100644 docs/user_guide/bigframes.pandas/remote_functions.rst delete mode 100644 docs/user_guide/index.rst delete mode 100644 notebooks/00 - Summary.ipynb delete mode 100644 notebooks/01 - Getting Started.ipynb rename notebooks/{02 - DataFrame.ipynb => dataframes/dataframe.ipynb} (99%) rename notebooks/{99 - Longer ML demo.ipynb => experimental/longer_ml_demo.ipynb} (100%) create mode 100644 notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb rename notebooks/{06 - Using ML - Large Language Models.ipynb => generative_ai/large_language_models.ipynb} (100%) create mode 100644 notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb create mode 100644 notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb create mode 100644 notebooks/getting_started/getting_started_bq_dataframes.ipynb rename notebooks/{03 - Using ML - ML fundamentals.ipynb => getting_started/ml_fundamentals.ipynb} (100%) rename notebooks/{10 - Regionalized.ipynb => location/regionalized.ipynb} (100%) rename notebooks/{05 - Using ML - Easy linear regression.ipynb => regression/easy_linear_regression.ipynb} (100%) rename notebooks/{04 - Using ML - SKLearn linear regression.ipynb => regression/sklearn_linear_regression.ipynb} (100%) rename notebooks/{50 - Remote Function.ipynb => remote_functions/remote_function.ipynb} (100%) create mode 100644 samples/snippets/remote_function.py create mode 100644 samples/snippets/remote_function_test.py create mode 100644 tests/data/hockey_players.json create mode 100644 tests/data/hockey_players.jsonl create mode 100644 tests/system/small/test_ipython.py create mode 100644 third_party/bigframes_vendored/pandas/io/pickle.py diff --git a/.kokoro/build.sh b/.kokoro/build.sh old mode 100644 new mode 100755 diff --git a/.kokoro/docker/docs/Dockerfile b/.kokoro/docker/docs/Dockerfile index d300bee2607..e8ee8191ee1 100644 --- a/.kokoro/docker/docs/Dockerfile +++ b/.kokoro/docker/docs/Dockerfile @@ -60,19 +60,16 @@ RUN apt-get update \ && rm -rf /var/lib/apt/lists/* \ && rm -f /var/cache/apt/archives/*.deb -###################### Install python 3.9.13 and 3.10.5 +###################### Install python 3.9.13 -# Download python 3.9.13 and 3.10.5 +# Download python 3.9.13 RUN wget https://www.python.org/ftp/python/3.9.13/Python-3.9.13.tgz -RUN wget https://www.python.org/ftp/python/3.9.13/Python-3.10.5.tgz # Extract files RUN tar -xvf Python-3.9.13.tgz -RUN tar -xvf Python-3.10.5.tgz -# Install python 3.9.13 and 3.10.5 +# Install python 3.9.13 RUN ./Python-3.9.13/configure --enable-optimizations -RUN ./Python-3.10.5/configure --enable-optimizations RUN make altinstall ###################### Install pip @@ -82,7 +79,5 @@ RUN wget -O /tmp/get-pip.py 'https://bootstrap.pypa.io/get-pip.py' \ # Test pip RUN python3 -m pip -RUN python3.9 -m pip -RUN python3.10 -m pip -CMD ["python3.10"] +CMD ["python3.9"] diff --git a/.kokoro/docs/common.cfg b/.kokoro/docs/common.cfg index 3e36916024a..ce84d7ec49a 100644 --- a/.kokoro/docs/common.cfg +++ b/.kokoro/docs/common.cfg @@ -11,7 +11,7 @@ action { 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" +build_file: "bigframes/.kokoro/trampoline_v2.sh" # Configure the docker image for kokoro-trampoline. env_vars: { @@ -20,7 +20,7 @@ env_vars: { } env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: "github/python-bigquery-dataframes/.kokoro/publish-docs.sh" + value: "git/bigframes/.kokoro/publish-docs.sh" } env_vars: { diff --git a/.kokoro/docs/docs-presubmit.cfg b/.kokoro/docs/docs-presubmit.cfg index 43ec87185e4..1d0dc4b4991 100644 --- a/.kokoro/docs/docs-presubmit.cfg +++ b/.kokoro/docs/docs-presubmit.cfg @@ -13,7 +13,7 @@ env_vars: { env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: "github/python-bigquery-dataframes/.kokoro/build.sh" + value: ".kokoro/build.sh" } # Only run this nox session. diff --git a/README.rst b/README.rst index c6dbb059576..6ae3753eed4 100644 --- a/README.rst +++ b/README.rst @@ -7,6 +7,9 @@ powered by the BigQuery engine. * ``bigframes.pandas`` provides a pandas-compatible API for analytics. * ``bigframes.ml`` provides a scikit-learn-like API for ML. +BigQuery DataFrames is an open-source package. You can run +``pip install --upgrade bigframes`` to install the latest version. + Documentation ------------- @@ -65,6 +68,127 @@ querying is not in the US multi-region. If you try to read a table from another location, you get a NotFound exception. +ML Capabilities +--------------- + +The ML capabilities in BigQuery DataFrames let you preprocess data, and +then train models on that data. You can also chain these actions together to +create data pipelines. + +Preprocess data +^^^^^^^^^^^^^^^^^^^^^^^^ + +Create transformers to prepare data for use in estimators (models) by +using the +`bigframes.ml.preprocessing module `_ +and the `bigframes.ml.compose module `_. +BigQuery DataFrames offers the following transformations: + +* Use the `OneHotEncoder class `_ + in the ``bigframes.ml.preprocessing`` module to transform categorical values into numeric format. +* Use the `StandardScaler class `_ + in the ``bigframes.ml.preprocessing`` module to standardize features by removing the mean and scaling to unit variance. +* Use the `ColumnTransformer class `_ + in the ``bigframes.ml.compose`` module to apply transformers to DataFrames columns. + + +Train models +^^^^^^^^^^^^ + +Create estimators to train models in BigQuery DataFrames. + +**Clustering models** + +Create estimators for clustering models by using the +`bigframes.ml.cluster module `_. + +* Use the `KMeans class `_ + to create K-means clustering models. Use these models for + data segmentation. For example, identifying customer segments. K-means is an + unsupervised learning technique, so model training doesn't require labels or split + data for training or evaluation. + +**Decomposition models** + +Create estimators for decomposition models by using the `bigframes.ml.decomposition module `_. + +* Use the `PCA class `_ + to create principal component analysis (PCA) models. Use these + models for computing principal components and using them to perform a change of + basis on the data. This provides dimensionality reduction by projecting each data + point onto only the first few principal components to obtain lower-dimensional + data while preserving as much of the data's variation as possible. + + +**Ensemble models** + +Create estimators for ensemble models by using the `bigframes.ml.ensemble module `_. + +* Use the `RandomForestClassifier class `_ + to create random forest classifier models. Use these models for constructing multiple + learning method decision trees for classification. +* Use the `RandomForestRegressor class `_ + to create random forest regression models. Use + these models for constructing multiple learning method decision trees for regression. +* Use the `XGBClassifier class `_ + to create gradient boosted tree classifier models. Use these models for additively + constructing multiple learning method decision trees for classification. +* Use the `XGBRegressor class `_ + to create gradient boosted tree regression models. Use these models for additively + constructing multiple learning method decision trees for regression. + + +**Forecasting models** + +Create estimators for forecasting models by using the `bigframes.ml.forecasting module `_. + +* Use the `ARIMAPlus class `_ + to create time series forecasting models. + +**Imported models** + +Create estimators for imported models by using the `bigframes.ml.imported module `_. + +* Use the `ONNXModel class `_ + to import Open Neural Network Exchange (ONNX) models. +* Use the `TensorFlowModel class `_ + to import TensorFlow models. + +**Linear models** + +Create estimators for linear models by using the `bigframes.ml.linear_model module `_. + +* Use the `LinearRegression class `_ + to create linear regression models. Use these models for forecasting. For example, + forecasting the sales of an item on a given day. +* Use the `LogisticRegression class `_ + to create logistic regression models. Use these models for the classification of two + or more possible values such as whether an input is ``low-value``, ``medium-value``, + or ``high-value``. + +**Large language models** + +Create estimators for LLMs by using the `bigframes.ml.llm module `_. + +* Use the `PaLM2TextGenerator class `_ to create PaLM2 text generator models. Use these models + for text generation tasks. +* Use the `PaLM2TextEmbeddingGenerator class `_ to create PaLM2 text embedding generator models. + Use these models for text embedding generation tasks. + + +Create pipelines +^^^^^^^^^^^^^^^^ + +Create ML pipelines by using +`bigframes.ml.pipeline module `_. +Pipelines let you assemble several ML steps to be cross-validated together while setting +different parameters. This simplifies your code, and allows you to deploy data preprocessing +steps and an estimator together. + +* Use the `Pipeline class `_ + to create a pipeline of transforms with a final estimator. + + ML locations ------------ @@ -181,7 +305,7 @@ following IAM roles: Quotas and limits ------------------ +------------------ `BigQuery quotas `_ including hardware, software, and network components. diff --git a/bigframes/__init__.py b/bigframes/__init__.py index 2ee745bc52e..3e54a6d0903 100644 --- a/bigframes/__init__.py +++ b/bigframes/__init__.py @@ -16,13 +16,16 @@ from bigframes._config import options from bigframes._config.bigquery_options import BigQueryOptions +from bigframes.core.global_session import get_global_session, reset_session from bigframes.session import connect, Session from bigframes.version import __version__ __all__ = [ + "options", "BigQueryOptions", + "get_global_session", + "reset_session", "connect", - "options", "Session", "__version__", ] diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index f696b8287bc..2731990febe 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -53,6 +53,10 @@ _BYTES_TO_KILOBYTES = 1024 _BYTES_TO_MEGABYTES = _BYTES_TO_KILOBYTES * 1024 +# This is the max limit of physical columns in BQ +# May choose to set smaller limit for number of block columns to allow overhead for ordering, etc. +_BQ_MAX_COLUMNS = 10000 + # All sampling method _HEAD = "head" _UNIFORM = "uniform" @@ -75,9 +79,9 @@ class Block: def __init__( self, expr: core.ArrayValue, - index_columns: Iterable[str] = (), - column_labels: Optional[Sequence[Label]] = None, - index_labels: Optional[Sequence[Label]] = None, + index_columns: Iterable[str], + column_labels: typing.Union[pd.Index, typing.Sequence[Label]], + index_labels: typing.Union[pd.Index, typing.Sequence[Label], None] = None, ): """Construct a block object, will create default index if no index columns specified.""" if index_labels and (len(index_labels) != len(list(index_columns))): @@ -88,15 +92,18 @@ def __init__( expr, new_index_col_id = expr.promote_offsets() index_columns = [new_index_col_id] self._index_columns = tuple(index_columns) + # Index labels don't need complicated hierarchical access so can store as tuple self._index_labels = ( tuple(index_labels) if index_labels else tuple([None for _ in index_columns]) ) self._expr = self._normalize_expression(expr, self._index_columns) - # TODO(tbergeron): Force callers to provide column labels + # Use pandas index to more easily replicate column indexing, especially for hierarchical column index self._column_labels = ( - tuple(column_labels) if column_labels else tuple(self.value_columns) + column_labels.copy() + if isinstance(column_labels, pd.Index) + else pd.Index(column_labels) ) if len(self.value_columns) != len(self._column_labels): raise ValueError( @@ -139,8 +146,8 @@ def value_columns(self) -> Sequence[str]: ] @property - def column_labels(self) -> List[Label]: - return list(self._column_labels) + def column_labels(self) -> pd.Index: + return self._column_labels @property def expr(self) -> core.ArrayValue: @@ -193,6 +200,24 @@ def index_name_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]: mapping[label] = (*mapping.get(label, ()), id) return mapping + def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]: + """ + Unlike label_to_col_id, this works with partial labels for multi-index. + + Only some methods, like __getitem__ can use a partial key to get columns + from a dataframe. These methods should use cols_matching_label, while + methods that require exact label matches should use label_to_col_id. + """ + # TODO(tbergeron): Refactor so that all label lookups use this method + if partial_label not in self.column_labels: + return [] + loc = self.column_labels.get_loc(partial_label) + if isinstance(loc, int): + return [self.value_columns[loc]] + if isinstance(loc, slice): + return self.value_columns[loc] + return [col for col, is_present in zip(self.value_columns, loc) if is_present] + def order_by( self, by: typing.Sequence[ordering.OrderingColumnReference], @@ -237,8 +262,9 @@ def reset_index(self, drop: bool = True) -> Block: index_labels=[None], ) else: + # Add index names to column index index_labels = self.index.names - index_labels_rewritten = [] + column_labels_modified = self.column_labels for level, label in enumerate(index_labels): if label is None: if "index" not in self.column_labels: @@ -248,12 +274,17 @@ def reset_index(self, drop: bool = True) -> Block: if label in self.column_labels: raise ValueError(f"cannot insert {label}, already exists") - index_labels_rewritten.append(label) + if isinstance(self.column_labels, pd.MultiIndex): + nlevels = self.column_labels.nlevels + label = tuple(label if i == 0 else "" for i in range(nlevels)) + # Create index copy with label inserted + # See: https://pandas.pydata.org/docs/reference/api/pandas.Index.insert.html + column_labels_modified = column_labels_modified.insert(level, label) block = Block( expr, index_columns=[new_index_col_id], - column_labels=[*index_labels_rewritten, *self.column_labels], + column_labels=column_labels_modified, index_labels=[None], ) return block @@ -568,8 +599,11 @@ def _apply_value_keys_to_expr(self, value_keys: Optional[Iterable[str]] = None): expr = expr.select_columns(itertools.chain(self._index_columns, value_keys)) return expr - def with_column_labels(self, value: typing.Iterable[Label]) -> Block: - label_list = tuple(value) + def with_column_labels( + self, + value: typing.Union[pd.Index, typing.Iterable[Label]], + ) -> Block: + label_list = value.copy() if isinstance(value, pd.Index) else pd.Index(value) if len(label_list) != len(self.value_columns): raise ValueError( f"The column labels size `{len(label_list)} ` should equal to the value" @@ -742,7 +776,9 @@ def create_constant( ) -> typing.Tuple[Block, str]: result_id = guid.generate_guid() expr = self.expr.assign_constant(result_id, scalar_constant, dtype=dtype) - labels = [*self.column_labels, label] + # Create index copy with label inserted + # See: https://pandas.pydata.org/docs/reference/api/pandas.Index.insert.html + labels = self.column_labels.insert(len(self.column_labels), label) return ( Block( expr, @@ -755,8 +791,11 @@ def create_constant( def assign_label(self, column_id: str, new_label: Label) -> Block: col_index = self.value_columns.index(column_id) - new_labels = list(self.column_labels) - new_labels[col_index] = new_label + # Create index copy with label inserted + # See: https://pandas.pydata.org/docs/reference/api/pandas.Index.insert.html + new_labels = self.column_labels.insert(col_index, new_label).delete( + col_index + 1 + ) return self.with_column_labels(new_labels) def filter(self, column_name: str, keep_null: bool = False): @@ -790,7 +829,7 @@ def aggregate_all_and_pivot( result_expr = self.expr.aggregate( aggregations, dropna=dropna ).unpivot_single_row( - row_labels=self.column_labels, + row_labels=self.column_labels.to_list(), index_col_id="index", unpivot_columns=[(value_col_id, self.value_columns)], dtype=dtype, @@ -818,11 +857,28 @@ def drop_columns(self, ids_to_drop: typing.Sequence[str]) -> Block: labels = self._get_labels_for_columns(remaining_value_col_ids) return Block(expr, self.index_columns, labels, self.index.names) - def rename(self, *, columns: typing.Mapping[Label, Label]): - # TODO(tbergeron) Support function(Callable) as columns parameter. - col_labels = [ - (columns.get(col_label, col_label)) for col_label in self.column_labels - ] + def rename( + self, + *, + columns: typing.Mapping[Label, Label] | typing.Callable[[typing.Any], Label], + ): + if isinstance(columns, typing.Mapping): + + def remap_f(x): + return columns.get(x, x) + + else: + remap_f = columns + if isinstance(self.column_labels, pd.MultiIndex): + col_labels: list[Label] = [] + for col_label in self.column_labels: + # Mapper applies to each level separately + modified_label = tuple(remap_f(part) for part in col_label) + col_labels.append(modified_label) + else: + col_labels = [] + for col_label in self.column_labels: + col_labels.append(remap_f(col_label)) return self.with_column_labels(col_labels) def aggregate( @@ -874,10 +930,16 @@ def aggregate( ] by_column_labels = self._get_labels_for_columns(by_value_columns) labels = (*by_column_labels, *aggregate_labels) - result_expr_pruned = result_expr.select_columns( + result_expr_pruned, offsets_id = result_expr.select_columns( [*by_value_columns, *output_col_ids] + ).promote_offsets() + + return ( + Block( + result_expr_pruned, index_columns=[offsets_id], column_labels=labels + ), + output_col_ids, ) - return Block(result_expr_pruned, column_labels=labels), output_col_ids def get_stat(self, column_id: str, stat: agg_ops.AggregateOp): """Gets aggregates immediately, and caches it""" @@ -891,7 +953,12 @@ def get_stat(self, column_id: str, stat: agg_ops.AggregateOp): aggregations = [(column_id, stat, stat.name) for stat in stats_to_fetch] expr = self.expr.aggregate(aggregations) - block = Block(expr, column_labels=[s.name for s in stats_to_fetch]) + expr, offset_index_id = expr.promote_offsets() + block = Block( + expr, + index_columns=[offset_index_id], + column_labels=[s.name for s in stats_to_fetch], + ) df, _ = block.to_pandas() # Carefully extract stats such that they aren't coerced to a common type @@ -988,6 +1055,10 @@ def slice( ) return block + # Using cache to optimize for Jupyter Notebook's behavior where both '__repr__' + # and '__repr_html__' are called in a single display action, reducing redundant + # queries. + @functools.cache def retrieve_repr_request_results( self, max_results: int ) -> Tuple[pd.DataFrame, int, bigquery.QueryJob]: @@ -1038,13 +1109,7 @@ def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: index_labels=self.index.names, ) if axis_number == 1: - expr = self._expr - return Block( - self._expr, - index_columns=self.index_columns, - column_labels=[f"{prefix}{label}" for label in self.column_labels], - index_labels=self.index.names, - ) + return self.rename(columns=lambda label: f"{prefix}{label}") def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block: axis_number = bigframes.core.utils.get_axis_number(axis) @@ -1061,13 +1126,110 @@ def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block: index_labels=self.index.names, ) if axis_number == 1: - expr = self._expr - return Block( - self._expr, - index_columns=self.index_columns, - column_labels=[f"{label}{suffix}" for label in self.column_labels], - index_labels=self.index.names, + return self.rename(columns=lambda label: f"{label}{suffix}") + + def pivot( + self, + *, + columns: Sequence[str], + values: Sequence[str], + values_in_index: typing.Optional[bool] = None, + ): + # Columns+index should uniquely identify rows + # Warning: This is not validated, breaking this constraint will result in silently non-deterministic behavior. + # -1 to allow for ordering column in addition to pivot columns + max_unique_value = (_BQ_MAX_COLUMNS - 1) // len(values) + columns_values = self._get_unique_values(columns, max_unique_value) + column_index = columns_values + + column_ids: list[str] = [] + block = self + for value in values: + for uvalue in columns_values: + block, masked_id = self._create_pivot_col(block, columns, value, uvalue) + column_ids.append(masked_id) + + block = block.select_columns(column_ids) + aggregations = [(col_id, agg_ops.AnyValueOp()) for col_id in column_ids] + result_block, _ = block.aggregate( + by_column_ids=self.index_columns, + aggregations=aggregations, + as_index=True, + dropna=True, + ) + + if values_in_index or len(values) > 1: + value_labels = self._get_labels_for_columns(values) + column_index = self._create_pivot_column_index(value_labels, columns_values) + else: + column_index = columns_values + + return result_block.with_column_labels(column_index) + + @staticmethod + def _create_pivot_column_index( + value_labels: Sequence[typing.Hashable], columns_values: pd.Index + ): + index_parts = [] + for value in value_labels: + as_frame = columns_values.to_frame() + as_frame.insert(0, None, value) # type: ignore + ipart = pd.MultiIndex.from_frame( + as_frame, names=(None, *columns_values.names) ) + index_parts.append(ipart) + return functools.reduce(lambda x, y: x.append(y), index_parts) + + @staticmethod + def _create_pivot_col( + block: Block, columns: typing.Sequence[str], value_col: str, value + ) -> typing.Tuple[Block, str]: + cond_id = "" + nlevels = len(columns) + for i in range(len(columns)): + uvalue_level = value[i] if nlevels > 1 else value + if pd.isna(uvalue_level): + block, eq_id = block.apply_unary_op( + columns[i], + ops.isnull_op, + ) + else: + block, eq_id = block.apply_unary_op( + columns[i], ops.partial_right(ops.eq_op, uvalue_level) + ) + if cond_id: + block, cond_id = block.apply_binary_op(eq_id, cond_id, ops.and_op) + else: + cond_id = eq_id + block, masked_id = block.apply_binary_op( + value_col, cond_id, ops.partial_arg3(ops.where_op, None) + ) + + return block, masked_id + + def _get_unique_values( + self, columns: Sequence[str], max_unique_values: int + ) -> pd.Index: + """Gets N unique values for a column immediately.""" + # Importing here to avoid circular import + import bigframes.core.block_transforms as block_tf + import bigframes.dataframe as df + + unique_value_block = block_tf.drop_duplicates( + self.select_columns(columns), columns + ) + pd_values = ( + df.DataFrame(unique_value_block).head(max_unique_values + 1).to_pandas() + ) + if len(pd_values) > max_unique_values: + raise ValueError(f"Too many unique values: {pd_values}") + + if len(columns) > 1: + return pd.MultiIndex.from_frame( + pd_values.sort_values(by=list(pd_values.columns), na_position="first") + ) + else: + return pd.Index(pd_values.squeeze(axis=1).sort_values(na_position="first")) def concat( self, @@ -1138,8 +1300,9 @@ def block_from_local(data, session=None, use_index=True) -> Block: ) else: keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) + keys_expr, offsets_id = keys_expr.promote_offsets() # Constructor will create default range index - return Block(keys_expr, column_labels=column_labels) + return Block(keys_expr, index_columns=[offsets_id], column_labels=column_labels) def _align_block_to_schema( diff --git a/bigframes/core/global_session.py b/bigframes/core/global_session.py new file mode 100644 index 00000000000..68529981cda --- /dev/null +++ b/bigframes/core/global_session.py @@ -0,0 +1,65 @@ +# 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. + +"""Utilities for managing a default, globally available Session object.""" + +import threading +from typing import Callable, Optional, TypeVar + +import bigframes._config +import bigframes.session + +_global_session: Optional[bigframes.session.Session] = None +_global_session_lock = threading.Lock() + + +def reset_session() -> None: + """Start a fresh session the next time a function requires a session. + + Closes the current session if it was already started. + + Returns: + None + """ + global _global_session + + with _global_session_lock: + if _global_session is not None: + _global_session.close() + _global_session = None + + bigframes._config.options.bigquery._session_started = False + + +def get_global_session(): + """Gets the global session. + + Creates the global session if it does not exist. + """ + global _global_session, _global_session_lock + + with _global_session_lock: + if _global_session is None: + _global_session = bigframes.session.connect( + bigframes._config.options.bigquery + ) + + return _global_session + + +_T = TypeVar("_T") + + +def with_default_session(func: Callable[..., _T], *args, **kwargs) -> _T: + return func(get_global_session(), *args, **kwargs) diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 10bee4f56d9..5b217effdd1 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -15,7 +15,8 @@ from __future__ import annotations import typing -import warnings + +import pandas as pd import bigframes.constants as constants import bigframes.core as core @@ -102,12 +103,12 @@ def __getitem__( def sum(self, numeric_only: bool = False, *args) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("sum") - return self._aggregate(agg_ops.sum_op, numeric_only=True) + return self._aggregate_all(agg_ops.sum_op, numeric_only=True) def mean(self, numeric_only: bool = False, *args) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("mean") - return self._aggregate(agg_ops.mean_op, numeric_only=True) + return self._aggregate_all(agg_ops.mean_op, numeric_only=True) def median( self, numeric_only: bool = False, *, exact: bool = False @@ -118,13 +119,13 @@ def median( ) if not numeric_only: self._raise_on_non_numeric("median") - return self._aggregate(agg_ops.median_op, numeric_only=True) + return self._aggregate_all(agg_ops.median_op, numeric_only=True) def min(self, numeric_only: bool = False, *args) -> df.DataFrame: - return self._aggregate(agg_ops.min_op, numeric_only=numeric_only) + return self._aggregate_all(agg_ops.min_op, numeric_only=numeric_only) def max(self, numeric_only: bool = False, *args) -> df.DataFrame: - return self._aggregate(agg_ops.max_op, numeric_only=numeric_only) + return self._aggregate_all(agg_ops.max_op, numeric_only=numeric_only) def std( self, @@ -133,7 +134,7 @@ def std( ) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("std") - return self._aggregate(agg_ops.std_op, numeric_only=True) + return self._aggregate_all(agg_ops.std_op, numeric_only=True) def var( self, @@ -142,16 +143,16 @@ def var( ) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("var") - return self._aggregate(agg_ops.var_op, numeric_only=True) + return self._aggregate_all(agg_ops.var_op, numeric_only=True) def all(self) -> df.DataFrame: - return self._aggregate(agg_ops.all_op) + return self._aggregate_all(agg_ops.all_op) def any(self) -> df.DataFrame: - return self._aggregate(agg_ops.any_op) + return self._aggregate_all(agg_ops.any_op) def count(self) -> df.DataFrame: - return self._aggregate(agg_ops.count_op) + return self._aggregate_all(agg_ops.count_op) def cumsum(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: if not numeric_only: @@ -168,71 +169,97 @@ def cumprod(self, *args, **kwargs) -> df.DataFrame: return self._apply_window_op(agg_ops.product_op, numeric_only=True) def agg(self, func=None, **kwargs) -> df.DataFrame: - column_labels = [] if func: - warnings.warn( - "DataFrameGroupby aggregate produces single-level column labels only currently. Subject to change in future versions." - ) if isinstance(func, str): - aggregations = [ - (col_id, agg_ops.AGGREGATIONS_LOOKUP[func]) - for col_id in self._aggregated_columns() - ] + return self._agg_string(func) elif utils.is_dict_like(func): - aggregations = [] - for label, funcs_for_id in func.items(): - col_id = self._resolve_label(label) - func_list = ( - funcs_for_id - if utils.is_list_like(funcs_for_id) - else [funcs_for_id] - ) - for f in func_list: - aggregations.append((col_id, agg_ops.AGGREGATIONS_LOOKUP[f])) - # Pandas creates multi-index here instead - column_labels.append(f"{label}_{f}") + return self._agg_dict(func) elif utils.is_list_like(func): - aggregations = [ - (col_id, agg_ops.AGGREGATIONS_LOOKUP[f]) - for col_id in self._aggregated_columns() - for f in func - ] - column_labels = [ - f"{self._block.col_id_to_label[col_id]}_{f}" - for col_id in self._aggregated_columns() - for f in func - ] + return self._agg_list(func) else: raise NotImplementedError( f"Aggregate with {func} not supported. {constants.FEEDBACK_LINK}" ) else: - aggregations = [] - for k, v in kwargs.items(): - if not isinstance(k, str): - raise NotImplementedError( - f"Only string aggregate names supported. {constants.FEEDBACK_LINK}" - ) - if not hasattr(v, "column") or not hasattr(v, "aggfunc"): - import bigframes.pandas as bpd - - raise NotImplementedError( - f"kwargs values must be {bpd.NamedAgg.__qualname__}" - ) - col_id = self._resolve_label(v.column) - aggregations.append((col_id, agg_ops.AGGREGATIONS_LOOKUP[v.aggfunc])) - column_labels.append(k) + return self._agg_named(**kwargs) + def _agg_string(self, func: str) -> df.DataFrame: + aggregations = [ + (col_id, agg_ops.lookup_agg_func(func)) + for col_id in self._aggregated_columns() + ] agg_block, _ = self._block.aggregate( by_column_ids=self._by_col_ids, aggregations=aggregations, as_index=self._as_index, dropna=self._dropna, ) + return df.DataFrame(agg_block) - if column_labels: - agg_block = agg_block.with_column_labels(column_labels) + def _agg_dict(self, func: typing.Mapping) -> df.DataFrame: + aggregations = [] + column_labels = [] + for label, funcs_for_id in func.items(): + col_id = self._resolve_label(label) + func_list = ( + funcs_for_id if utils.is_list_like(funcs_for_id) else [funcs_for_id] + ) + for f in func_list: + aggregations.append((col_id, agg_ops.lookup_agg_func(f))) + column_labels.append((col_id, f)) + agg_block, _ = self._block.aggregate( + by_column_ids=self._by_col_ids, + aggregations=aggregations, + as_index=self._as_index, + dropna=self._dropna, + ) + agg_block = agg_block.with_column_labels( + pd.MultiIndex.from_tuples(column_labels) + ) + return df.DataFrame(agg_block) + def _agg_list(self, func: typing.Sequence) -> df.DataFrame: + aggregations = [ + (col_id, agg_ops.lookup_agg_func(f)) + for col_id in self._aggregated_columns() + for f in func + ] + column_labels = [ + (col_id, f) for col_id in self._aggregated_columns() for f in func + ] + agg_block, _ = self._block.aggregate( + by_column_ids=self._by_col_ids, + aggregations=aggregations, + as_index=self._as_index, + dropna=self._dropna, + ) + agg_block = agg_block.with_column_labels( + pd.MultiIndex.from_tuples(column_labels) + ) + return df.DataFrame(agg_block) + + def _agg_named(self, **kwargs) -> df.DataFrame: + aggregations = [] + column_labels = [] + for k, v in kwargs.items(): + if not isinstance(k, str): + raise NotImplementedError( + f"Only string aggregate names supported. {constants.FEEDBACK_LINK}" + ) + if not hasattr(v, "column") or not hasattr(v, "aggfunc"): + import bigframes.pandas as bpd + + raise TypeError(f"kwargs values must be {bpd.NamedAgg.__qualname__}") + col_id = self._resolve_label(v.column) + aggregations.append((col_id, agg_ops.lookup_agg_func(v.aggfunc))) + column_labels.append(k) + agg_block, _ = self._block.aggregate( + by_column_ids=self._by_col_ids, + aggregations=aggregations, + as_index=self._as_index, + dropna=self._dropna, + ) + agg_block = agg_block.with_column_labels(column_labels) return df.DataFrame(agg_block) aggregate = agg @@ -261,7 +288,7 @@ def _column_type(self, col_id: str) -> dtypes.Dtype: dtype = self._block.dtypes[col_offset] return dtype - def _aggregate( + def _aggregate_all( self, aggregate_op: agg_ops.AggregateOp, numeric_only: bool = False ) -> df.DataFrame: aggregated_col_ids = self._aggregated_columns(numeric_only=numeric_only) @@ -359,11 +386,11 @@ def prod(self, *args) -> series.Series: def agg(self, func=None) -> typing.Union[df.DataFrame, series.Series]: column_names: list[str] = [] if isinstance(func, str): - aggregations = [(self._value_column, agg_ops.AGGREGATIONS_LOOKUP[func])] + aggregations = [(self._value_column, agg_ops.lookup_agg_func(func))] column_names = [func] elif utils.is_list_like(func): aggregations = [ - (self._value_column, agg_ops.AGGREGATIONS_LOOKUP[f]) for f in func + (self._value_column, agg_ops.lookup_agg_func(f)) for f in func ] column_names = list(func) else: diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 0aaf169bea4..46091f211ac 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -29,7 +29,7 @@ import bigframes.series if typing.TYPE_CHECKING: - LocSingleKey = bigframes.series.Series | indexes.Index | slice + LocSingleKey = typing.Union[bigframes.series.Series, indexes.Index, slice] class LocSeriesIndexer: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index d4b6e470255..5fbe5d1f9eb 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -174,7 +174,11 @@ def __init__( self._query_job: Optional[bigquery.QueryJob] = None def __dir__(self): - return dir(type(self)) + self._block.column_labels + return dir(type(self)) + [ + label + for label in self._block.column_labels + if label and isinstance(label, str) + ] def _ipython_key_completions_(self) -> List[str]: return list( @@ -201,13 +205,16 @@ def _find_indices( col_ids = self._sql_names(columns, tolerance) return [self._block.value_columns.index(col_id) for col_id in col_ids] - def _resolve_label_exact(self, label) -> str: + def _resolve_label_exact(self, label) -> Optional[str]: + """Returns the column id matching the label if there is exactly + one such column. If there are multiple columns with the same name, + raises an error. If there is no such column, returns None.""" matches = self._block.label_to_col_id.get(label, []) - if len(matches) != 1: + if len(matches) > 1: raise ValueError( - f"Index data must be 1-dimensional. {constants.FEEDBACK_LINK}" + f"Multiple columns matching id {label} were found. {constants.FEEDBACK_LINK}" ) - return matches[0] + return matches[0] if len(matches) != 0 else None def _sql_names( self, @@ -215,7 +222,11 @@ def _sql_names( tolerance: bool = False, ) -> Sequence[str]: """Retrieve sql name (column name in BQ schema) of column(s).""" - labels = columns if utils.is_list_like(columns) else [columns] # type:ignore + labels = ( + columns + if utils.is_list_like(columns) and not isinstance(columns, tuple) + else [columns] + ) # type:ignore results: Sequence[str] = [] for label in labels: col_ids = self._block.label_to_col_id.get(label, []) @@ -246,6 +257,11 @@ def dtypes(self) -> pandas.Series: def columns(self) -> pandas.Index: return self.dtypes.index + @columns.setter + def columns(self, labels: pandas.Index): + new_block = self._block.with_column_labels(labels) + self._set_block(new_block) + @property def shape(self) -> Tuple[int, int]: return self._block.shape @@ -295,7 +311,7 @@ def _to_sql_query( # Has to be unordered as it is impossible to order the sql without # including metadata columns in selection with ibis. ibis_expr = self._block.expr.to_ibis_expr(ordering_mode="unordered") - column_labels = self._block.column_labels + column_labels = list(self._block.column_labels) # TODO(swast): Need to have a better way of controlling when to include # the index or not. @@ -387,11 +403,8 @@ def __getitem__( if isinstance(key, bigframes.series.Series): return self._getitem_bool_series(key) - sql_names = self._sql_names(key) - # Only input is a single key and only find one column, returns a Series - if (not utils.is_list_like(key)) and len(sql_names) == 1: - return bigframes.series.Series(self._block.select_column(sql_names[0])) - + if isinstance(key, typing.Hashable): + return self._getitem_label(key) # Select a subset of columns or re-order columns. # In Ibis after you apply a projection, any column objects from the # table before the projection can't be combined with column objects @@ -410,11 +423,31 @@ def __getitem__( selected_ids: Tuple[str, ...] = () for label in key: - col_ids = self._block.label_to_col_id.get(label, []) + col_ids = self._block.label_to_col_id[label] selected_ids = (*selected_ids, *col_ids) return DataFrame(self._block.select_columns(selected_ids)) + def _getitem_label(self, key: blocks.Label): + col_ids = self._block.cols_matching_label(key) + if len(col_ids) == 0: + raise KeyError(key) + block = self._block.select_columns(col_ids) + if isinstance(self.columns, pandas.MultiIndex): + # Multiindex should drop-level if not selecting entire + key_levels = len(key) if isinstance(key, tuple) else 1 + index_levels = self.columns.nlevels + if key_levels < index_levels: + block = block.with_column_labels( + block.column_labels.droplevel(list(range(key_levels))) + ) + # Force return DataFrame in this case, even if only single column + return DataFrame(block) + + if len(col_ids) == 1: + return bigframes.series.Series(block) + return DataFrame(block) + # Bool Series selects rows def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame: if not key.dtype == pandas.BooleanDtype(): @@ -736,7 +769,7 @@ def drop( *, axis: typing.Union[int, str] = 0, index: typing.Any = None, - columns: Union[blocks.Label, Iterable[blocks.Label]] = None, + columns: Union[blocks.Label, Sequence[blocks.Label]] = None, level: typing.Optional[LevelType] = None, ) -> DataFrame: if labels: @@ -767,10 +800,6 @@ def drop( self._block.value_columns ) if columns: - if not utils.is_list_like(columns): - columns = [columns] # type:ignore - columns = list(columns) - block = block.drop_columns(self._sql_names(columns)) if not index and not columns: raise ValueError("Must specify 'labels' or 'index'/'columns") @@ -849,7 +878,7 @@ def _assign_single_item( def _assign_scalar(self, label: str, value: Union[int, float]) -> DataFrame: # TODO(swast): Make sure that k is the ID / SQL name, not a label, # which could be invalid SQL. - col_ids = self._sql_names(label, tolerance=True) + col_ids = self._block.cols_matching_label(label) block, constant_col_id = self._block.create_constant(value, label) for col_id in col_ids: @@ -868,7 +897,7 @@ def _assign_series_join_on_index( ) column_ids = [ - get_column_left(col_id) for col_id in self._sql_names(label, tolerance=True) + get_column_left(col_id) for col_id in self._block.cols_matching_label(label) ] block = joined_index._block source_column = get_column_right(series._value_column) @@ -903,7 +932,12 @@ def set_index( else: keys = typing.cast(typing.Sequence[blocks.Label], tuple(keys)) col_ids = [self._resolve_label_exact(key) for key in keys] - return DataFrame(self._block.set_index(col_ids, append=append, drop=drop)) + missing = [keys[i] for i in range(len(col_ids)) if col_ids[i] is None] + if len(missing) > 0: + raise KeyError(f"None of {missing} are in the columns") + # convert col_ids to non-optional strs since we just determined they are not None + col_ids_strs: List[str] = [col_id for col_id in col_ids if col_id is not None] + return DataFrame(self._block.set_index(col_ids_strs, append=append, drop=drop)) def sort_index( self, ascending: bool = True, na_position: Literal["first", "last"] = "last" @@ -932,7 +966,7 @@ def sort_values( if na_position not in {"first", "last"}: raise ValueError("Param na_position must be one of 'first' or 'last'") - sort_labels = tuple(by) if utils.is_list_like(by) else (by,) + sort_labels = list(by) if utils.is_list_like(by) else [by] sort_column_ids = self._sql_names(sort_labels) len_by = len(sort_labels) @@ -982,9 +1016,11 @@ def value_counts( return bigframes.series.Series(block) def add_prefix(self, prefix: str, axis: int | str | None = None) -> DataFrame: + axis = 1 if axis is None else axis return DataFrame(self._get_block().add_prefix(prefix, axis)) def add_suffix(self, suffix: str, axis: int | str | None = None) -> DataFrame: + axis = 1 if axis is None else axis return DataFrame(self._get_block().add_suffix(suffix, axis)) def dropna(self) -> DataFrame: @@ -1115,7 +1151,7 @@ def agg( raise NotImplementedError( f"Multiple aggregations only supported on numeric columns. {constants.FEEDBACK_LINK}" ) - aggregations = [agg_ops.AGGREGATIONS_LOOKUP[f] for f in func] + aggregations = [agg_ops.lookup_agg_func(f) for f in func] return DataFrame( self._block.summarize( self._block.value_columns, @@ -1125,7 +1161,7 @@ def agg( else: return bigframes.series.Series( self._block.aggregate_all_and_pivot( - agg_ops.AGGREGATIONS_LOOKUP[typing.cast(str, func)] + agg_ops.lookup_agg_func(typing.cast(str, func)) ) ) @@ -1142,6 +1178,37 @@ def describe(self) -> DataFrame: ) return typing.cast(DataFrame, result) + def pivot( + self, + *, + columns: typing.Union[blocks.Label, Sequence[blocks.Label]], + index: typing.Optional[ + typing.Union[blocks.Label, Sequence[blocks.Label]] + ] = None, + values: typing.Optional[ + typing.Union[blocks.Label, Sequence[blocks.Label]] + ] = None, + ) -> DataFrame: + if index: + block = self.set_index(index)._block + else: + block = self._block + + column_ids = self._sql_names(columns) + if values: + value_col_ids = self._sql_names(values) + else: + value_col_ids = [ + col for col in block.value_columns if col not in column_ids + ] + + pivot_block = block.pivot( + columns=column_ids, + values=value_col_ids, + values_in_index=utils.is_list_like(values), + ) + return DataFrame(pivot_block) + def _drop_non_numeric(self, keep_bool=True) -> DataFrame: types_to_keep = set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) if not keep_bool: @@ -1263,7 +1330,10 @@ def merge( ) # Constructs default index - block = blocks.Block(expr, column_labels=labels) + expr, offset_index_id = expr.promote_offsets() + block = blocks.Block( + expr, index_columns=[offset_index_id], column_labels=labels + ) return DataFrame(block) def _get_merged_col_labels( diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py index e4c68eb17c2..9f9d9f85d0e 100644 --- a/bigframes/ml/base.py +++ b/bigframes/ml/base.py @@ -22,9 +22,10 @@ """ import abc -from typing import cast, Optional, TypeVar +from typing import cast, Optional, TypeVar, Union from bigframes.ml import core +import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.base @@ -143,3 +144,24 @@ def score(self, X, y): @abc.abstractmethod def to_gbq(self, model_name, replace): pass + + +class Transformer(BaseEstimator): + """A BigQuery DataFrames Transformer base class that transforms data. + + Also the transformers can be attached to a pipeline with a predictor.""" + + @abc.abstractmethod + def fit(self, X, y): + pass + + @abc.abstractmethod + def transform(self, X): + pass + + def fit_transform( + self, + X: Union[bpd.DataFrame, bpd.Series], + y: Optional[Union[bpd.DataFrame, bpd.Series]] = None, + ) -> bpd.DataFrame: + return self.fit(X, y).transform(X) diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index 49b4899bebd..df01303ffac 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -34,7 +34,7 @@ class ColumnTransformer( third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, - base.BaseEstimator, + base.Transformer, ): __doc__ = ( third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer.__doc__ diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 500a9fcb244..ee46a370520 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -30,7 +30,7 @@ class StandardScaler( third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, - base.BaseEstimator, + base.Transformer, ): __doc__ = ( third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler.__doc__ @@ -106,7 +106,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: class OneHotEncoder( third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, - base.BaseEstimator, + base.Transformer, ): # BQML max value https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder#syntax TOP_K_DEFAULT = 1000000 diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 7bdd97812e9..58f19ea8e7f 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -76,6 +76,12 @@ def _as_ibis(self, x: ibis_types.Value): return x.notnull() +class HashOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.IntegerValue, x).hash() + + +## String Operation class ReverseOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.StringValue, x).reverse() @@ -100,7 +106,58 @@ class IsNumericOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): # catches all members of the Unicode number class, which matches pandas isnumeric # see https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#regexp_contains - return typing.cast(ibis_types.StringValue, x).re_search(r"^(\pN*)$") + # TODO: Validate correctness, my miss eg ⅕ character + return typing.cast(ibis_types.StringValue, x).re_search(r"^(\pN+)$") + + +class IsAlphaOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_search( + r"^(\p{Lm}|\p{Lt}|\p{Lu}|\p{Ll}|\p{Lo})+$" + ) + + +class IsDigitOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + # Based on docs, should include superscript/subscript-ed numbers + # Tests however pass only when set to Nd unicode class + return typing.cast(ibis_types.StringValue, x).re_search(r"^(\p{Nd})+$") + + +class IsDecimalOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_search(r"^(\p{Nd})+$") + + +class IsAlnumOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.StringValue, x).re_search( + r"^(\p{N}|\p{Lm}|\p{Lt}|\p{Lu}|\p{Ll}|\p{Lo})+$" + ) + + +class IsSpaceOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + # All characters are whitespace characters, False for empty string + return typing.cast(ibis_types.StringValue, x).re_search(r"^\s+$") + + +class IsLowerOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + # No upper case characters, min one cased character + # See: https://docs.python.org/3/library/stdtypes.html#str + return typing.cast(ibis_types.StringValue, x).re_search( + r"\p{Ll}" + ) & ~typing.cast(ibis_types.StringValue, x).re_search(r"\p{Lu}|\p{Lt}") + + +class IsUpperOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + # No lower case characters, min one cased character + # See: https://docs.python.org/3/library/stdtypes.html#str + return typing.cast(ibis_types.StringValue, x).re_search( + r"\p{Lu}" + ) & ~typing.cast(ibis_types.StringValue, x).re_search(r"\p{Ll}|\p{Lt}") class RstripOp(UnaryOp): @@ -227,11 +284,25 @@ def _as_ibis(self, x: ibis_types.Value): return any_match if any_match is not None else ibis_types.literal(False) -class HashOp(UnaryOp): +class ZfillOp(UnaryOp): + def __init__(self, width: int): + self._width = width + def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.IntegerValue, x).hash() + str_value = typing.cast(ibis_types.StringValue, x) + return ( + ibis.case() + .when( + str_value[0] == "-", + "-" + + StrPadOp(self._width - 1, "0", "left")._as_ibis(str_value.substr(1)), + ) + .else_(StrPadOp(self._width, "0", "left")._as_ibis(str_value)) + .end() + ) +## Datetime Ops class DayOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.TimestampValue, x).day() @@ -390,7 +461,14 @@ def _as_ibis(self, x: ibis_types.Value): lower_op = LowerOp() upper_op = UpperOp() strip_op = StripOp() +isalnum_op = IsAlnumOp() +isalpha_op = IsAlphaOp() +isdecimal_op = IsDecimalOp() +isdigit_op = IsDigitOp() isnumeric_op = IsNumericOp() +isspace_op = IsSpaceOp() +islower_op = IsLowerOp() +isupper_op = IsUpperOp() rstrip_op = RstripOp() lstrip_op = LstripOp() hash_op = HashOp() @@ -692,6 +770,18 @@ def clip_op( ) +def partial_arg1(op: TernaryOp, scalar: typing.Any) -> BinaryOp: + return lambda x, y: op(dtypes.literal_to_ibis_scalar(scalar, validate=False), x, y) + + +def partial_arg2(op: TernaryOp, scalar: typing.Any) -> BinaryOp: + return lambda x, y: op(x, dtypes.literal_to_ibis_scalar(scalar, validate=False), y) + + +def partial_arg3(op: TernaryOp, scalar: typing.Any) -> BinaryOp: + return lambda x, y: op(x, y, dtypes.literal_to_ibis_scalar(scalar, validate=False)) + + def is_null(value) -> bool: # float NaN/inf should be treated as distinct from 'true' null values return typing.cast(bool, pd.isna(value)) and not isinstance(value, float) diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py index 1687f705a1e..874c264194c 100644 --- a/bigframes/operations/aggregations.py +++ b/bigframes/operations/aggregations.py @@ -251,6 +251,21 @@ def skips_nulls(self): return False +class AnyValueOp(AggregateOp): + # Warning: only use if all values are equal. Non-deterministic otherwise. + # Do not expose to users. For special cases only (e.g. pivot). + name = "any_value" + + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + return _apply_window_if_present(column.arbitrary(), window) + + @property + def skips_nulls(self): + return True + + class RankOp(WindowOp): name = "rank" @@ -381,7 +396,7 @@ def _map_to_literal( # TODO: Alternative names and lookup from numpy function objects -AGGREGATIONS_LOOKUP: dict[str, AggregateOp] = { +_AGGREGATIONS_LOOKUP: dict[str, AggregateOp] = { op.name: op for op in [ sum_op, @@ -401,3 +416,18 @@ def _map_to_literal( ApproxQuartilesOp(3), ] } + + +def lookup_agg_func(key: str) -> AggregateOp: + if callable(key): + raise NotImplementedError( + "Aggregating with callable object not supported, pass method name as string instead (eg. 'sum' instead of np.sum)." + ) + if not isinstance(key, str): + raise ValueError( + f"Cannot aggregate using object of type: {type(key)}. Use string method name (eg. 'sum')" + ) + if key in _AGGREGATIONS_LOOKUP: + return _AGGREGATIONS_LOOKUP[key] + else: + raise ValueError(f"Unrecognize aggregate function: {key}") diff --git a/bigframes/operations/strings.py b/bigframes/operations/strings.py index 0f1395c78f2..0545ea34d6a 100644 --- a/bigframes/operations/strings.py +++ b/bigframes/operations/strings.py @@ -70,6 +70,41 @@ def upper(self) -> series.Series: def isnumeric(self) -> series.Series: return self._apply_unary_op(ops.isnumeric_op) + def isalpha( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isalpha_op) + + def isdigit( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isdigit_op) + + def isdecimal( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isdecimal_op) + + def isalnum( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isalnum_op) + + def isspace( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isspace_op) + + def islower( + self, + ) -> series.Series: + return self._apply_unary_op(ops.islower_op) + + def isupper( + self, + ) -> series.Series: + return self._apply_unary_op(ops.isupper_op) + def rstrip(self) -> series.Series: return self._apply_unary_op(ops.rstrip_op) @@ -183,6 +218,12 @@ def endswith( pat = (pat,) return self._apply_unary_op(ops.EndsWithOp(pat)) + def zfill(self, width: int) -> series.Series: + return self._apply_unary_op(ops.ZfillOp(width)) + + def center(self, width: int, fillchar: str = " ") -> series.Series: + return self._apply_unary_op(ops.StrPadOp(width, fillchar, "both")) + def cat( self, others: Union[str, series.Series], diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index ed7a09e7b71..b688c18723b 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -18,7 +18,6 @@ from collections import namedtuple import inspect -import threading import typing from typing import ( Any, @@ -32,15 +31,21 @@ Optional, Sequence, Tuple, - TypeVar, Union, ) from google.cloud import bigquery import numpy import pandas +from pandas._typing import ( + CompressionOptions, + FilePath, + ReadPickleBuffer, + StorageOptions, +) import bigframes._config as config +import bigframes.core.global_session as global_session import bigframes.core.indexes import bigframes.core.reshape import bigframes.dataframe @@ -49,14 +54,6 @@ import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat import third_party.bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile -# Support pandas dtype attribute -NA = pandas.NA -BooleanDtype = pandas.BooleanDtype -Float64Dtype = pandas.Float64Dtype -Int64Dtype = pandas.Int64Dtype -StringDtype = pandas.StringDtype -ArrowDtype = pandas.ArrowDtype - # Include method definition so that the method appears in our docs for # bigframes.pandas general functions. @@ -135,52 +132,6 @@ def cut( cut.__doc__ = vendored_pandas_tile.cut.__doc__ -options = config.options -"""Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames.""" - -_global_session: Optional[bigframes.session.Session] = None -_global_session_lock = threading.Lock() - - -def reset_session() -> None: - """Start a fresh session the next time a function requires a session. - - Closes the current session if it was already started. - - Returns: - None - """ - global _global_session - - with _global_session_lock: - if _global_session is not None: - _global_session.close() - _global_session = None - - options.bigquery._session_started = False - - -def get_global_session(): - """Gets the global session. - - Creates the global session if it does not exist. - """ - global _global_session, _global_session_lock - - with _global_session_lock: - if _global_session is None: - _global_session = bigframes.session.connect(options.bigquery) - - return _global_session - - -_T = TypeVar("_T") - - -def _with_default_session(func: Callable[..., _T], *args, **kwargs) -> _T: - return func(get_global_session(), *args, **kwargs) - - def _set_default_session_location_if_possible(query): # Set the location as per the query if this is the first query the user is # running and: @@ -257,7 +208,7 @@ def read_csv( encoding: Optional[str] = None, **kwargs, ) -> bigframes.dataframe.DataFrame: - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_csv, filepath_or_buffer=filepath_or_buffer, sep=sep, @@ -283,7 +234,7 @@ def read_gbq( max_results: Optional[int] = None, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query) - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_gbq, query, index_col=index_col, @@ -296,7 +247,7 @@ def read_gbq( def read_gbq_model(model_name: str): - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_gbq_model, model_name, ) @@ -313,7 +264,7 @@ def read_gbq_query( max_results: Optional[int] = None, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query) - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_gbq_query, query, index_col=index_col, @@ -333,7 +284,7 @@ def read_gbq_table( max_results: Optional[int] = None, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query) - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_gbq_table, query, index_col=index_col, @@ -346,7 +297,7 @@ def read_gbq_table( def read_pandas(pandas_dataframe: pandas.DataFrame) -> bigframes.dataframe.DataFrame: - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_pandas, pandas_dataframe, ) @@ -355,8 +306,24 @@ def read_pandas(pandas_dataframe: pandas.DataFrame) -> bigframes.dataframe.DataF read_pandas.__doc__ = inspect.getdoc(bigframes.session.Session.read_pandas) +def read_pickle( + filepath_or_buffer: FilePath | ReadPickleBuffer, + compression: CompressionOptions = "infer", + storage_options: StorageOptions = None, +): + return global_session.with_default_session( + bigframes.session.Session.read_pickle, + filepath_or_buffer=filepath_or_buffer, + compression=compression, + storage_options=storage_options, + ) + + +read_pickle.__doc__ = inspect.getdoc(bigframes.session.Session.read_pickle) + + def read_parquet(path: str | IO["bytes"]) -> bigframes.dataframe.DataFrame: - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_parquet, path, ) @@ -372,7 +339,7 @@ def remote_function( bigquery_connection: Optional[str] = None, reuse: bool = True, ): - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.remote_function, input_types=input_types, output_type=output_type, @@ -386,7 +353,7 @@ def remote_function( def read_gbq_function(function_name: str): - return _with_default_session( + return global_session.with_default_session( bigframes.session.Session.read_gbq_function, function_name=function_name, ) @@ -395,25 +362,58 @@ def read_gbq_function(function_name: str): read_gbq_function.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_function) -# Other aliases +# pandas dtype attributes +NA = pandas.NA +BooleanDtype = pandas.BooleanDtype +Float64Dtype = pandas.Float64Dtype +Int64Dtype = pandas.Int64Dtype +StringDtype = pandas.StringDtype +ArrowDtype = pandas.ArrowDtype + +# Class aliases +# TODO(swast): Make these real classes so we can refer to these in type +# checking and docstrings. DataFrame = bigframes.dataframe.DataFrame Index = bigframes.core.indexes.Index Series = bigframes.series.Series -# Used by DataFrameGroupby.agg +# Other public pandas attributes NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"]) +options = config.options +"""Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames.""" + +# Session management APIs +get_global_session = global_session.get_global_session +reset_session = global_session.reset_session + + # Use __all__ to let type checkers know what is part of the public API. __all___ = [ + # Functions "concat", - "DataFrame", - "options", "read_csv", "read_gbq", "read_gbq_function", "read_gbq_model", "read_pandas", + "read_pickle", "remote_function", + # pandas dtype attributes + "NA", + "BooleanDtype", + "Float64Dtype", + "Int64Dtype", + "StringDtype", + "ArrowDtype" + # Class aliases + "DataFrame", + "Index", "Series", + # Other public pandas attributes "NamedAgg", + "options", + # Session management APIs + "get_global_session", + "reset_session", ] diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 27e2b8f7c23..7cf74d63111 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -35,15 +35,10 @@ import cloudpickle import google.api_core.exceptions from google.cloud import bigquery, bigquery_connection_v1, functions_v2 -from google.cloud.bigquery.routine import Routine -from google.cloud.bigquery.standard_sql import StandardSqlTypeNames from ibis.backends.bigquery.compiler import compiles from ibis.backends.bigquery.datatypes import BigQueryType -from ibis.expr.datatypes.core import boolean from ibis.expr.datatypes.core import DataType as IbisDataType from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type -from ibis.expr.datatypes.core import float64, int64 -from ibis.expr.datatypes.core import string as ibis_string import ibis.expr.operations as ops import ibis.expr.rules as rlz @@ -63,11 +58,16 @@ # Input and output types supported by BigQuery DataFrames remote functions. # TODO(shobs): Extend the support to all types supported by BQ remote functions # https://cloud.google.com/bigquery/docs/remote-functions#limitations -_supported_io_ibis_types = {boolean, float64, int64, ibis_string} -TYPE_ERROR_MESSAGE_FORMAT = ( - f"Type {{}} not supported, supported types are {_supported_io_ibis_types}. " - f"{constants.FEEDBACK_LINK}" -) +SUPPORTED_IO_PYTHON_TYPES = {bool, float, int, str} +SUPPORTED_IO_BIGQUERY_TYPEKINDS = { + "BOOLEAN", + "BOOL", + "FLOAT", + "FLOAT64", + "INT64", + "INTEGER", + "STRING", +} def get_remote_function_locations(bq_location): @@ -116,7 +116,7 @@ def routine_ref_to_string_for_query(routine_ref: bigquery.RoutineReference) -> s class IbisSignature(NamedTuple): parameter_names: List[str] - input_types: List[IbisDataType] + input_types: List[Optional[IbisDataType]] output_type: IbisDataType @@ -512,7 +512,7 @@ def remote_function_node( """Creates an Ibis node representing a remote function call.""" fields = { - name: rlz.value(type_) + name: rlz.value(type_) if type_ else rlz.any for name, type_ in zip( ibis_signature.parameter_names, ibis_signature.input_types ) @@ -538,20 +538,22 @@ def f(*args, **kwargs): return f +class UnsupportedTypeError(ValueError): + def __init__(self, type_, supported_types): + self.type = type_ + self.supported_types = supported_types + + def ibis_type_from_python_type(t: type) -> IbisDataType: - ibis_type = python_type_to_bigquery_type(t) - assert ibis_type in _supported_io_ibis_types, TYPE_ERROR_MESSAGE_FORMAT.format( - ibis_type - ) - return ibis_type + if t not in SUPPORTED_IO_PYTHON_TYPES: + raise UnsupportedTypeError(t, SUPPORTED_IO_PYTHON_TYPES) + return python_type_to_bigquery_type(t) -def ibis_type_from_type_kind(tk: StandardSqlTypeNames) -> IbisDataType: - ibis_type = BigQueryType.to_ibis(tk) - assert ibis_type in _supported_io_ibis_types, TYPE_ERROR_MESSAGE_FORMAT.format( - ibis_type - ) - return ibis_type +def ibis_type_from_type_kind(tk: bigquery.StandardSqlTypeNames) -> IbisDataType: + if tk not in SUPPORTED_IO_BIGQUERY_TYPEKINDS: + raise UnsupportedTypeError(tk, SUPPORTED_IO_BIGQUERY_TYPEKINDS) + return BigQueryType.to_ibis(tk) def ibis_signature_from_python_signature( @@ -566,13 +568,18 @@ def ibis_signature_from_python_signature( ) -def ibis_signature_from_routine( - routine: Routine, -) -> IbisSignature: +class ReturnTypeMissingError(ValueError): + pass + + +def ibis_signature_from_routine(routine: bigquery.Routine) -> IbisSignature: + if not routine.return_type: + raise ReturnTypeMissingError + return IbisSignature( parameter_names=[arg.name for arg in routine.arguments], input_types=[ - ibis_type_from_type_kind(arg.data_type.type_kind) + ibis_type_from_type_kind(arg.data_type.type_kind) if arg.data_type else None for arg in routine.arguments ], output_type=ibis_type_from_type_kind(routine.return_type.type_kind), @@ -584,9 +591,7 @@ class DatasetMissingError(ValueError): def get_routine_reference( - routine_ref_str: str, - bigquery_client: bigquery.Client, - session: Optional[Session], + routine_ref_str: str, bigquery_client: bigquery.Client, session: Optional[Session] ) -> bigquery.RoutineReference: try: # Handle cases ".." and @@ -859,7 +864,21 @@ def read_gbq_function( ) # Find the routine and get its arguments. - routine = bigquery_client.get_routine(routine_ref) - ibis_signature = ibis_signature_from_routine(routine) + try: + routine = bigquery_client.get_routine(routine_ref) + except google.api_core.exceptions.NotFound: + raise ValueError(f"Unknown function '{routine_ref}'. {constants.FEEDBACK_LINK}") + + try: + ibis_signature = ibis_signature_from_routine(routine) + except ReturnTypeMissingError: + raise ValueError( + "Function return type must be specified. {constants.FEEDBACK_LINK}" + ) + except UnsupportedTypeError as e: + raise ValueError( + f"Type {e.type} not supported, supported types are {e.supported_types}. " + f"{constants.FEEDBACK_LINK}" + ) return remote_function_node(routine_ref, ibis_signature) diff --git a/bigframes/series.py b/bigframes/series.py index f8f44dc2e64..a1da93dee38 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -651,7 +651,7 @@ def agg(self, func: str | typing.Sequence[str]) -> scalars.Scalar | Series: raise NotImplementedError( f"Multiple aggregations only supported on numeric series. {constants.FEEDBACK_LINK}" ) - aggregations = [agg_ops.AGGREGATIONS_LOOKUP[f] for f in func] + aggregations = [agg_ops.lookup_agg_func(f) for f in func] return Series( self._block.summarize( [self._value_column], @@ -661,7 +661,7 @@ def agg(self, func: str | typing.Sequence[str]) -> scalars.Scalar | Series: else: return self._apply_aggregation( - agg_ops.AGGREGATIONS_LOOKUP[typing.cast(str, func)] + agg_ops.lookup_agg_func(typing.cast(str, func)) ) def skew(self): diff --git a/bigframes/session.py b/bigframes/session.py index 73fdd731063..3ef52507460 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -54,6 +54,12 @@ import ibis.expr.types as ibis_types import numpy as np import pandas +from pandas._typing import ( + CompressionOptions, + FilePath, + ReadPickleBuffer, + StorageOptions, +) import pydata_google_auth import bigframes._config.bigquery_options as bigquery_options @@ -75,6 +81,7 @@ import third_party.bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq import third_party.bigframes_vendored.pandas.io.parquet as third_party_pandas_parquet import third_party.bigframes_vendored.pandas.io.parsers.readers as third_party_pandas_readers +import third_party.bigframes_vendored.pandas.io.pickle as third_party_pandas_pickle _ENV_DEFAULT_PROJECT = "GOOGLE_CLOUD_PROJECT" _APPLICATION_NAME = f"bigframes/{bigframes.version.__version__}" @@ -194,6 +201,7 @@ def _create_cloud_clients( class Session( third_party_pandas_gbq.GBQIOMixin, third_party_pandas_parquet.ParquetIOMixin, + third_party_pandas_pickle.PickleIOMixin, third_party_pandas_readers.ReaderIOMixin, ): """Establishes a BigQuery connection to capture a group of job activities related to @@ -252,6 +260,8 @@ def _create_and_bind_bq_session(self): """Create a BQ session and bind the session id with clients to capture BQ activities: go/bigframes-transient-data""" job_config = bigquery.QueryJobConfig(create_session=True) + # Make sure the session is a new one, not one associated with another query. + job_config.use_query_cache = False query_job = self.bqclient.query( "SELECT 1", job_config=job_config, location=self._location ) @@ -458,8 +468,8 @@ def read_gbq_table( {self.ibis_client.compile(distinct_table)} ) - SELECT (SELECT COUNT(*) FROM full_table) AS total_count, - (SELECT COUNT(*) FROM distinct_table) AS distinct_count + SELECT (SELECT COUNT(*) FROM full_table) AS `total_count`, + (SELECT COUNT(*) FROM distinct_table) AS `distinct_count` """ results, query_job = self._start_query(is_unique_sql) row = next(iter(results)) @@ -467,6 +477,7 @@ def read_gbq_table( total_count = row["total_count"] distinct_count = row["distinct_count"] is_total_ordering = total_count == distinct_count + ordering = core.ExpressionOrdering( ordering_value_columns=[ core.OrderingColumnReference(column_id) for column_id in index_cols @@ -477,7 +488,6 @@ def read_gbq_table( # We have a total ordering, so query via "time travel" so that # the underlying data doesn't mutate. if is_total_ordering: - # Get the timestamp from the job metadata rather than the query # text so that the query for determining uniqueness of the ID # columns can be cached. @@ -663,7 +673,8 @@ def _read_ibis( core.ArrayValue( self, table_expression, columns, hidden_ordering_columns, ordering ), - [index_col.get_name() for index_col in index_cols], + index_columns=[index_col.get_name() for index_col in index_cols], + column_labels=column_keys, index_labels=index_labels, ) @@ -887,6 +898,25 @@ def read_csv( ) return self.read_pandas(pandas_df) + def read_pickle( + self, + filepath_or_buffer: FilePath | ReadPickleBuffer, + compression: CompressionOptions = "infer", + storage_options: StorageOptions = None, + ): + pandas_obj = pandas.read_pickle( + filepath_or_buffer, + compression=compression, + storage_options=storage_options, + ) + + if isinstance(pandas_obj, pandas.Series): + if pandas_obj.name is None: + pandas_obj.name = "0" + bigframes_df = self.read_pandas(pandas_obj.to_frame()) + return bigframes_df[bigframes_df.columns[0]] + return self.read_pandas(pandas_obj) + def read_parquet( self, path: str | IO["bytes"], @@ -1086,6 +1116,10 @@ def read_gbq_function( Then it can be applied to a DataFrame or Series. + .. note:: + The return type of the function must be explicitly specified in the + function's original definition even if not otherwise required. + Args: function_name (str): the function's name in BigQuery in the format diff --git a/docs/index.rst b/docs/index.rst index ff1cd09eb7d..d239ea3a785 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -6,7 +6,6 @@ API reference .. toctree:: :maxdepth: 3 - user_guide/index reference/index Changelog diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 91abc59bc0a..891f15a51b4 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -128,3 +128,4 @@ name: preprocessing name: bigframes.ml name: BigQuery DataFrames + status: beta diff --git a/docs/user_guide/bigframes.pandas/remote_functions.rst b/docs/user_guide/bigframes.pandas/remote_functions.rst deleted file mode 100644 index 7540ba8a28c..00000000000 --- a/docs/user_guide/bigframes.pandas/remote_functions.rst +++ /dev/null @@ -1,134 +0,0 @@ - -Using the Remote Functions -========================== - -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 on it via `help` command. - -.. code-block:: python - - import bigframes.pandas as bpd - help(bpd.remote_function) - -Read a table and inspect the column of interest. - -.. code-block:: python - - 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 would be -using a pre-created connection named `bigframes-rf-conn`. 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. - -.. code-block:: python - - @bpd.remote_function([float], str, bigquery_connection='bigframes-rf-conn') - def get_bucket(num): - 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. - -.. code-block:: python - - 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. - -.. code-block:: python - - df[['body_mass_g', 'body_mass_bucket']].head(10) - -This 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. - -.. warning:: - The deployed cloud function may be visible to other users with sufficient - privilege in the project. The user should be careful about having any - sensitive data in the code that will be deployed as a remote function. - -The cloud function can be located from a property set in the remote function object. - -.. code-block:: python - - get_bucket.bigframes_cloud_function - -and then there is a BigQuery remote function created configured to call into the -cloud function via the BigQuery connection. That can also be located from -another property set in the remote function object. - -.. code-block:: python - - get_bucket.bigframes_remote_function - -The cloud assets created are persistant and the user can manage them directy -from the Google Cloud Console. - -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 decorated it as a remote function: - -.. code-block:: python - - @bpd.remote_function([str], str, bigquery_connection='bigframes-rf-conn') - def get_hash(input): - import hashlib - # handle missing value - if input is None: - input = "" - encoded_input = input.encode() - hash = hashlib.md5(encoded_input) - return hash.hexdigest() - -We can use this remote function in another `pandas`-like API `map` that can be -applied on a DataFrame: - -.. code-block:: python - - df_redacted = df[["species", "island", "sex"]].map(get_hash) - df_redacted.head(10). - -Using Existing Functions -======================== - -If you have already defined a custom function in BigQuery, either in the -BigQuery Google Cloud Console or with the `remote_function` decorator above or -otherwise, you may use it with BigQuery DataFrames with the `read_gbq_function` -method. - -More details are available via the `help` command: - -.. code-block:: python - - import bigframes.pandas as pd - help(pd.read_gbq_function) - -Here is an example of using `read_gbq_function` to load an existing function -named `get_bucket`: - -.. code-block:: python - - import bigframes.pandas as pd - - df = pd.read_gbq("bigquery-public-data.ml_datasets.penguins") - get_bucket = pd.read_gbq_function("get_bucket") - - df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket)) - df.head(10) - -Note: As mentioned above, if a function is created using the `remote_function` -decorator, its generated name (including project and dataset) is accessible -immediately afterward in the function's `bigframes_remote_function` attribute. -The same string can be passed to `read_gbq_function` later in another context. diff --git a/docs/user_guide/index.rst b/docs/user_guide/index.rst deleted file mode 100644 index 915e172159c..00000000000 --- a/docs/user_guide/index.rst +++ /dev/null @@ -1,9 +0,0 @@ -.. _user_guide: - -User Guide -========== - -.. toctree:: - :maxdepth: 2 - - bigframes.pandas/remote_functions diff --git a/notebooks/00 - Summary.ipynb b/notebooks/00 - Summary.ipynb deleted file mode 100644 index 66ac9a8de84..00000000000 --- a/notebooks/00 - Summary.ipynb +++ /dev/null @@ -1,2060 +0,0 @@ -{ - "cells": [ - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Using the BigQuery DataFrames API" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Set BigQuery DataFrames options" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "import bigframes.pandas\n", - "\n", - "bigframes.pandas.options.bigquery.project = \"bigframes-dev\"\n", - "bigframes.pandas.options.bigquery.location = \"us\"" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Initialize a dataframe for a BigQuery table" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "df = bigframes.pandas.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## View the DataFrame" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Adelie Penguin (Pygoscelis adeliae)Dream36.618.4184.03475.0FEMALE
1Adelie Penguin (Pygoscelis adeliae)Dream39.819.1184.04650.0MALE
2Adelie Penguin (Pygoscelis adeliae)Dream40.918.9184.03900.0MALE
3Chinstrap penguin (Pygoscelis antarctica)Dream46.517.9192.03500.0FEMALE
4Adelie Penguin (Pygoscelis adeliae)Dream37.316.8192.03000.0FEMALE
5Adelie Penguin (Pygoscelis adeliae)Dream43.218.5192.04100.0MALE
6Chinstrap penguin (Pygoscelis antarctica)Dream46.916.6192.02700.0FEMALE
7Chinstrap penguin (Pygoscelis antarctica)Dream50.518.4200.03400.0FEMALE
8Chinstrap penguin (Pygoscelis antarctica)Dream49.519.0200.03800.0MALE
9Adelie Penguin (Pygoscelis adeliae)Dream40.220.1200.03975.0MALE
10Adelie Penguin (Pygoscelis adeliae)Dream40.818.9208.04300.0MALE
11Adelie Penguin (Pygoscelis adeliae)Dream39.018.7185.03650.0MALE
12Adelie Penguin (Pygoscelis adeliae)Dream37.016.9185.03000.0FEMALE
13Chinstrap penguin (Pygoscelis antarctica)Dream47.017.3185.03700.0FEMALE
14Adelie Penguin (Pygoscelis adeliae)Dream34.017.1185.03400.0FEMALE
15Adelie Penguin (Pygoscelis adeliae)Dream37.016.5185.03400.0FEMALE
16Chinstrap penguin (Pygoscelis antarctica)Dream45.717.3193.03600.0FEMALE
17Chinstrap penguin (Pygoscelis antarctica)Dream50.619.4193.03800.0MALE
18Adelie Penguin (Pygoscelis adeliae)Dream39.717.9193.04250.0MALE
19Adelie Penguin (Pygoscelis adeliae)Dream37.818.1193.03750.0MALE
\n", - "
[344 rows x 7 columns in total]" - ], - "text/plain": [ - " species island culmen_length_mm \\\n", - "0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 \n", - "1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 \n", - "2 Adelie Penguin (Pygoscelis adeliae) Dream 40.9 \n", - "3 Chinstrap penguin (Pygoscelis antarctica) Dream 46.5 \n", - "4 Adelie Penguin (Pygoscelis adeliae) Dream 37.3 \n", - "5 Adelie Penguin (Pygoscelis adeliae) Dream 43.2 \n", - "6 Chinstrap penguin (Pygoscelis antarctica) Dream 46.9 \n", - "7 Chinstrap penguin (Pygoscelis antarctica) Dream 50.5 \n", - "8 Chinstrap penguin (Pygoscelis antarctica) Dream 49.5 \n", - "9 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", - "10 Adelie Penguin (Pygoscelis adeliae) Dream 40.8 \n", - "11 Adelie Penguin (Pygoscelis adeliae) Dream 39.0 \n", - "12 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", - "13 Chinstrap penguin (Pygoscelis antarctica) Dream 47.0 \n", - "14 Adelie Penguin (Pygoscelis adeliae) Dream 34.0 \n", - "15 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", - "16 Chinstrap penguin (Pygoscelis antarctica) Dream 45.7 \n", - "17 Chinstrap penguin (Pygoscelis antarctica) Dream 50.6 \n", - "18 Adelie Penguin (Pygoscelis adeliae) Dream 39.7 \n", - "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.8 \n", - "20 Chinstrap penguin (Pygoscelis antarctica) Dream 46.6 \n", - "21 Chinstrap penguin (Pygoscelis antarctica) Dream 51.3 \n", - "22 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", - "23 Adelie Penguin (Pygoscelis adeliae) Dream 36.8 \n", - "24 Chinstrap penguin (Pygoscelis antarctica) Dream 49.6 \n", - "\n", - " culmen_depth_mm flipper_length_mm body_mass_g sex \n", - "0 18.4 184.0 3475.0 FEMALE \n", - "1 19.1 184.0 4650.0 MALE \n", - "2 18.9 184.0 3900.0 MALE \n", - "3 17.9 192.0 3500.0 FEMALE \n", - "4 16.8 192.0 3000.0 FEMALE \n", - "5 18.5 192.0 4100.0 MALE \n", - "6 16.6 192.0 2700.0 FEMALE \n", - "7 18.4 200.0 3400.0 FEMALE \n", - "8 19.0 200.0 3800.0 MALE \n", - "9 20.1 200.0 3975.0 MALE \n", - "10 18.9 208.0 4300.0 MALE \n", - "11 18.7 185.0 3650.0 MALE \n", - "12 16.9 185.0 3000.0 FEMALE \n", - "13 17.3 185.0 3700.0 FEMALE \n", - "14 17.1 185.0 3400.0 FEMALE \n", - "15 16.5 185.0 3400.0 FEMALE \n", - "16 17.3 193.0 3600.0 FEMALE \n", - "17 19.4 193.0 3800.0 MALE \n", - "18 17.9 193.0 4250.0 MALE \n", - "19 18.1 193.0 3750.0 MALE \n", - "20 17.8 193.0 3800.0 FEMALE \n", - "21 19.2 193.0 3650.0 MALE \n", - "22 17.1 193.0 3400.0 FEMALE \n", - "23 18.5 193.0 3500.0 FEMALE \n", - "24 18.2 193.0 3775.0 MALE \n", - "...\n", - "\n", - "[344 rows x 7 columns]" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### View the column names in the dataframe (aka columns names in the table)" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "Index(['species', 'island', 'culmen_length_mm', 'culmen_depth_mm',\n", - " 'flipper_length_mm', 'body_mass_g', 'sex'],\n", - " dtype='object')" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df.columns" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### View the table schema" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "species string[pyarrow]\n", - "island string[pyarrow]\n", - "culmen_length_mm Float64\n", - "culmen_depth_mm Float64\n", - "flipper_length_mm Float64\n", - "body_mass_g Float64\n", - "sex string[pyarrow]\n", - "dtype: object" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df.dtypes" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Select a subset of columns" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
speciesislandbody_mass_g
0Adelie Penguin (Pygoscelis adeliae)Dream3475.0
1Adelie Penguin (Pygoscelis adeliae)Dream4650.0
2Adelie Penguin (Pygoscelis adeliae)Dream3900.0
3Chinstrap penguin (Pygoscelis antarctica)Dream3500.0
4Adelie Penguin (Pygoscelis adeliae)Dream3000.0
5Adelie Penguin (Pygoscelis adeliae)Dream4100.0
6Chinstrap penguin (Pygoscelis antarctica)Dream2700.0
7Chinstrap penguin (Pygoscelis antarctica)Dream3400.0
8Chinstrap penguin (Pygoscelis antarctica)Dream3800.0
9Adelie Penguin (Pygoscelis adeliae)Dream3975.0
10Adelie Penguin (Pygoscelis adeliae)Dream4300.0
11Adelie Penguin (Pygoscelis adeliae)Dream3650.0
12Adelie Penguin (Pygoscelis adeliae)Dream3000.0
13Chinstrap penguin (Pygoscelis antarctica)Dream3700.0
14Adelie Penguin (Pygoscelis adeliae)Dream3400.0
15Adelie Penguin (Pygoscelis adeliae)Dream3400.0
16Chinstrap penguin (Pygoscelis antarctica)Dream3600.0
17Chinstrap penguin (Pygoscelis antarctica)Dream3800.0
18Adelie Penguin (Pygoscelis adeliae)Dream4250.0
19Adelie Penguin (Pygoscelis adeliae)Dream3750.0
\n", - "
[344 rows x 3 columns in total]" - ], - "text/plain": [ - " species island body_mass_g\n", - "0 Adelie Penguin (Pygoscelis adeliae) Dream 3475.0\n", - "1 Adelie Penguin (Pygoscelis adeliae) Dream 4650.0\n", - "2 Adelie Penguin (Pygoscelis adeliae) Dream 3900.0\n", - "3 Chinstrap penguin (Pygoscelis antarctica) Dream 3500.0\n", - "4 Adelie Penguin (Pygoscelis adeliae) Dream 3000.0\n", - "5 Adelie Penguin (Pygoscelis adeliae) Dream 4100.0\n", - "6 Chinstrap penguin (Pygoscelis antarctica) Dream 2700.0\n", - "7 Chinstrap penguin (Pygoscelis antarctica) Dream 3400.0\n", - "8 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", - "9 Adelie Penguin (Pygoscelis adeliae) Dream 3975.0\n", - "10 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", - "11 Adelie Penguin (Pygoscelis adeliae) Dream 3650.0\n", - "12 Adelie Penguin (Pygoscelis adeliae) Dream 3000.0\n", - "13 Chinstrap penguin (Pygoscelis antarctica) Dream 3700.0\n", - "14 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", - "15 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", - "16 Chinstrap penguin (Pygoscelis antarctica) Dream 3600.0\n", - "17 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", - "18 Adelie Penguin (Pygoscelis adeliae) Dream 4250.0\n", - "19 Adelie Penguin (Pygoscelis adeliae) Dream 3750.0\n", - "20 Chinstrap penguin (Pygoscelis antarctica) Dream 3800.0\n", - "21 Chinstrap penguin (Pygoscelis antarctica) Dream 3650.0\n", - "22 Adelie Penguin (Pygoscelis adeliae) Dream 3400.0\n", - "23 Adelie Penguin (Pygoscelis adeliae) Dream 3500.0\n", - "24 Chinstrap penguin (Pygoscelis antarctica) Dream 3775.0\n", - "...\n", - "\n", - "[344 rows x 3 columns]" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = df[[\n", - " \"species\",\n", - " \"island\",\n", - " \"body_mass_g\",\n", - "]]\n", - "df" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### View the first ten values of a series" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "0 3475.0\n", - "1 4650.0\n", - "2 3900.0\n", - "3 3500.0\n", - "4 3000.0\n", - "5 4100.0\n", - "6 2700.0\n", - "7 3400.0\n", - "8 3800.0\n", - "9 3975.0\n", - "Name: body_mass_g, dtype: Float64" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df['body_mass_g'].head(10)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Compute the mean of a series" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "4201.7543859649095" - ] - }, - "execution_count": 8, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df['body_mass_g'].mean()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Filter the DataFrame" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
speciesislandbody_mass_g
1Adelie Penguin (Pygoscelis adeliae)Dream4650.0
5Adelie Penguin (Pygoscelis adeliae)Dream4100.0
10Adelie Penguin (Pygoscelis adeliae)Dream4300.0
18Adelie Penguin (Pygoscelis adeliae)Dream4250.0
25Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
26Adelie Penguin (Pygoscelis adeliae)Dream4000.0
27Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
28Chinstrap penguin (Pygoscelis antarctica)Dream4300.0
30Chinstrap penguin (Pygoscelis antarctica)Dream4450.0
36Adelie Penguin (Pygoscelis adeliae)Dream4450.0
44Chinstrap penguin (Pygoscelis antarctica)Dream4100.0
45Chinstrap penguin (Pygoscelis antarctica)Dream4800.0
57Chinstrap penguin (Pygoscelis antarctica)Dream4400.0
61Chinstrap penguin (Pygoscelis antarctica)Dream4150.0
66Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
67Chinstrap penguin (Pygoscelis antarctica)Dream4050.0
68Chinstrap penguin (Pygoscelis antarctica)Dream4100.0
74Adelie Penguin (Pygoscelis adeliae)Dream4350.0
77Adelie Penguin (Pygoscelis adeliae)Dream4150.0
78Adelie Penguin (Pygoscelis adeliae)Dream4400.0
\n", - "
[177 rows x 3 columns in total]" - ], - "text/plain": [ - " species island body_mass_g\n", - "1 Adelie Penguin (Pygoscelis adeliae) Dream 4650.0\n", - "5 Adelie Penguin (Pygoscelis adeliae) Dream 4100.0\n", - "10 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", - "18 Adelie Penguin (Pygoscelis adeliae) Dream 4250.0\n", - "25 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", - "26 Adelie Penguin (Pygoscelis adeliae) Dream 4000.0\n", - "27 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", - "28 Chinstrap penguin (Pygoscelis antarctica) Dream 4300.0\n", - "30 Chinstrap penguin (Pygoscelis antarctica) Dream 4450.0\n", - "36 Adelie Penguin (Pygoscelis adeliae) Dream 4450.0\n", - "44 Chinstrap penguin (Pygoscelis antarctica) Dream 4100.0\n", - "45 Chinstrap penguin (Pygoscelis antarctica) Dream 4800.0\n", - "57 Chinstrap penguin (Pygoscelis antarctica) Dream 4400.0\n", - "61 Chinstrap penguin (Pygoscelis antarctica) Dream 4150.0\n", - "66 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", - "67 Chinstrap penguin (Pygoscelis antarctica) Dream 4050.0\n", - "68 Chinstrap penguin (Pygoscelis antarctica) Dream 4100.0\n", - "74 Adelie Penguin (Pygoscelis adeliae) Dream 4350.0\n", - "77 Adelie Penguin (Pygoscelis adeliae) Dream 4150.0\n", - "78 Adelie Penguin (Pygoscelis adeliae) Dream 4400.0\n", - "80 Chinstrap penguin (Pygoscelis antarctica) Dream 4300.0\n", - "90 Chinstrap penguin (Pygoscelis antarctica) Dream 4150.0\n", - "92 Chinstrap penguin (Pygoscelis antarctica) Dream 4500.0\n", - "93 Adelie Penguin (Pygoscelis adeliae) Dream 4300.0\n", - "94 Chinstrap penguin (Pygoscelis antarctica) Dream 4550.0\n", - "...\n", - "\n", - "[177 rows x 3 columns]" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df[df['body_mass_g'] >= 4000.0]" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Using the Remote Functions" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### BigQuery DataFrames gives you the ability to turn your custom scalar functions into a BigQuery remote function.\n", - "\n", - "It requires the GCP project to be set up appropriately and the user having sufficient privileges to use them. One can find more details on it via `help` command." - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Help on function remote_function in module bigframes.pandas:\n", - "\n", - "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True)\n", - " Decorator to turn a user defined function into a BigQuery remote function.\n", - " \n", - " Args:\n", - " input_types (list(type)):\n", - " List of input data types in the user defined function.\n", - " output_type (type):\n", - " Data type of the output in the user defined function.\n", - " dataset (str, Optional):\n", - " Dataset to use to create a BigQuery function. It should be in\n", - " `.` or `` format. If this\n", - " param is not provided then session dataset id would be used.\n", - " bigquery_connection (str, Optional):\n", - " Name of the BigQuery connection. If it is pre created in the same\n", - " location as the `bigquery_client.location` then it would be used,\n", - " otherwise it would be created dynamically assuming the user has\n", - " necessary priviliges. If this param is not provided then the\n", - " bigquery connection from the session would be used.\n", - " reuse (bool, Optional):\n", - " Reuse the remote function if already exists.\n", - " `True` by default, which will result in reusing an existing remote\n", - " function (if any) that was previously created for the same udf.\n", - " Setting it to false would force creating a unique remote function.\n", - " If the required remote function does not exist then it would be\n", - " created irrespective of this param.\n", - " \n", - " Notes:\n", - " Please make sure following is setup before using this API:\n", - " \n", - " 1. Have the below APIs enabled for your project:\n", - " a. BigQuery Connection API\n", - " b. Cloud Functions API\n", - " c. Cloud Run API\n", - " d. Cloud Build API\n", - " e. Artifact Registry API\n", - " f. Cloud Resource Manager API\n", - " \n", - " This can be done from the cloud console (change PROJECT_ID to yours):\n", - " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", - " Or from the gcloud CLI:\n", - " $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com\n", - " \n", - " 2. Have following IAM roles enabled for you:\n", - " a. BigQuery Data Editor (roles/bigquery.dataEditor)\n", - " b. BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", - " c. Cloud Functions Developer (roles/cloudfunctions.developer)\n", - " d. Service Account User (roles/iam.serviceAccountUser)\n", - " e. Storage Object Viewer (roles/storage.objectViewer)\n", - " f. Project IAM Admin (roles/resourcemanager.projectIamAdmin)\n", - " (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", - " \n", - " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", - " a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", - " b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", - " Alternatively, the IAM could also be setup via the gcloud CLI:\n", - " $ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"\n", - "\n" - ] - } - ], - "source": [ - "import bigframes.pandas as pd\n", - "help(pd.remote_function)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Define a custom function, and specify the intent to turn it into a remote function.\n", - "\n", - "It requires a BigQuery connection. If the connection is not already created,\n", - "the BigQuery DataFrames package attempts to create one assuming the necessary\n", - "APIs and IAM permissions are setup in the project." - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "[INFO][2023-06-28 23:31:49,355][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-b3fab64f5997ad6a516379defe8d4202 --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmp9w5e89lh --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", - "Preparing function...\n", - ".done.\n", - "Deploying function...\n", - "[Build]..........................................................................................................................................................................................................................................................................................................................................................................................................................done\n", - "[Service].........................................................................................................................................................................................................done\n", - "Done.\n", - "You can view your function in the Cloud Console here: https://console.cloud.google.com/functions/details/us-central1/bigframes-b3fab64f5997ad6a516379defe8d4202?project=bigframes-dev\n", - "\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "buildConfig:\n", - " build: projects/1084210331973/locations/us-central1/builds/780b1780-9b38-4515-ae60-89d05454ef83\n", - " entryPoint: udf_http\n", - " runtime: python310\n", - " source:\n", - " storageSource:\n", - " bucket: gcf-v2-sources-1084210331973-us-central1\n", - " object: bigframes-b3fab64f5997ad6a516379defe8d4202/function-source.zip\n", - " sourceProvenance:\n", - " resolvedStorageSource:\n", - " bucket: gcf-v2-sources-1084210331973-us-central1\n", - " generation: '1687995112300727'\n", - " object: bigframes-b3fab64f5997ad6a516379defe8d4202/function-source.zip\n", - "environment: GEN_2\n", - "labels:\n", - " deployment-tool: cli-gcloud\n", - "name: projects/bigframes-dev/locations/us-central1/functions/bigframes-b3fab64f5997ad6a516379defe8d4202\n", - "serviceConfig:\n", - " allTrafficOnLatestRevision: true\n", - " availableCpu: '0.1666'\n", - " availableMemory: 256M\n", - " ingressSettings: ALLOW_ALL\n", - " maxInstanceCount: 100\n", - " maxInstanceRequestConcurrency: 1\n", - " revision: bigframes-b3fab64f5997ad6a516379defe8d4202-00001-tut\n", - " service: projects/bigframes-dev/locations/us-central1/services/bigframes-b3fab64f5997ad6a516379defe8d4202\n", - " serviceAccountEmail: 1084210331973-compute@developer.gserviceaccount.com\n", - " timeoutSeconds: 60\n", - " uri: https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app\n", - "state: ACTIVE\n", - "updateTime: '2023-06-28T23:32:51.911131997Z'\n", - "url: https://us-central1-bigframes-dev.cloudfunctions.net/bigframes-b3fab64f5997ad6a516379defe8d4202\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "[INFO][2023-06-28 23:32:55,330][bigframes.remote_function] Successfully created cloud function bigframes-b3fab64f5997ad6a516379defe8d4202 with uri (https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app)\n", - "[INFO][2023-06-28 23:32:59,378][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", - "[INFO][2023-06-28 23:32:59,379][bigframes.remote_function] Creating BQ remote function: \n", - " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_b3fab64f5997ad6a516379defe8d4202(num FLOAT64)\n", - " RETURNS STRING\n", - " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", - " OPTIONS (\n", - " endpoint = \"https://bigframes-b3fab64f5997ad6a516379defe8d4202-7krlje3eoq-uc.a.run.app\"\n", - " )\n", - "[INFO][2023-06-28 23:33:00,338][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_b3fab64f5997ad6a516379defe8d4202\n" - ] - } - ], - "source": [ - "@pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", - "def get_bucket(num):\n", - " if not num: return \"NA\"\n", - " boundary = 4000\n", - " return \"at_or_above_4000\" if num >= boundary else \"below_4000\"" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Run the custom function on the BigQuery-backed dataframe" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
body_mass_gbody_mass_bucket
03475.0below_4000
14650.0at_or_above_4000
23900.0below_4000
33500.0below_4000
43000.0below_4000
54100.0at_or_above_4000
62700.0below_4000
73400.0below_4000
83800.0below_4000
93975.0below_4000
\n", - "
[10 rows x 2 columns in total]" - ], - "text/plain": [ - " body_mass_g body_mass_bucket\n", - "0 3475.0 below_4000\n", - "1 4650.0 at_or_above_4000\n", - "2 3900.0 below_4000\n", - "3 3500.0 below_4000\n", - "4 3000.0 below_4000\n", - "5 4100.0 at_or_above_4000\n", - "6 2700.0 below_4000\n", - "7 3400.0 below_4000\n", - "8 3800.0 below_4000\n", - "9 3975.0 below_4000\n", - "\n", - "[10 rows x 2 columns]" - ] - }, - "execution_count": 12, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = df.assign(body_mass_bucket=df['body_mass_g'].apply(get_bucket))\n", - "df[['body_mass_g', 'body_mass_bucket']].head(10)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Using the ML API" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Initialize a DataFrame from a BigQuery table" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
speciesislandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Adelie Penguin (Pygoscelis adeliae)Dream36.618.4184.03475.0FEMALE
1Adelie Penguin (Pygoscelis adeliae)Dream39.819.1184.04650.0MALE
2Adelie Penguin (Pygoscelis adeliae)Dream40.918.9184.03900.0MALE
3Chinstrap penguin (Pygoscelis antarctica)Dream46.517.9192.03500.0FEMALE
4Adelie Penguin (Pygoscelis adeliae)Dream37.316.8192.03000.0FEMALE
5Adelie Penguin (Pygoscelis adeliae)Dream43.218.5192.04100.0MALE
6Chinstrap penguin (Pygoscelis antarctica)Dream46.916.6192.02700.0FEMALE
7Chinstrap penguin (Pygoscelis antarctica)Dream50.518.4200.03400.0FEMALE
8Chinstrap penguin (Pygoscelis antarctica)Dream49.519.0200.03800.0MALE
9Adelie Penguin (Pygoscelis adeliae)Dream40.220.1200.03975.0MALE
10Adelie Penguin (Pygoscelis adeliae)Dream40.818.9208.04300.0MALE
11Adelie Penguin (Pygoscelis adeliae)Dream39.018.7185.03650.0MALE
12Adelie Penguin (Pygoscelis adeliae)Dream37.016.9185.03000.0FEMALE
13Chinstrap penguin (Pygoscelis antarctica)Dream47.017.3185.03700.0FEMALE
14Adelie Penguin (Pygoscelis adeliae)Dream34.017.1185.03400.0FEMALE
15Adelie Penguin (Pygoscelis adeliae)Dream37.016.5185.03400.0FEMALE
16Chinstrap penguin (Pygoscelis antarctica)Dream45.717.3193.03600.0FEMALE
17Chinstrap penguin (Pygoscelis antarctica)Dream50.619.4193.03800.0MALE
18Adelie Penguin (Pygoscelis adeliae)Dream39.717.9193.04250.0MALE
19Adelie Penguin (Pygoscelis adeliae)Dream37.818.1193.03750.0MALE
\n", - "
[344 rows x 7 columns in total]" - ], - "text/plain": [ - " species island culmen_length_mm \\\n", - "0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 \n", - "1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 \n", - "2 Adelie Penguin (Pygoscelis adeliae) Dream 40.9 \n", - "3 Chinstrap penguin (Pygoscelis antarctica) Dream 46.5 \n", - "4 Adelie Penguin (Pygoscelis adeliae) Dream 37.3 \n", - "5 Adelie Penguin (Pygoscelis adeliae) Dream 43.2 \n", - "6 Chinstrap penguin (Pygoscelis antarctica) Dream 46.9 \n", - "7 Chinstrap penguin (Pygoscelis antarctica) Dream 50.5 \n", - "8 Chinstrap penguin (Pygoscelis antarctica) Dream 49.5 \n", - "9 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", - "10 Adelie Penguin (Pygoscelis adeliae) Dream 40.8 \n", - "11 Adelie Penguin (Pygoscelis adeliae) Dream 39.0 \n", - "12 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", - "13 Chinstrap penguin (Pygoscelis antarctica) Dream 47.0 \n", - "14 Adelie Penguin (Pygoscelis adeliae) Dream 34.0 \n", - "15 Adelie Penguin (Pygoscelis adeliae) Dream 37.0 \n", - "16 Chinstrap penguin (Pygoscelis antarctica) Dream 45.7 \n", - "17 Chinstrap penguin (Pygoscelis antarctica) Dream 50.6 \n", - "18 Adelie Penguin (Pygoscelis adeliae) Dream 39.7 \n", - "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.8 \n", - "20 Chinstrap penguin (Pygoscelis antarctica) Dream 46.6 \n", - "21 Chinstrap penguin (Pygoscelis antarctica) Dream 51.3 \n", - "22 Adelie Penguin (Pygoscelis adeliae) Dream 40.2 \n", - "23 Adelie Penguin (Pygoscelis adeliae) Dream 36.8 \n", - "24 Chinstrap penguin (Pygoscelis antarctica) Dream 49.6 \n", - "\n", - " culmen_depth_mm flipper_length_mm body_mass_g sex \n", - "0 18.4 184.0 3475.0 FEMALE \n", - "1 19.1 184.0 4650.0 MALE \n", - "2 18.9 184.0 3900.0 MALE \n", - "3 17.9 192.0 3500.0 FEMALE \n", - "4 16.8 192.0 3000.0 FEMALE \n", - "5 18.5 192.0 4100.0 MALE \n", - "6 16.6 192.0 2700.0 FEMALE \n", - "7 18.4 200.0 3400.0 FEMALE \n", - "8 19.0 200.0 3800.0 MALE \n", - "9 20.1 200.0 3975.0 MALE \n", - "10 18.9 208.0 4300.0 MALE \n", - "11 18.7 185.0 3650.0 MALE \n", - "12 16.9 185.0 3000.0 FEMALE \n", - "13 17.3 185.0 3700.0 FEMALE \n", - "14 17.1 185.0 3400.0 FEMALE \n", - "15 16.5 185.0 3400.0 FEMALE \n", - "16 17.3 193.0 3600.0 FEMALE \n", - "17 19.4 193.0 3800.0 MALE \n", - "18 17.9 193.0 4250.0 MALE \n", - "19 18.1 193.0 3750.0 MALE \n", - "20 17.8 193.0 3800.0 FEMALE \n", - "21 19.2 193.0 3650.0 MALE \n", - "22 17.1 193.0 3400.0 FEMALE \n", - "23 18.5 193.0 3500.0 FEMALE \n", - "24 18.2 193.0 3775.0 MALE \n", - "...\n", - "\n", - "[344 rows x 7 columns]" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = bigframes.pandas.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")\n", - "df" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Clean and prepare the data" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
islandculmen_length_mmculmen_depth_mmflipper_length_mmbody_mass_gsex
0Dream36.618.4184.03475.0FEMALE
1Dream39.819.1184.04650.0MALE
2Dream40.918.9184.03900.0MALE
4Dream37.316.8192.03000.0FEMALE
5Dream43.218.5192.04100.0MALE
9Dream40.220.1200.03975.0MALE
10Dream40.818.9208.04300.0MALE
11Dream39.018.7185.03650.0MALE
12Dream37.016.9185.03000.0FEMALE
14Dream34.017.1185.03400.0FEMALE
15Dream37.016.5185.03400.0FEMALE
18Dream39.717.9193.04250.0MALE
19Dream37.818.1193.03750.0MALE
22Dream40.217.1193.03400.0FEMALE
23Dream36.818.5193.03500.0FEMALE
26Dream41.518.5201.04000.0MALE
31Dream33.116.1178.02900.0FEMALE
32Dream37.218.1178.03900.0MALE
33Dream39.516.7178.03250.0FEMALE
35Dream36.018.5186.03100.0FEMALE
\n", - "
[146 rows x 6 columns in total]" - ], - "text/plain": [ - " island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g \\\n", - "0 Dream 36.6 18.4 184.0 3475.0 \n", - "1 Dream 39.8 19.1 184.0 4650.0 \n", - "2 Dream 40.9 18.9 184.0 3900.0 \n", - "4 Dream 37.3 16.8 192.0 3000.0 \n", - "5 Dream 43.2 18.5 192.0 4100.0 \n", - "9 Dream 40.2 20.1 200.0 3975.0 \n", - "10 Dream 40.8 18.9 208.0 4300.0 \n", - "11 Dream 39.0 18.7 185.0 3650.0 \n", - "12 Dream 37.0 16.9 185.0 3000.0 \n", - "14 Dream 34.0 17.1 185.0 3400.0 \n", - "15 Dream 37.0 16.5 185.0 3400.0 \n", - "18 Dream 39.7 17.9 193.0 4250.0 \n", - "19 Dream 37.8 18.1 193.0 3750.0 \n", - "22 Dream 40.2 17.1 193.0 3400.0 \n", - "23 Dream 36.8 18.5 193.0 3500.0 \n", - "26 Dream 41.5 18.5 201.0 4000.0 \n", - "31 Dream 33.1 16.1 178.0 2900.0 \n", - "32 Dream 37.2 18.1 178.0 3900.0 \n", - "33 Dream 39.5 16.7 178.0 3250.0 \n", - "35 Dream 36.0 18.5 186.0 3100.0 \n", - "36 Dream 39.6 18.1 186.0 4450.0 \n", - "38 Dream 41.3 20.3 194.0 3550.0 \n", - "41 Dream 35.7 18.0 202.0 3550.0 \n", - "51 Dream 38.1 17.6 187.0 3425.0 \n", - "53 Dream 36.0 17.1 187.0 3700.0 \n", - "\n", - " sex \n", - "0 FEMALE \n", - "1 MALE \n", - "2 MALE \n", - "4 FEMALE \n", - "5 MALE \n", - "9 MALE \n", - "10 MALE \n", - "11 MALE \n", - "12 FEMALE \n", - "14 FEMALE \n", - "15 FEMALE \n", - "18 MALE \n", - "19 MALE \n", - "22 FEMALE \n", - "23 FEMALE \n", - "26 MALE \n", - "31 FEMALE \n", - "32 MALE \n", - "33 FEMALE \n", - "35 FEMALE \n", - "36 MALE \n", - "38 MALE \n", - "41 FEMALE \n", - "51 FEMALE \n", - "53 FEMALE \n", - "...\n", - "\n", - "[146 rows x 6 columns]" - ] - }, - "execution_count": 14, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# filter down to the data we want to analyze\n", - "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", - "\n", - "# drop the columns we don't care about\n", - "adelie_data = adelie_data.drop(columns=[\"species\"])\n", - "\n", - "# drop rows with nulls to get our training data\n", - "training_data = adelie_data.dropna()\n", - "\n", - "# take a peek at the training data\n", - "training_data" - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "metadata": {}, - "outputs": [], - "source": [ - "# pick feature columns and label column\n", - "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", - "label_columns = training_data[['body_mass_g']]\n", - "\n", - "# also get the rows that we want to make predictions for (i.e. where the feature column is null)\n", - "missing_body_mass = adelie_data[adelie_data.body_mass_g.isnull()]" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Train and evaluate a linear regression model using the ML API" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
mean_absolute_errormean_squared_errormean_squared_log_errormedian_absolute_errorr2_scoreexplained_variance
0223.87876378553.6016340.005614181.3309110.6239510.623951
\n", - "
[1 rows x 6 columns in total]" - ], - "text/plain": [ - " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", - "0 223.878763 78553.601634 0.005614 \n", - "\n", - " median_absolute_error r2_score explained_variance \n", - "0 181.330911 0.623951 0.623951 \n", - "\n", - "[1 rows x 6 columns]" - ] - }, - "execution_count": 16, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from bigframes.ml.linear_model import LinearRegression\n", - "\n", - "# as in scikit-learn, a newly created model is just a bundle of parameters\n", - "# default parameters are fine here\n", - "model = LinearRegression()\n", - "\n", - "# this will train a temporary model in BigQuery Machine Learning\n", - "model.fit(feature_columns, label_columns)\n", - "\n", - "# check how the model performed\n", - "model.score(feature_columns, label_columns)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Make predictions using the model" - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
predicted_body_mass_g
2923603.735118
\n", - "
[1 rows x 1 columns in total]" - ], - "text/plain": [ - " predicted_body_mass_g\n", - "292 3603.735118\n", - "\n", - "[1 rows x 1 columns]" - ] - }, - "execution_count": 17, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "model.predict(missing_body_mass)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Save the trained model to BigQuery, so we can load it later" - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "LinearRegression()" - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "model.to_gbq(\"bqml_tutorial.penguins_model\", replace=True)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.9" - }, - "orig_nbformat": 4 - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/notebooks/01 - Getting Started.ipynb b/notebooks/01 - Getting Started.ipynb deleted file mode 100644 index 473bdd8ceae..00000000000 --- a/notebooks/01 - Getting Started.ipynb +++ /dev/null @@ -1,1190 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "d7a03de2-c0ef-4f80-9cd5-f96e87cf2d54", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "# On the instance where you are running jupyter,\n", - "# authenticate with gcloud first:\n", - "#\n", - "# gcloud auth application-default login\n", - "\n", - "import bigframes.pandas as bpd\n", - "\n", - "# Change this location to the location of your datasets.\n", - "# We use \"us\" as that is the location of the sample data.\n", - "bpd.options.bigquery.location = \"us\"" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "941cb6c3-8c54-42ce-a945-4fa604176b2e", - "metadata": { - "tags": [] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "5c480b60490940d3a45fa6b9ca2cecdb", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job 411d90c8-8b22-40b0-ad42-04f9e38c074e is DONE. 0 Bytes processed. \n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
taxi_idtrip_end_timestamptrip_secondstrip_milespickup_census_tractdropoff_census_tractpickup_community_areadropoff_community_areafaretips...extrastrip_totalpayment_typecompanypickup_latitudepickup_longitudepickup_locationdropoff_latitudedropoff_longitudedropoff_location
trip_start_timestampunique_key
2023-07-01 00:00:00+00:00fb004dfe98302ccf34e5e11bf1081568a94843a275cf3a53aae5e5858361a7ca64f75d3407dc0a44d7bc42...2023-07-01 00:00:00+00:005621.82<NA><NA>22229.52.53...0.012.03MobileCity Service41.922761-87.699155POINT (-87.6991553432 41.9227606205)41.922761-87.699155POINT (-87.6991553432 41.9227606205)
\n", - "

1 rows × 21 columns

\n", - "[1 rows x 21 columns in total]" - ], - "text/plain": [ - " taxi_id \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 75cf3a53aae5e5858361a7ca64f75d3407dc0a44d7bc42... \n", - "\n", - " trip_end_timestamp \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 2023-07-01 00:00:00+00:00 \n", - "\n", - " trip_seconds \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 562 \n", - "\n", - " trip_miles \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 1.82 \n", - "\n", - " pickup_census_tract \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 \n", - "\n", - " dropoff_census_tract \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 \n", - "\n", - " pickup_community_area \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 22 \n", - "\n", - " dropoff_community_area \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 22 \n", - "\n", - " fare \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 9.5 \n", - "\n", - " tips ... \\\n", - "trip_start_timestamp unique_key ... \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 2.53 ... \n", - "\n", - " extras \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 0.0 \n", - "\n", - " trip_total \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 12.03 \n", - "\n", - " payment_type \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 Mobile \n", - "\n", - " company \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 City Service \n", - "\n", - " pickup_latitude \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 41.922761 \n", - "\n", - " pickup_longitude \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 -87.699155 \n", - "\n", - " pickup_location \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 POINT (-87.6991553432 41.9227606205) \n", - "\n", - " dropoff_latitude \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 41.922761 \n", - "\n", - " dropoff_longitude \\\n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 -87.699155 \n", - "\n", - " dropoff_location \n", - "trip_start_timestamp unique_key \n", - "2023-07-01 00:00:00+00:00 fb004dfe98302ccf34e5e11bf1081568a94843a2 POINT (-87.6991553432 41.9227606205) \n", - "\n", - "[1 rows x 21 columns]" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df.tail(n=1)" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "f1eee1c6-7214-440e-bc17-6839d53a6718", - "metadata": {}, - "outputs": [], - "source": [ - "df = df[[\n", - " \"company\",\n", - " \"trip_miles\",\n", - " \"fare\",\n", - " \"tips\",\n", - "]]" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "66071984-e371-4161-8d7e-00d810b5cfab", - "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "0ce3a9e3bd90477dac718b8e60fd3458", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job b105f8f6-fe3d-4f90-85c7-2cbede41413b is DONE. 0 Bytes processed. \n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
companytrip_milesfaretips
trip_start_timestampunique_key
2013-01-01 00:00:00+00:00006d53094ff5b3c81140c00aa91fdb5467db0802<NA>2.37.850.0
00a723cf08cf29c56faa8ef77be519f6590b3405<NA>1.05.650.0
01396f0debefe45ccce23e48dd471edd67bd32e8Dispatch Taxi Affiliation0.08.850.0
01e9a03fd793670ed35ef7195eeb99775895611fDispatch Taxi Affiliation1.87.450.0
024f0f183c43ad2be33f232fd8dca0a86a3e4925<NA>4.914.050.0
04373ef74229c2122ff90d92642fcd94c88302e4Dispatch Taxi Affiliation0.06.650.0
0501a22fba176594d6f974360e20737af1b6cb66<NA>2.27.650.0
05faff951f2a2703fe3eb24efac026c2a1669a92<NA>3.39.650.0
0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbeBlue Ribbon Taxi Association Inc.0.06.250.0
06b4dc6dcd6295e69284e1c9d5a98015b0588ea6Choice Taxi Association0.013.850.0
07483a0baa6df0af68696b6f2483b7f429ff1827Northwest Management LLC0.64.650.0
07585d085f65cd512b94de09684253fe0ffb31a2Taxi Affiliation Services7.04.650.0
090ea9dc5dc26374ee06cdab1b809895beb1befcDispatch Taxi Affiliation0.010.050.0
091eda4ec41520d6ba96003f05eb10744f73aa82Northwest Management LLC0.015.250.0
097f2efed248cb5900792a89bba96a64918345ceBlue Ribbon Taxi Association Inc.0.012.250.0
09ffa3f40d36dbf8d51ca903738d9bba92557ed7Taxi Affiliation Services0.06.450.0
0a3446fb4f72d6485548883cd7f6cfe068d44829Taxi Affiliation Services0.014.250.0
0a7ea707a8700555e898489200a66b546a834170<NA>5.0413.852.0
0be545d9934dd343a2d929fca2e53d3d2851137bDispatch Taxi Affiliation1.15.650.0
0c2b3705c9748a5949bf5b7d727b2d1f2e18805f<NA>12.326.250.0
0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4Taxi Affiliation Services0.015.250.0
0dafcce34426b5377fd3925418a4c24b3e45c7d9<NA>1.26.250.0
0f5f7162b289dd9743b9d2e7fc75bf69696bb7bcTaxi Affiliation Services0.09.850.0
0f63007e437d589c9139acb026f03a832da3240fNorthwest Management LLC1.56.650.0
10508745a1069e9cdaad8760a6a243dc924cc0f2<NA>5.513.250.0
\n", - "

25 rows × 4 columns

\n", - "[208395018 rows x 4 columns in total]" - ], - "text/plain": [ - " company \\\n", - "trip_start_timestamp unique_key \n", - "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 \n", - " 00a723cf08cf29c56faa8ef77be519f6590b3405 \n", - " 01396f0debefe45ccce23e48dd471edd67bd32e8 Dispatch Taxi Affiliation \n", - " 01e9a03fd793670ed35ef7195eeb99775895611f Dispatch Taxi Affiliation \n", - " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 \n", - " 04373ef74229c2122ff90d92642fcd94c88302e4 Dispatch Taxi Affiliation \n", - " 0501a22fba176594d6f974360e20737af1b6cb66 \n", - " 05faff951f2a2703fe3eb24efac026c2a1669a92 \n", - " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe Blue Ribbon Taxi Association Inc. \n", - " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 Choice Taxi Association \n", - " 07483a0baa6df0af68696b6f2483b7f429ff1827 Northwest Management LLC \n", - " 07585d085f65cd512b94de09684253fe0ffb31a2 Taxi Affiliation Services \n", - " 090ea9dc5dc26374ee06cdab1b809895beb1befc Dispatch Taxi Affiliation \n", - " 091eda4ec41520d6ba96003f05eb10744f73aa82 Northwest Management LLC \n", - " 097f2efed248cb5900792a89bba96a64918345ce Blue Ribbon Taxi Association Inc. \n", - " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 Taxi Affiliation Services \n", - " 0a3446fb4f72d6485548883cd7f6cfe068d44829 Taxi Affiliation Services \n", - " 0a7ea707a8700555e898489200a66b546a834170 \n", - " 0be545d9934dd343a2d929fca2e53d3d2851137b Dispatch Taxi Affiliation \n", - " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f \n", - " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 Taxi Affiliation Services \n", - " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 \n", - " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc Taxi Affiliation Services \n", - " 0f63007e437d589c9139acb026f03a832da3240f Northwest Management LLC \n", - " 10508745a1069e9cdaad8760a6a243dc924cc0f2 \n", - "\n", - " trip_miles \\\n", - "trip_start_timestamp unique_key \n", - "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 2.3 \n", - " 00a723cf08cf29c56faa8ef77be519f6590b3405 1.0 \n", - " 01396f0debefe45ccce23e48dd471edd67bd32e8 0.0 \n", - " 01e9a03fd793670ed35ef7195eeb99775895611f 1.8 \n", - " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 4.9 \n", - " 04373ef74229c2122ff90d92642fcd94c88302e4 0.0 \n", - " 0501a22fba176594d6f974360e20737af1b6cb66 2.2 \n", - " 05faff951f2a2703fe3eb24efac026c2a1669a92 3.3 \n", - " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 0.0 \n", - " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 0.0 \n", - " 07483a0baa6df0af68696b6f2483b7f429ff1827 0.6 \n", - " 07585d085f65cd512b94de09684253fe0ffb31a2 7.0 \n", - " 090ea9dc5dc26374ee06cdab1b809895beb1befc 0.0 \n", - " 091eda4ec41520d6ba96003f05eb10744f73aa82 0.0 \n", - " 097f2efed248cb5900792a89bba96a64918345ce 0.0 \n", - " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 0.0 \n", - " 0a3446fb4f72d6485548883cd7f6cfe068d44829 0.0 \n", - " 0a7ea707a8700555e898489200a66b546a834170 5.04 \n", - " 0be545d9934dd343a2d929fca2e53d3d2851137b 1.1 \n", - " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 12.3 \n", - " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 0.0 \n", - " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 1.2 \n", - " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 0.0 \n", - " 0f63007e437d589c9139acb026f03a832da3240f 1.5 \n", - " 10508745a1069e9cdaad8760a6a243dc924cc0f2 5.5 \n", - "\n", - " fare \\\n", - "trip_start_timestamp unique_key \n", - "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 7.85 \n", - " 00a723cf08cf29c56faa8ef77be519f6590b3405 5.65 \n", - " 01396f0debefe45ccce23e48dd471edd67bd32e8 8.85 \n", - " 01e9a03fd793670ed35ef7195eeb99775895611f 7.45 \n", - " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 14.05 \n", - " 04373ef74229c2122ff90d92642fcd94c88302e4 6.65 \n", - " 0501a22fba176594d6f974360e20737af1b6cb66 7.65 \n", - " 05faff951f2a2703fe3eb24efac026c2a1669a92 9.65 \n", - " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 6.25 \n", - " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 13.85 \n", - " 07483a0baa6df0af68696b6f2483b7f429ff1827 4.65 \n", - " 07585d085f65cd512b94de09684253fe0ffb31a2 4.65 \n", - " 090ea9dc5dc26374ee06cdab1b809895beb1befc 10.05 \n", - " 091eda4ec41520d6ba96003f05eb10744f73aa82 15.25 \n", - " 097f2efed248cb5900792a89bba96a64918345ce 12.25 \n", - " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 6.45 \n", - " 0a3446fb4f72d6485548883cd7f6cfe068d44829 14.25 \n", - " 0a7ea707a8700555e898489200a66b546a834170 13.85 \n", - " 0be545d9934dd343a2d929fca2e53d3d2851137b 5.65 \n", - " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 26.25 \n", - " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 15.25 \n", - " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 6.25 \n", - " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 9.85 \n", - " 0f63007e437d589c9139acb026f03a832da3240f 6.65 \n", - " 10508745a1069e9cdaad8760a6a243dc924cc0f2 13.25 \n", - "\n", - " tips \n", - "trip_start_timestamp unique_key \n", - "2013-01-01 00:00:00+00:00 006d53094ff5b3c81140c00aa91fdb5467db0802 0.0 \n", - " 00a723cf08cf29c56faa8ef77be519f6590b3405 0.0 \n", - " 01396f0debefe45ccce23e48dd471edd67bd32e8 0.0 \n", - " 01e9a03fd793670ed35ef7195eeb99775895611f 0.0 \n", - " 024f0f183c43ad2be33f232fd8dca0a86a3e4925 0.0 \n", - " 04373ef74229c2122ff90d92642fcd94c88302e4 0.0 \n", - " 0501a22fba176594d6f974360e20737af1b6cb66 0.0 \n", - " 05faff951f2a2703fe3eb24efac026c2a1669a92 0.0 \n", - " 0651a5e1eb4792a53395c3b3b4fd9f7f8383cfbe 0.0 \n", - " 06b4dc6dcd6295e69284e1c9d5a98015b0588ea6 0.0 \n", - " 07483a0baa6df0af68696b6f2483b7f429ff1827 0.0 \n", - " 07585d085f65cd512b94de09684253fe0ffb31a2 0.0 \n", - " 090ea9dc5dc26374ee06cdab1b809895beb1befc 0.0 \n", - " 091eda4ec41520d6ba96003f05eb10744f73aa82 0.0 \n", - " 097f2efed248cb5900792a89bba96a64918345ce 0.0 \n", - " 09ffa3f40d36dbf8d51ca903738d9bba92557ed7 0.0 \n", - " 0a3446fb4f72d6485548883cd7f6cfe068d44829 0.0 \n", - " 0a7ea707a8700555e898489200a66b546a834170 2.0 \n", - " 0be545d9934dd343a2d929fca2e53d3d2851137b 0.0 \n", - " 0c2b3705c9748a5949bf5b7d727b2d1f2e18805f 0.0 \n", - " 0d37bbfeb207fbb2353e1911fdf878f3c24ed4c4 0.0 \n", - " 0dafcce34426b5377fd3925418a4c24b3e45c7d9 0.0 \n", - " 0f5f7162b289dd9743b9d2e7fc75bf69696bb7bc 0.0 \n", - " 0f63007e437d589c9139acb026f03a832da3240f 0.0 \n", - " 10508745a1069e9cdaad8760a6a243dc924cc0f2 0.0 \n", - "...\n", - "\n", - "[208395018 rows x 4 columns]" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "e51b687e-0282-459c-8fd0-0ca22cfb153d", - "metadata": { - "tags": [] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "bc4d2a7cee5e4fc1a9513ec2fb54d293", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job 494b3016-0582-45e0-9fcf-bb38eb00722e is RUNNING. \n", + " \n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\n", + " \n", + " \n", + "" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "24743cf4a1e1" + }, + "source": [ + "**_NOTE_**: This notebook has been tested in the following environment:\n", + "\n", + "* Python version = 3.9" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "The goal of this notebook is to demonstrate an enterprise generative AI use case. A marketing user can provide information about a new pharmaceutical drug and its generic name, and receive ideas on marketing-oriented brand names for that drug.\n", + "\n", + "Learn more about [BigQuery DataFrames](https://cloud.google.com/bigquery/docs/dataframes-quickstart)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn about Generative AI concepts such as prompting and few-shot learning, as well as how to use BigFrames ML for performing these tasks simply using an intuitive dataframe API.\n", + "\n", + "The steps performed include:\n", + "\n", + "1. Ask the user for the generic name and usage for the drug.\n", + "1. Use `bigframes` to query the FDA dataset of over 100,000 drugs, filtered on the brand name, generic name, and indications & usage columns.\n", + "1. Filter this dataset to find prototypical brand names that can be used as examples in prompt tuning.\n", + "1. Create a prompt with the user input, general instructions, examples and counter-examples for the desired brand name.\n", + "1. Use the `bigframes.ml.llm.PaLM2TextGenerator` to generate choices of brand names." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This notebook uses the [FDA dataset](https://cloud.google.com/blog/topics/healthcare-life-sciences/fda-mystudies-comes-to-google-cloud) available at [`bigquery-public-data.fda_drug`](https://console.cloud.google.com/bigquery?ws=!1m4!1m3!3m2!1sbigquery-public-data!2sfda_drug)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery (compute)\n", + "* BigQuery ML\n", + "\n", + "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models),\n", + "and [BigQuery ML pricing](https://cloud.google.com/bigquery/pricing#bqml),\n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages required to execute this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "!pip install -U --quiet bigframes" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# # Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bpd\n", + "from google.cloud import bigquery_connection_v1 as bq_connection\n", + "from bigframes.ml.llm import PaLM2TextGenerator\n", + "from IPython.display import Markdown" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**1. Vertex AI Workbench**\n", + "* Do nothing as you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**2. Local JupyterLab instance, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**3. Colab, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Enable the BigQuery API](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com).\n", + "\n", + "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "evsJaAj5te0X" + }, + "source": [ + "#### BigFrames configuration\n", + "\n", + "Next, we will specify a [BigQuery connection](https://cloud.google.com/bigquery/docs/working-with-connections). If you already have a connection, you can simplify provide the name and skip the following creation steps.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "G1vVsPiMsL2X" + }, + "outputs": [], + "source": [ + "# Please fill in these values.\n", + "LOCATION = \"us\" # @param {type:\"string\"}\n", + "CONNECTION = \"\" # @param {type:\"string\"}\n", + "\n", + "connection_name = f\"{PROJECT_ID}.{LOCATION}.{CONNECTION}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WGS_TzhWlPBN" + }, + "source": [ + "We will now try to use the provided connection, and if it doesn't exist, create a new one. We will also print the service account used." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "56Hw42m6kFrj" + }, + "outputs": [], + "source": [ + "# Initialize client and set request parameters\n", + "client = bq_connection.ConnectionServiceClient()\n", + "new_conn_parent = f\"projects/{PROJECT_ID}/locations/{LOCATION}\"\n", + "exists_conn_parent = f\"projects/{PROJECT_ID}/locations/{LOCATION}/connections/{CONNECTION}\"\n", + "cloud_resource_properties = bq_connection.CloudResourceProperties({})\n", + "\n", + "# Try to connect using provided connection\n", + "try:\n", + " request = client.get_connection(\n", + " request=bq_connection.GetConnectionRequest(name=exists_conn_parent)\n", + " )\n", + " CONN_SERVICE_ACCOUNT = f\"serviceAccount:{request.cloud_resource.service_account_id}\"\n", + "# Create a new connection on error\n", + "except Exception:\n", + " connection = bq_connection.types.Connection(\n", + " {\"friendly_name\": CONNECTION, \"cloud_resource\": cloud_resource_properties}\n", + " )\n", + " request = bq_connection.CreateConnectionRequest(\n", + " {\n", + " \"parent\": new_conn_parent,\n", + " \"connection_id\": CONNECTION,\n", + " \"connection\": connection,\n", + " }\n", + " )\n", + " response = client.create_connection(request)\n", + " CONN_SERVICE_ACCOUNT = (\n", + " f\"serviceAccount:{response.cloud_resource.service_account_id}\"\n", + " )\n", + "# Set service account permissions\n", + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/bigquery.connectionUser'\n", + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/aiplatform.user'\n", + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/run.invoker'\n", + "\n", + "print(CONN_SERVICE_ACCOUNT)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "### Initialize BigFrames client\n", + "\n", + "Here, we set the project configuration based on the provided parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OCccLirpkSRz" + }, + "outputs": [], + "source": [ + "bpd.options.bigquery.project = PROJECT_ID\n", + "bpd.options.bigquery.location = LOCATION" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "m8UCEtX9uLn6" + }, + "source": [ + "## Generate a name\n", + "\n", + "Let's start with entering a generic name and description of the drug." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oxphj2gnuKou" + }, + "outputs": [], + "source": [ + "GENERIC_NAME = \"Entropofloxacin\" # @param {type:\"string\"}\n", + "USAGE = \"Entropofloxacin is a fluoroquinolone antibiotic that is used to treat a variety of bacterial infections, including: pneumonia, streptococcus infections, salmonella infections, escherichia coli infections, and pseudomonas aeruginosa infections It is taken by mouth or by injection. The dosage and frequency of administration will vary depending on the type of infection being treated. It should be taken for the full course of treatment, even if symptoms improve after a few days. Stopping the medication early may increase the risk of the infection coming back.\" # @param {type:\"string\"}\n", + "NUM_NAMES = 10 # @param {type:\"integer\"}\n", + "TEMPERATURE = 0.5 # @param {type: \"number\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1q-vlbalzu1Q" + }, + "source": [ + "We can now create a prompt string, and populate it with the name and description." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "0knz5ZWMzed-" + }, + "outputs": [], + "source": [ + "zero_shot_prompt = f\"\"\"Provide {NUM_NAMES} unique and modern brand names in Markdown bullet point format. Do not provide any additional explanation.\n", + "\n", + "Be creative with the brand names. Don't use English words directly; use variants or invented words.\n", + "\n", + "The generic name is: {GENERIC_NAME}\n", + "\n", + "The indications and usage are: {USAGE}.\"\"\"\n", + "\n", + "print(zero_shot_prompt)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "LCRE2L720f5y" + }, + "source": [ + "Next, let's create a helper function to predict with our model. It will take a string input, and add it to a temporary BigFrames `DataFrame`. It will also return the string extracted from the response `DataFrame`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "LB3xgDroIxlx" + }, + "outputs": [], + "source": [ + "def predict(prompt: str, temperature: float = TEMPERATURE) -> str:\n", + " # Create dataframe\n", + " input = bpd.DataFrame(\n", + " {\n", + " \"prompt\": [prompt],\n", + " }\n", + " )\n", + "\n", + " # Return response\n", + " return model.predict(input, temperature).ml_generate_text_llm_result.iloc[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b1ZapNZsJW2p" + }, + "source": [ + "We can now initialize the model, and get a response to our prompt!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "UW2fQ2k5Hsic" + }, + "outputs": [], + "source": [ + "# Get BigFrames session\n", + "session = bpd.get_global_session()\n", + "\n", + "# Define the model\n", + "model = PaLM2TextGenerator(session=session, connection_name=connection_name)\n", + "\n", + "# Invoke LLM with prompt\n", + "response = predict(zero_shot_prompt)\n", + "\n", + "# Print results as Markdown\n", + "Markdown(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "o3yIhHV2jsUT" + }, + "source": [ + "We're off to a great start! Let's see if we can refine our response." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mBroUzWS8xOL" + }, + "source": [ + "## Few-shot learning\n", + "\n", + "Let's try using [few-shot learning](https://paperswithcode.com/task/few-shot-learning). We will provide a few examples of what we're looking for along with our prompt.\n", + "\n", + "Our prompt will consist of 3 parts:\n", + "* General instructions (e.g. generate $n$ brand names)\n", + "* Multiple examples\n", + "* Information about the drug we'd like to generate a name for\n", + "\n", + "Let's walk through how to construct this prompt.\n", + "\n", + "Our first step will be to define how many examples we want to provide in the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MXdI78SOElyt" + }, + "outputs": [], + "source": [ + "# Specify number of examples to include\n", + "\n", + "NUM_EXAMPLES = 3 # @param {type:\"integer\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "U8w4puVM_892" + }, + "source": [ + "Next, let's define a prefix that will set the overall context." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "aQ2iscnhF2cx" + }, + "outputs": [], + "source": [ + "prefix_prompt = f\"\"\"Provide {NUM_NAMES} unique and modern brand names in Markdown bullet point format, related to the drug at the bottom of this prompt.\n", + "\n", + "Be creative with the brand names. Don't use English words directly; use variants or invented words.\n", + "\n", + "First, we will provide {NUM_EXAMPLES} examples to help with your thought process.\n", + "\n", + "Then, we will provide the generic name and usage for the drug we'd like you to generate brand names for.\n", + "\"\"\"\n", + "\n", + "print(prefix_prompt)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "VI0Spv-axN7d" + }, + "source": [ + "Our next step will be to include examples into the prompt.\n", + "\n", + "We will start out by retrieving the raw data for the examples, by querying the BigQuery public dataset." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "IoO_Bp8wA07N" + }, + "outputs": [], + "source": [ + "# Query 3 columns of interest from drug label dataset\n", + "df = bpd.read_gbq(\"bigquery-public-data.fda_drug.drug_label\",\n", + " col_order=[\"openfda_generic_name\", \"openfda_brand_name\", \"indications_and_usage\"])\n", + "\n", + "# Exclude any rows with missing data\n", + "df = df.dropna()\n", + "\n", + "# Drop duplicate rows\n", + "df = df.drop_duplicates()\n", + "\n", + "# Print values\n", + "df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W5kOtbNGBTI2" + }, + "source": [ + "Let's now filter the results to remove atypical names." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "95WDe2eCCeLx" + }, + "outputs": [], + "source": [ + "# Remove names with spaces\n", + "df = df[df[\"openfda_brand_name\"].str.find(\" \") == -1]\n", + "\n", + "# Remove names with 5 or fewer characters\n", + "df = df[df[\"openfda_brand_name\"].str.len() > 5]\n", + "\n", + "# Remove names where the generic and brand name match (case-insensitive)\n", + "df = df[df[\"openfda_generic_name\"].str.lower() != df[\"openfda_brand_name\"].str.lower()]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "FZD89ep4EyYc" + }, + "source": [ + "Let's take `NUM_EXAMPLES` samples to include in the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2ohZYg7QEyJV" + }, + "outputs": [], + "source": [ + "# Take a sample and convert to a Pandas dataframe for local usage.\n", + "df_examples = df.sample(NUM_EXAMPLES, random_state=3).to_pandas()\n", + "\n", + "df_examples" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "J-Qa1_SCImXy" + }, + "source": [ + "Let's now convert the data to a JSON structure, to enable embedding into a prompt. For consistency, we'll capitalize each example brand name." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PcJdSaw0EGcW" + }, + "outputs": [], + "source": [ + "examples = [\n", + " {\n", + " \"brand_name\": brand_name.capitalize(),\n", + " \"generic_name\": generic_name,\n", + " \"usage\": usage,\n", + " }\n", + " for brand_name, generic_name, usage in zip(\n", + " df_examples[\"openfda_brand_name\"],\n", + " df_examples[\"openfda_generic_name\"],\n", + " df_examples[\"indications_and_usage\"],\n", + " )\n", + "]\n", + "\n", + "print(examples)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oU4mb1Dwgq64" + }, + "source": [ + "We'll create a prompt template for each example, and view the first one." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "kzAVsF6wJ93S" + }, + "outputs": [], + "source": [ + "example_prompt = \"\"\n", + "for example in examples:\n", + " example_prompt += f\"Generic name: {example['generic_name']}\\nUsage: {example['usage']}\\nBrand name: {example['brand_name']}\\n\\n\"\n", + "\n", + "example_prompt" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "kbV2X1CXAyLV" + }, + "source": [ + "Finally, we can create a suffix to our prompt. This will contain the generic name of the drug, its usage, ending with a request for brand names." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OYp6W_XfHTlo" + }, + "outputs": [], + "source": [ + "suffix_prompt = f\"\"\"Generic name: {GENERIC_NAME}\n", + "Usage: {USAGE}\n", + "Brand names:\"\"\"\n", + "\n", + "print(suffix_prompt)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "RiaisW1nihJP" + }, + "source": [ + "Let's pull it altogether into a few shot prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "99xdU7l8C1h8" + }, + "outputs": [], + "source": [ + "# Define the prompt\n", + "few_shot_prompt = prefix_prompt + example_prompt + suffix_prompt\n", + "\n", + "# Print the prompt\n", + "print(few_shot_prompt)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "nbUWdHtfitWn" + }, + "source": [ + "Now, let's pass our prompt to the LLM, and get a response!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d4ODRJdvLhlQ" + }, + "outputs": [], + "source": [ + "response = predict(few_shot_prompt)\n", + "\n", + "Markdown(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "pFakjrTElOBs" + }, + "source": [ + "# Bulk generation\n", + "\n", + "Let's take these experiments to the next level by generating many names in bulk. We'll see how to leverage BigFrames at scale!\n", + "\n", + "We can start by finding drugs that are missing brand names. There are approximately 4,000 drugs that meet this criteria. We'll put a limit of 100 in this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8eAutS41mx6U" + }, + "outputs": [], + "source": [ + "# Query 3 columns of interest from drug label dataset\n", + "df_missing = bpd.read_gbq(\"bigquery-public-data.fda_drug.drug_label\",\n", + " col_order=[\"openfda_generic_name\", \"openfda_brand_name\", \"indications_and_usage\"])\n", + "\n", + "# Exclude any rows with missing data\n", + "df_missing = df_missing.dropna()\n", + "\n", + "# Include rows in which openfda_brand_name equals openfda_generic_name\n", + "df_missing = df_missing[df_missing[\"openfda_generic_name\"] == df_missing[\"openfda_brand_name\"]]\n", + "\n", + "# Limit the number of rows for demonstration purposes\n", + "df_missing = df_missing.head(100)\n", + "\n", + "# Print values\n", + "df_missing.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Fm6L8S7eVnCI" + }, + "source": [ + "We will create a column `prompt` with a customized prompt for each row." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "19TvGN1PVmVX" + }, + "outputs": [], + "source": [ + "df_missing[\"prompt\"] = (\n", + " \"Provide a unique and modern brand name related to this pharmaceutical drug.\"\n", + " + \"Don't use English words directly; use variants or invented words. The generic name is: \"\n", + " + df_missing[\"openfda_generic_name\"]\n", + " + \". The indications and usage are: \"\n", + " + df_missing[\"indications_and_usage\"]\n", + " + \".\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "njxwBvCKgMPE" + }, + "source": [ + "We'll create a new helper method, `batch_predict()` and query the LLM. The job may take a couple minutes to execute." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tiSHa5B4aFhw" + }, + "outputs": [], + "source": [ + "def batch_predict(\n", + " input: bpd.DataFrame, temperature: float = TEMPERATURE\n", + ") -> bpd.DataFrame:\n", + " return model.predict(input, temperature).ml_generate_text_llm_result\n", + "\n", + "\n", + "response = batch_predict(df_missing[\"prompt\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "K5a2nHdLgZEj" + }, + "source": [ + "Let's check the results for one of our responses!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "TnizdeqBdbZj" + }, + "outputs": [], + "source": [ + "# Pick a sample\n", + "k = 0\n", + "\n", + "# Gather the prompt and response details\n", + "prompt_generic = df_missing[\"openfda_generic_name\"][k].iloc[0]\n", + "prompt_usage = df_missing[\"indications_and_usage\"][k].iloc[0]\n", + "response_str = response[k].iloc[0]\n", + "\n", + "# Print details\n", + "print(f\"Generic name: {prompt_generic}\")\n", + "print(f\"Brand name: {prompt_usage}\")\n", + "print(f\"Response: {response_str}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W4MviwyMI-Qh" + }, + "source": [ + "Congratulations! You have learned how to use generative AI to jumpstart the creative process.\n", + "\n", + "You've also seen how BigFrames can manage each step of the process, including gathering data, data manipulation, and querying the LLM." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Bys6--dVmq7R" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cIODjOLump_-" + }, + "outputs": [], + "source": [ + "# Delete the BigQuery Connection\n", + "from google.cloud import bigquery_connection_v1 as bq_connection\n", + "client = bq_connection.ConnectionServiceClient()\n", + "CONNECTION_ID = f\"projects/{PROJECT_ID}/locations/{LOCATION}/connections/{CONNECTION}\"\n", + "client.delete_connection(name=CONNECTION_ID)\n", + "print(f\"Deleted connection {CONNECTION_ID}.\")" + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/06 - Using ML - Large Language Models.ipynb b/notebooks/generative_ai/large_language_models.ipynb similarity index 100% rename from notebooks/06 - Using ML - Large Language Models.ipynb rename to notebooks/generative_ai/large_language_models.ipynb diff --git a/notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb b/notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb new file mode 100644 index 00000000000..39e2ef535cc --- /dev/null +++ b/notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb @@ -0,0 +1,891 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2022 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "## Use BigQuery DataFrames with Generative AI for code generation\n", + "\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "24743cf4a1e1" + }, + "source": [ + "**_NOTE_**: This notebook has been tested in the following environment:\n", + "\n", + "* Python version = 3.10" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "Use this notebook to walk through an example use case of generating sample code by using BigQuery DataFrames and its integration with Generative AI support on Vertex AI.\n", + "\n", + "Learn more about [BigQuery DataFrames](https://cloud.google.com/python/docs/reference/bigframes/latest)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you create a CSV file containing sample code for calling a given set of APIs.\n", + "\n", + "The steps include:\n", + "\n", + "- Defining an LLM model in BigQuery DataFrames, specifically the [`text-bison` model of the PaLM API](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text), using `bigframes.ml.llm`.\n", + "- Creating a DataFrame by reading in data from Cloud Storage.\n", + "- Manipulating data in the DataFrame to build LLM prompts.\n", + "- Sending DataFrame prompts to the LLM model using the `predict` method.\n", + "- Creating and using a custom function to transform the output provided by the LLM model response.\n", + "- Exporting the resulting transformed DataFrame as a CSV file." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses a dataset listing the names of various pandas DataFrame and Series APIs." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery\n", + "* Generative AI support on Vertex AI\n", + "* Cloud Functions\n", + "\n", + "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models),\n", + "[Generative AI support on Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing#generative_ai_models), and [Cloud Functions pricing](https://cloud.google.com/functions/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages, which are required to run this notebook:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "!pip install bigframes --upgrade --quiet" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "Complete the tasks in this section to set up your environment." + ] + }, + { + "cell_type": "markdown", + "source": [ + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Click here](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com,bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com,cloudresourcemanager.googleapis.com) to enable the following APIs:\n", + "\n", + " * BigQuery API\n", + " * BigQuery Connection API\n", + " * Cloud Functions API\n", + " * Cloud Run API\n", + " * Artifact Registry API\n", + " * Cloud Build API\n", + " * Cloud Resource Manager API\n", + " * Vertex AI API\n", + "\n", + "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)." + ], + "metadata": { + "id": "Wbr2aVtFQBcg" + } + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "If you don't know your project ID, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Set the region\n", + "\n", + "You can also change the `REGION` variable used by BigQuery. Learn more about [BigQuery regions](https://cloud.google.com/bigquery/docs/locations#supported_locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "eF-Twtc4XGem" + }, + "outputs": [], + "source": [ + "REGION = \"US\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you might have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**Vertex AI Workbench**\n", + "\n", + "Do nothing, you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**Local JupyterLab instance**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**Colab**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "from google.cloud import bigquery\n", + "from google.cloud import bigquery_connection_v1 as bq_connection" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "### Set BigQuery DataFrames options" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NPPMuw2PXGeo" + }, + "outputs": [], + "source": [ + "bf.options.bigquery.project = PROJECT_ID\n", + "bf.options.bigquery.location = REGION" + ] + }, + { + "cell_type": "markdown", + "source": [ + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + ], + "metadata": { + "id": "DTVtFlqeFbrU" + } + }, + { + "cell_type": "markdown", + "source": [ + "# Define the LLM model\n", + "\n", + "BigQuery DataFrames provides integration with [`text-bison` model of the PaLM API](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text) via Vertex AI.\n", + "\n", + "This section walks through a few steps required in order to use the model in your notebook." + ], + "metadata": { + "id": "6eytf4xQHzcF" + } + }, + { + "cell_type": "markdown", + "source": [ + "## Create a BigQuery Cloud resource connection\n", + "\n", + "You need to create a [Cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) to enable BigQuery DataFrames to interact with Vertex AI services." + ], + "metadata": { + "id": "rS4VO1TGiO4G" + } + }, + { + "cell_type": "code", + "source": [ + "CONN_NAME = \"bqdf-llm\"\n", + "\n", + "client = bq_connection.ConnectionServiceClient()\n", + "new_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n", + "exists_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}\"\n", + "cloud_resource_properties = bq_connection.CloudResourceProperties({})\n", + "\n", + "try:\n", + " request = client.get_connection(\n", + " request=bq_connection.GetConnectionRequest(name=exists_conn_parent)\n", + " )\n", + " CONN_SERVICE_ACCOUNT = f\"serviceAccount:{request.cloud_resource.service_account_id}\"\n", + "except Exception:\n", + " connection = bq_connection.types.Connection(\n", + " {\"friendly_name\": CONN_NAME, \"cloud_resource\": cloud_resource_properties}\n", + " )\n", + " request = bq_connection.CreateConnectionRequest(\n", + " {\n", + " \"parent\": new_conn_parent,\n", + " \"connection_id\": CONN_NAME,\n", + " \"connection\": connection,\n", + " }\n", + " )\n", + " response = client.create_connection(request)\n", + " CONN_SERVICE_ACCOUNT = (\n", + " f\"serviceAccount:{response.cloud_resource.service_account_id}\"\n", + " )\n", + "print(CONN_SERVICE_ACCOUNT)" + ], + "metadata": { + "id": "KFPjDM4LVh96" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Set permissions for the service account\n", + "\n", + "The resource connection service account requires certain project-level permissions:\n", + " - `roles/aiplatform.user` and `roles/bigquery.connectionUser`: These roles are required for the connection to create a model definition using the LLM model in Vertex AI ([documentation](https://cloud.google.com/bigquery/docs/generate-text#give_the_service_account_access)).\n", + " - `roles/run.invoker`: This role is required for the connection to have read-only access to Cloud Run services that back custom/remote functions ([documentation](https://cloud.google.com/bigquery/docs/remote-functions#grant_permission_on_function)).\n", + "\n", + "Set these permissions by running the following `gcloud` commands:" + ], + "metadata": { + "id": "W6l6Ol2biU9h" + } + }, + { + "cell_type": "code", + "source": [ + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/bigquery.connectionUser'\n", + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/aiplatform.user'\n", + "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/run.invoker'" + ], + "metadata": { + "id": "d8wja24SVq6s" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Define the model\n", + "\n", + "Use `bigframes.ml.llm` to define the model:" + ], + "metadata": { + "id": "qUjT8nw-jIXp" + } + }, + { + "cell_type": "code", + "source": [ + "from bigframes.ml.llm import PaLM2TextGenerator\n", + "\n", + "session = bf.get_global_session()\n", + "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", + "model = PaLM2TextGenerator(session=session, connection_name=connection)" + ], + "metadata": { + "id": "sdjeXFwcHfl7" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Read data from Cloud Storage into BigQuery DataFrames\n", + "\n", + "You can create a BigQuery DataFrames DataFrame by reading data from any of the following locations:\n", + "\n", + "* A local data file\n", + "* Data stored in a BigQuery table\n", + "* A data file stored in Cloud Storage\n", + "* An in-memory pandas DataFrame\n", + "\n", + "In this tutorial, you create BigQuery DataFrames DataFrames by reading two CSV files stored in Cloud Storage, one containing a list of DataFrame API names and one containing a list of Series API names." + ], + "metadata": { + "id": "GbW0oCnU1s1N" + } + }, + { + "cell_type": "code", + "source": [ + "df_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/df.csv\")\n", + "series_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/series.csv\")" + ], + "metadata": { + "id": "SchiTkQGIJog" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Take a peek at a few rows of data for each file:" + ], + "metadata": { + "id": "7OBjw2nmQY3-" + } + }, + { + "cell_type": "code", + "source": [ + "df_api.head(2)" + ], + "metadata": { + "id": "QCqgVCIsGGuv" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "series_api.head(2)" + ], + "metadata": { + "id": "BGJnZbgEGS5-" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Generate code using the LLM model\n", + "\n", + "Prepare the prompts and send them to the LLM model for prediction." + ], + "metadata": { + "id": "m3ZJEsi7SUKV" + } + }, + { + "cell_type": "markdown", + "source": [ + "## Prompt design in BigQuery DataFrames\n", + "\n", + "Designing prompts for LLMs is a fast growing area and you can read more in [this documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/introduction-prompt-design).\n", + "\n", + "For this tutorial, you use a simple prompt to ask the LLM model for sample code for each of the API methods (or rows) from the last step's DataFrames. The output is the new DataFrames `df_prompt` and `series_prompt`, which contain the full prompt text." + ], + "metadata": { + "id": "9EMAqR37AfLS" + } + }, + { + "cell_type": "code", + "source": [ + "df_prompt_prefix = \"Generate Pandas sample code for DataFrame.\"\n", + "series_prompt_prefix = \"Generate Pandas sample code for Series.\"\n", + "\n", + "df_prompt = (df_prompt_prefix + df_api['API'])\n", + "series_prompt = (series_prompt_prefix + series_api['API'])\n", + "\n", + "df_prompt.head(2)" + ], + "metadata": { + "id": "EDAaIwHpQCDZ" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Make predictions using the LLM model\n", + "\n", + "Use the BigQuery DataFrames DataFrame containing the full prompt text as the input to the `predict` method. The `predict` method calls the LLM model and returns its generated text output back to two new BigQuery DataFrames DataFrames, `df_pred` and `series_pred`.\n", + "\n", + "Note: The predictions might take a few minutes to run." + ], + "metadata": { + "id": "rwPLjqW2Ajzh" + } + }, + { + "cell_type": "code", + "source": [ + "df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024)\n", + "series_pred = model.predict(series_prompt.to_frame(), max_output_tokens=1024)" + ], + "metadata": { + "id": "6i6HkFJZa8na" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Once the predictions are processed, take a look at the sample output from the LLM, which provides code samples for the API names listed in the DataFrames dataset." + ], + "metadata": { + "id": "89cB8MW4UIdV" + } + }, + { + "cell_type": "code", + "source": [ + "print(df_pred['ml_generate_text_llm_result'].iloc[0])" + ], + "metadata": { + "id": "9A2gw6hP_2nX" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Manipulate LLM output using a remote function\n", + "\n", + "The output that the LLM provides often contains additional text beyond the code sample itself. Using BigQuery DataFrames, you can deploy custom Python functions that process and transform this output.\n", + "\n" + ], + "metadata": { + "id": "Fx4lsNqMorJ-" + } + }, + { + "cell_type": "markdown", + "source": [ + "Running the cell below creates a custom function that you can use to process the LLM output data in two ways:\n", + "1. Strip the LLM text output to include only the code block.\n", + "2. Substitute `import pandas as pd` with `import bigframes.pandas as bf` so that the resulting code block works with BigQuery DataFrames." + ], + "metadata": { + "id": "d8L7SN03VByG" + } + }, + { + "cell_type": "code", + "source": [ + "@bf.remote_function([str], str, bigquery_connection=CONN_NAME)\n", + "def extract_code(text: str):\n", + " try:\n", + " res = text[text.find('\\n')+1:text.find('```', 3)]\n", + " res = res.replace(\"import pandas as pd\", \"import bigframes.pandas as bf\")\n", + " if \"import bigframes.pandas as bf\" not in res:\n", + " res = \"import bigframes.pandas as bf\\n\" + res\n", + " return res\n", + " except:\n", + " return \"\"" + ], + "metadata": { + "id": "GskyyUQPowBT" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "The custom function is deployed as a Cloud Function, and then integrated with BigQuery as a [remote function](https://cloud.google.com/bigquery/docs/remote-functions). Save both of the function names so that you can clean them up at the end of this notebook." + ], + "metadata": { + "id": "hVQAoqBUOJQf" + } + }, + { + "cell_type": "code", + "source": [ + "CLOUD_FUNCTION_NAME = format(extract_code.bigframes_cloud_function)\n", + "print(\"Cloud Function Name \" + CLOUD_FUNCTION_NAME)\n", + "REMOTE_FUNCTION_NAME = format(extract_code.bigframes_remote_function)\n", + "print(\"Remote Function Name \" + REMOTE_FUNCTION_NAME)" + ], + "metadata": { + "id": "PBlp-C-DOHRO" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Apply the custom function to each LLM output DataFrame to get the processed results:" + ], + "metadata": { + "id": "4FEucaiqVs3H" + } + }, + { + "cell_type": "code", + "source": [ + "df_code = df_pred.assign(code=df_pred['ml_generate_text_llm_result'].apply(extract_code))\n", + "series_code = series_pred.assign(code=series_pred['ml_generate_text_llm_result'].apply(extract_code))" + ], + "metadata": { + "id": "bsQ9cmoWo0Ps" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "You can see the differences by inspecting the first row of data:" + ], + "metadata": { + "id": "ujQVVuhfWA3y" + } + }, + { + "cell_type": "code", + "source": [ + "print(df_code['code'].iloc[0])" + ], + "metadata": { + "id": "7yWzjhGy_zcy" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Save the results to Cloud Storage\n", + "\n", + "BigQuery DataFrames lets you save a BigQuery DataFrames DataFrame as a CSV file in Cloud Storage for further use. Try that now with your processed LLM output data." + ], + "metadata": { + "id": "GTRdUw-Ro5R1" + } + }, + { + "cell_type": "markdown", + "source": [ + "Create a new Cloud Storage bucket with a unique name:" + ], + "metadata": { + "id": "9DQ7eiQxPTi3" + } + }, + { + "cell_type": "code", + "source": [ + "import uuid\n", + "BUCKET_ID = \"code-samples-\" + str(uuid.uuid1())\n", + "\n", + "!gsutil mb gs://{BUCKET_ID}" + ], + "metadata": { + "id": "-J5LHgS6LLZ0" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Use `to_csv` to write each BigQuery DataFrames DataFrame as a CSV file in the Cloud Storage bucket:" + ], + "metadata": { + "id": "tyxZXj0UPYUv" + } + }, + { + "cell_type": "code", + "source": [ + "df_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/df_code*.csv\")\n", + "series_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/series_code*.csv\")" + ], + "metadata": { + "id": "Zs_b5L-4IvER" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "You can navigate to the Cloud Storage bucket browser to download the two files and view them.\n", + "\n", + "Run the following cell, and then follow the link to your Cloud Storage bucket browser:" + ], + "metadata": { + "id": "UDBtDlrTuuh8" + } + }, + { + "cell_type": "code", + "source": [ + "print(f'https://console.developers.google.com/storage/browser/{BUCKET_ID}/')" + ], + "metadata": { + "id": "PspCXu-qu_ND" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Summary and next steps\n", + "\n", + "You've used BigQuery DataFrames' integration with LLM models (`bigframes.ml.llm`) to generate code samples, and have tranformed LLM output by creating and using a custom function in BigQuery DataFrames.\n", + "\n", + "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks)." + ], + "metadata": { + "id": "RGSvUk48RK20" + } + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "source": [ + "# # Delete the BigQuery Connection\n", + "# from google.cloud import bigquery_connection_v1 as bq_connection\n", + "# client = bq_connection.ConnectionServiceClient()\n", + "# CONNECTION_ID = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}\"\n", + "# client.delete_connection(name=CONNECTION_ID)\n", + "# print(f\"Deleted connection '{CONNECTION_ID}'.\")" + ], + "metadata": { + "id": "yw7A461XLjvW" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "# # Delete the Cloud Function\n", + "# ! gcloud functions delete {CLOUD_FUNCTION_NAME} --quiet\n", + "# # Delete the Remote Function\n", + "# REMOTE_FUNCTION_NAME = REMOTE_FUNCTION_NAME.replace(PROJECT_ID + \".\", \"\")\n", + "# ! bq rm --routine --force=true {REMOTE_FUNCTION_NAME}" + ] + }, + { + "cell_type": "code", + "source": [ + "# # Delete the Google Cloud Storage bucket and files\n", + "# ! gsutil rm -r gs://{BUCKET_ID}\n", + "# print(f\"Deleted bucket '{BUCKET_ID}'.\")" + ], + "metadata": { + "id": "iQFo6OUBLmi3" + }, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "colab": { + "toc_visible": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb b/notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb new file mode 100644 index 00000000000..9a74beaad86 --- /dev/null +++ b/notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb @@ -0,0 +1,743 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "## Train a linear regression model with BigQuery DataFrames ML\n", + "\n", + "\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "24743cf4a1e1" + }, + "source": [ + "**_NOTE_**: This notebook has been tested in the following environment:\n", + "\n", + "* Python version = 3.10" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "Use this notebook to learn how to train a linear regression model by using BigQuery DataFrames ML. BigQuery DataFrames ML provides a provides a scikit-learn-like API for ML powered by the BigQuery engine.\n", + "\n", + "This example is adapted from the [BQML linear regression tutorial](https://cloud.google.com/bigquery-ml/docs/linear-regression-tutorial).\n", + "\n", + "Learn more about [BigQuery DataFrames](https://cloud.google.com/python/docs/reference/bigframes/latest)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you use BigQuery DataFrames to create a linear regression model that predicts the weight of an Adelie penguin based on the penguin's island of residence, culmen length and depth, flipper length, and sex.\n", + "\n", + "The steps include:\n", + "\n", + "- Creating a DataFrame from a BigQuery table.\n", + "- Cleaning and preparing data using pandas.\n", + "- Creating a linear regression model using `bigframes.ml`.\n", + "- Saving the ML model to BigQuery for future use." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses the [```penguins``` table](https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=ml_datasets&t=penguins) (a BigQuery Public Dataset) which includes data on a set of penguins including species, island of residence, weight, culmen length and depth, flipper length, and sex." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery (compute)\n", + "* BigQuery ML\n", + "\n", + "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models)\n", + "and [BigQuery ML pricing](https://cloud.google.com/bigquery/pricing#bqml),\n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages, which are required to run this notebook:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9O0Ka4W2MNF3" + }, + "outputs": [], + "source": [ + "!pip install bigframes" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only\n", + "\n", + "Uncomment and run the following cell to restart the kernel:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "Complete the tasks in this section to set up your environment." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oDfTjfACBvJk" + }, + "source": [ + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Enable the BigQuery API](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com).\n", + "\n", + "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "If you don't know your project ID, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Set the region\n", + "\n", + "You can also change the `REGION` variable used by BigQuery. Learn more about [BigQuery regions](https://cloud.google.com/bigquery/docs/locations#supported_locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "eF-Twtc4XGem" + }, + "outputs": [], + "source": [ + "REGION = \"US\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you might have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**Vertex AI Workbench**\n", + "\n", + "Do nothing, you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**Local JupyterLab instance**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**Colab**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "### Set BigQuery DataFrames options" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NPPMuw2PXGeo" + }, + "outputs": [], + "source": [ + "bf.options.bigquery.project = PROJECT_ID\n", + "bf.options.bigquery.location = REGION" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "D21CoOlfFTYI" + }, + "source": [ + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9EMAqR37AfLS" + }, + "source": [ + "## Read a BigQuery table into a BigQuery DataFrames DataFrame\n", + "\n", + "Read the [```penguins``` table](https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=ml_datasets&t=penguins) into a BigQuery DataFrames DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EDAaIwHpQCDZ" + }, + "outputs": [], + "source": [ + "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DJu837YEXD7B" + }, + "source": [ + "Take a look at the DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "_gPD0Zn1Stdb" + }, + "outputs": [], + "source": [ + "df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "rwPLjqW2Ajzh" + }, + "source": [ + "## Clean and prepare data\n", + "\n", + "You can use pandas as you normally would on the BigQuery DataFrames DataFrame, but calculations happen in the BigQuery query engine instead of your local environment.\n", + "\n", + "Because this model will focus on the Adelie Penguin species, you need to filter the data for only those rows representing Adelie penguins. Then you drop the `species` column because it is no longer needed.\n", + "\n", + "As these functions are applied, only the new DataFrame object `adelie_data` is modified. The source table and the original DataFrame object `df` don't change." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6i6HkFJZa8na" + }, + "outputs": [], + "source": [ + "# Filter down to the data to the Adelie Penguin species\n", + "adelie_data = df[df.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", + "\n", + "# Drop the species column\n", + "adelie_data = adelie_data.drop(columns=[\"species\"])\n", + "\n", + "# Take a look at the filtered DataFrame\n", + "adelie_data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "jhK2OlyMbY4L" + }, + "source": [ + "Drop rows with `NULL` values in order to create a BigQuery DataFrames DataFrame for the training data:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "0am3hdlXZfxZ" + }, + "outputs": [], + "source": [ + "# Drop rows with nulls to get training data\n", + "training_data = adelie_data.dropna()\n", + "\n", + "# Take a peek at the training data\n", + "training_data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "M_-0X7NxYK5f" + }, + "source": [ + "Specify your feature (or input) columns and the label (or output) column:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "YKwCW7Nsavap" + }, + "outputs": [], + "source": [ + "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "label_columns = training_data[['body_mass_g']]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "CjyM7vZJZ0sQ" + }, + "source": [ + "There is a row within the `adelie_data` BigQuery DataFrames DataFrame that has a `NULL` value for the `body mass` column. `body mass` is the label column, which is the value that the model you are creating is trying to predict.\n", + "\n", + "Create a new BigQuery DataFrames DataFrame, `test_data`, for this row so that you can use it as test data on which to make a prediction later:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "wej78IDUaRW9" + }, + "outputs": [], + "source": [ + "test_data = adelie_data[adelie_data.body_mass_g.isnull()]\n", + "\n", + "test_data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Fx4lsNqMorJ-" + }, + "source": [ + "## Create the linear regression model\n", + "\n", + "BigQuery DataFrames ML lets you move from exploring data to creating machine learning models through its scikit-learn-like API, `bigframes.ml`. BigQuery DataFrames ML supports several types of [ML models](https://cloud.google.com/python/docs/reference/bigframes/latest#ml-capabilities).\n", + "\n", + "In this notebook, you create a linear regression model, a type of regression model that generates a continuous value from a linear combination of input features.\n", + "\n", + "When you create a model with BigQuery DataFrames ML, it is saved locally and limited to the BigQuery session. However, as you'll see in the next section, you can use `to_gbq` to save the model permanently to your BigQuery project." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "EloGtMnverFF" + }, + "source": [ + "### Create the model using `bigframes.ml`\n", + "\n", + "When you pass the feature columns without transforms, BigQuery ML uses\n", + "[automatic preprocessing](https://cloud.google.com/bigquery/docs/auto-preprocessing) to encode string values and scale numeric values.\n", + "\n", + "BigQuery ML also [automatically splits the data for training and evaluation](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-glm#data_split_method), although for datasets with less than 500 rows (such as this one), all rows are used for training." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "GskyyUQPowBT" + }, + "outputs": [], + "source": [ + "from bigframes.ml.linear_model import LinearRegression\n", + "\n", + "model = LinearRegression()\n", + "\n", + "model.fit(feature_columns, label_columns)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "UGjeMPC2caKK" + }, + "source": [ + "### Score the model\n", + "\n", + "Check how the model performed by using the `score` method. More information on model scoring can be found [here](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#mlevaluate_output)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "kGBJKafpo0dl" + }, + "outputs": [], + "source": [ + "model.score(feature_columns, label_columns)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "P2lUiZZ_cjri" + }, + "source": [ + "### Predict using the model\n", + "\n", + "Use the model to predict the body mass of the data row you saved earlier to the `test_data` DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bsQ9cmoWo0Ps" + }, + "outputs": [], + "source": [ + "model.predict(test_data)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "GTRdUw-Ro5R1" + }, + "source": [ + "## Save the model in BigQuery\n", + "\n", + "The model is saved locally within this session. You can save the model permanently to BigQuery for use in future sessions, and to make the model sharable with others." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "K0mPaoGpcwwy" + }, + "source": [ + "Create a BigQuery dataset to house the model, adding a name for your dataset as the `DATASET_ID` variable:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZSP7gt13QrQt" + }, + "outputs": [], + "source": [ + "DATASET_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "from google.cloud import bigquery\n", + "client = bigquery.Client(project=PROJECT_ID)\n", + "dataset = bigquery.Dataset(PROJECT_ID + \".\" + DATASET_ID)\n", + "dataset.location = REGION\n", + "dataset = client.create_dataset(dataset, exists_ok=True)\n", + "print(f\"Dataset {dataset.dataset_id} created.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zqAIWWgJczp-" + }, + "source": [ + "Save the model using the `to_gbq` method:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "QE_GD4Byo_jb" + }, + "outputs": [], + "source": [ + "model.to_gbq(DATASET_ID + \".penguin_weight\" , replace=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f7uHacAy49rT" + }, + "source": [ + "You can view the saved model in the BigQuery console under the dataset you created in the first step. Run the following cell and follow the link to view your BigQuery console:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "qDBoiA_0488Z" + }, + "outputs": [], + "source": [ + "print(f'https://console.developers.google.com/bigquery?p={PROJECT_ID}')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "G_wjSfXpWTuy" + }, + "source": [ + "# Summary and next steps\n", + "\n", + "You've created a linear regression model using `bigframes.ml`.\n", + "\n", + "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "# # Delete the BigQuery dataset and associated ML model\n", + "# from google.cloud import bigquery\n", + "# client = bigquery.Client(project=PROJECT_ID)\n", + "# client.delete_dataset(\n", + "# DATASET_ID, delete_contents=True, not_found_ok=True\n", + "# )\n", + "# print(\"Deleted dataset '{}'.\".format(DATASET_ID))" + ] + } + ], + "metadata": { + "colab": { + "provenance": [], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/getting_started/getting_started_bq_dataframes.ipynb b/notebooks/getting_started/getting_started_bq_dataframes.ipynb new file mode 100644 index 00000000000..7815182e54d --- /dev/null +++ b/notebooks/getting_started/getting_started_bq_dataframes.ipynb @@ -0,0 +1,971 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "# Get started with BigQuery DataFrames\n", + "\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "24743cf4a1e1" + }, + "source": [ + "**_NOTE_**: This notebook has been tested in the following environment:\n", + "\n", + "* Python version = 3.10" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "Use this notebook to get started with BigQuery DataFrames, including setup, installation, and basic tutorials.\n", + "\n", + "BigQuery DataFrames provides a Pythonic DataFrame and machine learning (ML) API powered by the BigQuery engine.\n", + "\n", + "* `bigframes.pandas` provides a pandas-like API for analytics.\n", + "* `bigframes.ml` provides a scikit-learn-like API for ML.\n", + "\n", + "Learn more about [BigQuery DataFrames](https://cloud.google.com/python/docs/reference/bigframes/latest)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn how to install BigQuery DataFrames, load data into a BigQuery DataFrames DataFrame, and inspect and manipulate the data using pandas and a custom Python function, running at BigQuery scale.\n", + "\n", + "The steps include:\n", + "\n", + "- Creating a BigQuery DataFrames DataFrame: Access data from a local CSV to create a BigQuery DataFrames DataFrame.\n", + "- Inspecting and manipulating data: Use pandas to perform data cleaning and preparation on the DataFrame.\n", + "- Deploying a custom function: Deploy a [remote function ](https://cloud.google.com/bigquery/docs/remote-functions)that runs a scalar Python function at BigQuery scale." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses the [```penguins``` table](https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=ml_datasets&t=penguins) (a BigQuery public dataset), which contains data on a set of penguins including species, island of residence, weight, culmen length and depth, flipper length, and sex.\n", + "\n", + "The same dataset is also stored in a public Cloud Storage bucket as a CSV file so that you can use it to try ingesting data from a local environment." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery (storage and compute)\n", + "* Cloud Functions\n", + "\n", + "Learn about [BigQuery storage pricing](https://cloud.google.com/bigquery/pricing#storage),\n", + "[BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models),\n", + "and [Cloud Functions pricing](https://cloud.google.com/functions/pricing),\n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages, which are required to run this notebook:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "mfPoOwPLGpSr" + }, + "outputs": [], + "source": [ + "!pip install bigframes" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only\n", + "\n", + "Uncomment and run the following cell to restart the kernel:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "Complete the tasks in this section to set up your environment." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Yq7zKYWelRQP" + }, + "source": [ + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Click here](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com,bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com,cloudresourcemanager.googleapis.com) to enable the following APIs:\n", + "\n", + " * BigQuery API\n", + " * BigQuery Connection API\n", + " * Cloud Functions API\n", + " * Cloud Run API\n", + " * Artifact Registry API\n", + " * Cloud Build API\n", + " * Cloud Resource Manager API\n", + "\n", + "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "If you don't know your project ID, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Set the region\n", + "\n", + "You can also change the `REGION` variable used by BigQuery. Learn more about [BigQuery regions](https://cloud.google.com/bigquery/docs/locations#supported_locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "eF-Twtc4XGem" + }, + "outputs": [], + "source": [ + "REGION = \"US\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you might have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**Vertex AI Workbench**\n", + "\n", + "Do nothing, you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**Local JupyterLab instance**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**Colab**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "\n", + "### Set BigQuery DataFrames options" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NPPMuw2PXGeo" + }, + "outputs": [], + "source": [ + "bf.options.bigquery.project = PROJECT_ID\n", + "bf.options.bigquery.location = REGION" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "pDfrKwMKE_dK" + }, + "source": [ + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-19Uiwoo9pP4" + }, + "source": [ + "## See the power of BigQuery DataFrames first-hand\n", + "\n", + "BigQuery DataFrames enables you to interact with datasets of any size, so that you can explore, transform, and understand even your biggest datasets using familiar tools like pandas and scikit-learn." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "KMX4D2uoBwM0" + }, + "source": [ + "For example, take the BigQuery sample table `bigquery-samples.wikipedia_pageviews.200809h`, which is ~60 GB is size. This is not a dataset you'd likely be able process in pandas without extra infrastructure.\n", + "\n", + "With BigQuery DataFrames, however, computation is handled by BigQuery's highly scalable compute engine, meaning you can focus on doing data science without hitting size limitations." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i98c46p1CXoV" + }, + "source": [ + "If you'd like to try creating a BigQuery DataFrames DataFrame from this table, uncomment and run the next cell to load the table using the `read_gbq` method.\n", + "\n", + "> Note: Keep in mind that running these operations will count against your monthly [free tier allowance in BigQuery](https://cloud.google.com/bigquery/pricing#free-tier)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Vyex9BQI-BNa" + }, + "outputs": [], + "source": [ + "# bq_df_sample = bf.read_gbq(\"bigquery-samples.wikipedia_pageviews.200809h\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gE6CEALjDZZV" + }, + "source": [ + "No problem! BigQuery DataFrames makes a DataFrame, `bq_df_sample`, containing the entirety of the source table of data." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "T6lAIeelDwLz" + }, + "source": [ + "Uncomment and run the following cell to see pandas in action over your new BigQuery DataFrames DataFrame.\n", + "\n", + "This code uses regex to filter the DataFrame to include only rows with Wikipedia page titles containing the word \"Google\", sums the total views by page title, and then returns the top 100 results." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "XfGq5apK-D_e" + }, + "outputs": [], + "source": [ + "# bq_df_sample[bq_df_sample.title.str.contains(r\"[Gg]oogle\")]\\\n", + "# .groupby(['title'], as_index=False)['views'].sum(numeric_only=True)\\\n", + "# .sort_values('views', ascending=False)\\\n", + "# .head(100)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i6XV-HTN-IFF" + }, + "source": [ + "In addition to giving you access to pandas, BigQuery DataFrames also enables you to build ML models, run inference, and deploy and run your own Python functions at scale. You'll see examples throughout this and other notebooks in this GitHub repo.\n", + "\n", + "Now you'll move to the smaller `penguins` dataset for the remainder of this getting started guide." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9EMAqR37AfLS" + }, + "source": [ + "## Create a BigQuery DataFrames DataFrame\n", + "\n", + "You can create a BigQuery DataFrames DataFrame by reading data from any of the following locations:\n", + "\n", + "* A local data file\n", + "* Data stored in a BigQuery table\n", + "* A data file stored in Cloud Storage\n", + "* An in-memory pandas DataFrame\n", + "\n", + "The following sections show how to use the first two options." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "iZDjzglh9eWZ" + }, + "source": [ + "### Create a DataFrame from a local file\n", + "\n", + "Use the instructions in the following sections to create a BigQuery DataFrames DataFrame from a local file.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8Jry3NoFv3Wm" + }, + "source": [ + "#### Get the CSV file\n", + "\n", + "First, copy and paste the following link into a new browser window to download the CSV file of the penguin data to your local machine:\n", + "\n", + "> http://storage.googleapis.com/cloud-samples-data/vertex-ai/bigframe/penguins.csv\n", + "\n", + "Next, upload the local CSV file to your notebook environment, using the relevant instructions for your environment:\n", + "\n", + "**Vertex AI Workbench or a local JupyterLab instance**\n", + "\n", + "1. Follow these [directions](https://jupyterlab.readthedocs.io/en/latest/user/files.html#uploading-and-downloading) to upload the file from your machine to your notebook environment by using the UI.\n", + "2. Uncomment the next cell, set the variable `fn` to match the path to your file, and then run the cell." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "SvyXzkRl783u" + }, + "outputs": [], + "source": [ + "# fn = 'penguins.csv'" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "yqcuF1JNvFse" + }, + "source": [ + "**Colab**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3QHQYlnoBLpt" + }, + "outputs": [], + "source": [ + "# from google.colab import files\n", + "# uploaded = files.upload()\n", + "# for fn in uploaded.keys():\n", + "# print('User uploaded file \"{name}\" with length {length} bytes'.format(\n", + "# name=fn, length=len(uploaded[fn])))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sJsrwAQY_H6g" + }, + "source": [ + "#### Create a DataFrame\n", + "\n", + "Create a BigQuery DataFrames DataFrame from the uploaded CSV file:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EDAaIwHpQCDZ" + }, + "outputs": [], + "source": [ + "df_from_local = bf.read_csv(fn)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "U-RVfNCu_h_h" + }, + "source": [ + "Take a look at the first few rows of the DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "_gPD0Zn1Stdb" + }, + "outputs": [], + "source": [ + "df_from_local.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "rK0lNJmz_xkA" + }, + "source": [ + "### Ingest data from a DataFrame to a BigQuery table\n", + "\n", + "BigQuery DataFrames lets you create a BigQuery table from a BigQuery DataFrames DataFrame on-the-fly." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "V1DWpmSCAEql" + }, + "source": [ + "First, create a BigQuery dataset to house the table. Choose a name for your dataset, or keep the suggestion of `birds`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZSP7gt13QrQt" + }, + "outputs": [], + "source": [ + "DATASET_ID = \"birds\" # @param {type:\"string\"}\n", + "\n", + "from google.cloud import bigquery\n", + "client = bigquery.Client(project=PROJECT_ID)\n", + "dataset = bigquery.Dataset(PROJECT_ID + \".\" + DATASET_ID)\n", + "dataset.location = REGION\n", + "dataset = client.create_dataset(dataset, exists_ok=True)\n", + "print(f\"Dataset {dataset.dataset_id} created.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Jd0dFISwAPPa" + }, + "source": [ + "Next, use the `to_gbq` method to create a BigQuery table from the DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oP1NIAmUBjop" + }, + "outputs": [], + "source": [ + "df_from_local.to_gbq(PROJECT_ID + \".\" + DATASET_ID + \".penguins\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "kfF6fnmmAZEK" + }, + "source": [ + "### Create a DataFrame from BigQuery data\n", + "You can create a BigQuery DataFrames DataFrame from a BigQuery table by using the `read_gbq` method and referencing either an entire table or a SQL query." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TEy5jHJDD6hx" + }, + "source": [ + "Create a BigQuery DataFrames DataFrame from the BigQuery table you created in the previous section, and view a few rows:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "IBuo-d6dWfsA" + }, + "outputs": [], + "source": [ + "query_or_table = f\"\"\"{PROJECT_ID}.{DATASET_ID}.penguins\"\"\"\n", + "bq_df = bf.read_gbq(query_or_table)\n", + "bq_df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "rwPLjqW2Ajzh" + }, + "source": [ + "## Inspect and manipulate data in BigQuery DataFrames" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "bExmYlL_ELtV" + }, + "source": [ + "### Using pandas\n", + "\n", + "You can use pandas as you normally would on the BigQuery DataFrames DataFrame, but calculations happen in the BigQuery query engine instead of your local environment. There are 150+ pandas functions supported in BigQuery DataFrames. You can view the list in [the documentation](https://cloud.google.com/python/docs/reference/bigframes/latest)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ZHFUc3Q_FHc1" + }, + "source": [ + "To see this in action, inspect one of the columns (or series) of the BigQuery DataFrames DataFrame:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6i6HkFJZa8na" + }, + "outputs": [], + "source": [ + "bq_df[\"body_mass_g\"].head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "EJIZJaNXFQzh" + }, + "source": [ + "Compute the mean of this series:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "YKwCW7Nsavap" + }, + "outputs": [], + "source": [ + "average_body_mass = bq_df[\"body_mass_g\"].mean()\n", + "print(f\"average_body_mass: {average_body_mass}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DSs1cnca-MOU" + }, + "source": [ + "Calculate the mean `body_mass_g` by `species` using the `groupby` operation:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4PyKMR61-Mjy" + }, + "outputs": [], + "source": [ + "bq_df[\"species\", \"body_mass_g\"].groupby(by=bq_df[\"species\"]).mean(numeric_only=True).head()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6sf9kZ2C9Ixe" + }, + "source": [ + "You can confirm that the calculations were run in BigQuery by clicking \"Open job\" from the previous cells' output. This takes you to the BigQuery console to view the SQL statement and job details." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cWVNZ8D_FUtT" + }, + "source": [ + "### Using custom functions\n", + "\n", + "Running your own Python functions (or being able to bring your packages) and using them at scale is a challenge many data scientists face. BigQuery DataFrames makes it easy to deploy [remote functions](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.pandas#bigframes_pandas_remote_function) that run scalar Python functions at BigQuery scale. These functions are persisted as [BigQuery remote functions](https://cloud.google.com/bigquery/docs/remote-functions) that you can then re-use." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zjw8toUbHuRD" + }, + "source": [ + "Running the cell below creates a custom function using the `remote_function` method. This function categorizes a value into one of two buckets: >= 4000 or <4000.\n", + "\n", + "> Note: Creating a function requires a [BigQuery connection](https://cloud.google.com/bigquery/docs/remote-functions#create_a_remote_function). This code assumes a pre-created connection named `bigframes-rf-conn`. If\n", + "the connection is not already created, BigQuery DataFrames attempts to create one assuming the [necessary APIs\n", + "and IAM permissions](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.pandas#bigframes_pandas_remote_function) are set up in the project.\n", + "\n", + "This cell takes a few minutes to run because it creates the BigQuery connection (if applicable) and deploys the Cloud Function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "rSWTOG-vb2Fc" + }, + "outputs": [], + "source": [ + "@bf.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", + "def get_bucket(num):\n", + " if not num: return \"NA\"\n", + " boundary = 4000\n", + " return \"at_or_above_4000\" if num >= boundary else \"below_4000\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "N7JH0BI5IOpK" + }, + "source": [ + "The custom function is deployed as a Cloud Function, and is then integrated with BigQuery as a remote function.\n", + "\n", + "Save both of the function names so that you can clean them up at the end of this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6ejPXoyEQpWE" + }, + "outputs": [], + "source": [ + "CLOUD_FUNCTION_NAME = format(get_bucket.bigframes_cloud_function)\n", + "print(\"Cloud Function Name \" + CLOUD_FUNCTION_NAME)\n", + "REMOTE_FUNCTION_NAME = format(get_bucket.bigframes_remote_function)\n", + "print(\"Remote Function Name \" + REMOTE_FUNCTION_NAME)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vHV3JqKjJHsH" + }, + "source": [ + "Apply the custom function to the BigQuery DataFrames DataFrame to bucketize the `body_mass_g` value of the penguins:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NxSd9WZFcIji" + }, + "outputs": [], + "source": [ + "bq_df = bq_df.assign(body_mass_bucket=bq_df['body_mass_g'].apply(get_bucket))\n", + "bq_df[['body_mass_g', 'body_mass_bucket']].head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "wCsmt0IwFkDy" + }, + "source": [ + "## Summary and next steps\n", + "\n", + "You've created BigQuery DataFrames DataFrames, and inspected and manipulated data with pandas and custom remote functions at BigQuery scale and speed.\n", + "\n", + "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks), including an introductory notebook for `bigframes.ml`." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "### Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "# # Delete the BigQuery dataset\n", + "# from google.cloud import bigquery\n", + "# client = bigquery.Client(project=PROJECT_ID)\n", + "# client.delete_dataset(\n", + "# DATASET_ID, delete_contents=True, not_found_ok=True\n", + "# )\n", + "# print(\"Deleted dataset '{}'.\".format(DATASET_ID))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "_dTCXvCxtPw9" + }, + "outputs": [], + "source": [ + "# # Delete the BigQuery Connection\n", + "# from google.cloud import bigquery_connection_v1 as bq_connection\n", + "# client = bq_connection.ConnectionServiceClient()\n", + "# CONNECTION_ID = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/bigframes-rf-conn\"\n", + "# client.delete_connection(name=CONNECTION_ID)\n", + "# print(\"Deleted connection '{}'.\".format(CONNECTION_ID))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EDAIIfcpwNOF" + }, + "outputs": [], + "source": [ + "# # Delete the Cloud Function\n", + "# ! gcloud functions delete {CLOUD_FUNCTION_NAME} --quiet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "QwumLUKmVpuH" + }, + "outputs": [], + "source": [ + "# # Delete the Remote Function\n", + "# REMOTE_FUNCTION_NAME = REMOTE_FUNCTION_NAME.replace(PROJECT_ID + \".\", \"\")\n", + "# ! bq rm --routine --force=true {REMOTE_FUNCTION_NAME}" + ] + } + ], + "metadata": { + "colab": { + "provenance": [], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/03 - Using ML - ML fundamentals.ipynb b/notebooks/getting_started/ml_fundamentals.ipynb similarity index 100% rename from notebooks/03 - Using ML - ML fundamentals.ipynb rename to notebooks/getting_started/ml_fundamentals.ipynb diff --git a/notebooks/10 - Regionalized.ipynb b/notebooks/location/regionalized.ipynb similarity index 100% rename from notebooks/10 - Regionalized.ipynb rename to notebooks/location/regionalized.ipynb diff --git a/notebooks/05 - Using ML - Easy linear regression.ipynb b/notebooks/regression/easy_linear_regression.ipynb similarity index 100% rename from notebooks/05 - Using ML - Easy linear regression.ipynb rename to notebooks/regression/easy_linear_regression.ipynb diff --git a/notebooks/04 - Using ML - SKLearn linear regression.ipynb b/notebooks/regression/sklearn_linear_regression.ipynb similarity index 100% rename from notebooks/04 - Using ML - SKLearn linear regression.ipynb rename to notebooks/regression/sklearn_linear_regression.ipynb diff --git a/notebooks/50 - Remote Function.ipynb b/notebooks/remote_functions/remote_function.ipynb similarity index 100% rename from notebooks/50 - Remote Function.ipynb rename to notebooks/remote_functions/remote_function.ipynb diff --git a/noxfile.py b/noxfile.py index 7d4cb1c61b1..ec5a1b16510 100644 --- a/noxfile.py +++ b/noxfile.py @@ -19,6 +19,7 @@ from multiprocessing import Process import os import pathlib +from pathlib import Path import re import shutil from typing import Dict, List @@ -30,6 +31,10 @@ BLACK_VERSION = "black==22.3.0" ISORT_VERSION = "isort==5.12.0" SPHINX_VERSION = "sphinx==4.5.0" + +# pytest-retry 1.4.0 fails on Python 3.11. +# https://github.com/str0zzapreti/pytest-retry/issues/17 +PYTEST_RETRY_VERSION = "pytest-retry<=1.3.0" LINT_PATHS = ["docs", "bigframes", "tests", "noxfile.py", "setup.py"] DEFAULT_PYTHON_VERSION = "3.10" @@ -55,7 +60,7 @@ "openpyxl", "pytest", "pytest-cov", - "pytest-retry", + PYTEST_RETRY_VERSION, "pytest-timeout", "pytest-xdist", "google-cloud-testutils", @@ -590,23 +595,37 @@ def system_prerelease(session: nox.sessions.Session): @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) def notebook(session): session.install("-e", ".[all]") - session.install("pytest", "pytest-xdist", "pytest-retry", "nbmake") - - notebooks = [ - "00 - Summary.ipynb", - "01 - Getting Started.ipynb", - "02 - DataFrame.ipynb", - "03 - Using ML - ML fundamentals.ipynb", - "04 - Using ML - SKLearn linear regression.ipynb", - "05 - Using ML - Easy linear regression.ipynb", - "06 - Using ML - Large Language Models.ipynb", - "50 - Remote Function.ipynb", + session.install("pytest", "pytest-xdist", PYTEST_RETRY_VERSION, "nbmake") + + notebooks_list = list(Path("notebooks/").glob("*/*.ipynb")) + + denylist = [ + # Regionalized testing is manually added later. + "notebooks/location/regionalized.ipynb", + # These notebooks contain special colab `param {type:"string"}` + # comments, which make it easy for customers to fill in their + # own information. + # TODO(ashleyxu): Test these notebooks by replacing parameters with + # appropriate values and omitting cleanup logic that may break + # our test infrastructure. + "notebooks/getting_started/getting_started_bq_dataframes.ipynb", + "notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb", + "notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb", + "notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb", + # The experimental notebooks imagine features that don't yet + # exist or only exist as temporary prototypes. + "notebooks/experimental/longer_ml_demo.ipynb", ] - notebooks = [os.path.join("notebooks", nb) for nb in notebooks] + + # Convert each Path notebook object to a string using a list comprehension. + notebooks = [str(nb) for nb in notebooks_list] + + # Remove tests that we choose not to test. + notebooks = list(filter(lambda nb: nb not in denylist, notebooks)) # Regionalized notebooks notebooks_reg = { - "10 - Regionalized.ipynb": [ + "regionalized.ipynb": [ "asia-southeast1", "eu", "europe-west4", @@ -616,7 +635,8 @@ def notebook(session): ] } notebooks_reg = { - os.path.join("notebooks", nb): regions for nb, regions in notebooks_reg.items() + os.path.join("notebooks/location", nb): regions + for nb, regions in notebooks_reg.items() } # For some reason nbmake exits silently with "no tests ran" message if diff --git a/pytest.ini b/pytest.ini index 693439f47ca..204c743bbfa 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,4 @@ [pytest] +doctest_optionflags = NORMALIZE_WHITESPACE filterwarnings = ignore::pandas.errors.SettingWithCopyWarning diff --git a/samples/snippets/remote_function.py b/samples/snippets/remote_function.py new file mode 100644 index 00000000000..37972672c32 --- /dev/null +++ b/samples/snippets/remote_function.py @@ -0,0 +1,147 @@ +# 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): + 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 would be using a pre-created connection named + # `bigframes-rf-conn`. 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([float], str, bigquery_connection="bigframes-rf-conn") + def get_bucket(num): + 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 decorated it + # as a remote function + @bpd.remote_function([str], str, bigquery_connection="bigframes-rf-conn") + def get_hash(input): + import hashlib + + # handle missing value + if input is None: + input = "" + encoded_input = input.encode() + hash = hashlib.md5(encoded_input) + return hash.hexdigest() + + # 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..8b51e46b45a --- /dev/null +++ b/samples/snippets/remote_function_test.py @@ -0,0 +1,32 @@ +# 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 remote_function + + +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.reset_session() + + # TODO(swast): Get project from environment so contributors can run tests. + remote_function.run_remote_function_and_read_gbq_function("bigframes-dev") + out, _ = capsys.readouterr() + assert "Created BQ remote function:" in out + assert "Created cloud function:" in out diff --git a/tests/data/hockey_players.json b/tests/data/hockey_players.json new file mode 100644 index 00000000000..8a9b252992d --- /dev/null +++ b/tests/data/hockey_players.json @@ -0,0 +1,37 @@ +[ + { + "mode": "NULLABLE", + "name": "team_name", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "position", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "player_name", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "goals", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "assists", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "number", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "season", + "type": "INTEGER" + } +] diff --git a/tests/data/hockey_players.jsonl b/tests/data/hockey_players.jsonl new file mode 100644 index 00000000000..d2b26cffdd0 --- /dev/null +++ b/tests/data/hockey_players.jsonl @@ -0,0 +1,10 @@ +{"team_name":"Canucks", "position":"C", "player_name":"Elias Petterson", "goals":39, "assists":63, "number":40, "season":2023} +{"team_name":"Canucks", "position":"LW", "player_name":"Ilya Mikheyev", "goals":13, "assists":15, "number":65, "season":2023} +{"team_name":"Canucks", "position":"RW", "player_name":"Andrei Kuzmenko", "goals":39, "assists":35, "number":40, "season":2023} +{"team_name":"Kraken", "position":"C", "player_name":"Jared McCann", "goals":40, "assists":30, "number":19, "season":2023} +{"team_name":"Kraken", "position":"LW", "player_name":"Yanni Gourde", "goals":14, "assists":34, "number":37, "season":2023} +{"team_name":"Kraken", "position":"RW", "player_name":"Jordan Eberle", "goals":20, "assists":43, "number":7, "season":2023} +{"team_name":"Canucks", "position":"C", "player_name":"Elias Petterson", "goals":32, "assists":36, "number":40, "season":2022} +{"team_name":"Kraken", "position":"C", "player_name":"Jared McCann", "goals":27, "assists":23, "number":19, "season":2022} +{"team_name":"Kraken", "position":"LW", "player_name":"Yanni Gourde", "goals":21, "assists":27, "number":37, "season":2022} +{"team_name":"Kraken", "position":"RW", "player_name":"Jordan Eberle", "goals":21, "assists":23, "number":7, "season":2022} diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 41c8eaffd78..bf5cf12c74b 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -211,6 +211,7 @@ def load_test_data_tables( ("scalars_too", "scalars_schema.json", "scalars.jsonl"), ("penguins", "penguins_schema.json", "penguins.jsonl"), ("time_series", "time_series_schema.json", "time_series.jsonl"), + ("hockey_players", "hockey_players.json", "hockey_players.jsonl"), ]: test_data_hash = hashlib.md5() _hash_digest_file(test_data_hash, DATA_DIR / schema_filename) @@ -255,6 +256,11 @@ def scalars_table_id(test_data_tables) -> str: return test_data_tables["scalars"] +@pytest.fixture(scope="session") +def hockey_table_id(test_data_tables) -> str: + return test_data_tables["hockey_players"] + + @pytest.fixture(scope="session") def scalars_table_id_2(test_data_tables) -> str: return test_data_tables["scalars_too"] @@ -354,6 +360,34 @@ def scalars_dfs( return scalars_df_index, scalars_pandas_df_index +@pytest.fixture(scope="session") +def hockey_df( + hockey_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at test data.""" + return session.read_gbq(hockey_table_id) + + +@pytest.fixture(scope="session") +def hockey_pandas_df() -> pd.DataFrame: + """pd.DataFrame pointing at test data.""" + df = pd.read_json( + DATA_DIR / "hockey_players.jsonl", + lines=True, + dtype={ + "team_name": pd.StringDtype(storage="pyarrow"), + "position": pd.StringDtype(storage="pyarrow"), + "player_name": pd.StringDtype(storage="pyarrow"), + "goals": pd.Int64Dtype(), + "assists": pd.Int64Dtype(), + "number": pd.Int64Dtype(), + "season": pd.Int64Dtype(), + }, + ) + df.index = df.index.astype("Int64") + return df + + @pytest.fixture(scope="session") def penguins_df_default_index( penguins_table_id: str, session: bigframes.Session @@ -721,3 +755,48 @@ def restore_sampling_settings(): yield bigframes.options.sampling.enable_downsampling = enable_downsampling bigframes.options.sampling.max_download_size = max_download_size + + +@pytest.fixture() +def weird_strings_pd(): + df = pd.DataFrame( + { + "string_col": [ + "٠١٢٣٤٥٦٧٨٩", + "", + "0", + "字", + "五", + "0123456789", + pd.NA, + "abc 123 mixed letters and numbers", + "no numbers here", + "123a", + "23!", + " 45", + "a45", + "Dž", + "tT", + "-123", + "-123.4", + "-0", + "-.0", + ".0", + ".1", + "⅙", + "²", + "\t", + "a\ta", + "p1\np2", + " ", + ] + }, + dtype=pd.StringDtype(storage="pyarrow"), + ) + df.index = df.index.astype("Int64") + return df.string_col + + +@pytest.fixture() +def weird_strings(session, weird_strings_pd): + return session.read_pandas(weird_strings_pd.to_frame()).string_col diff --git a/tests/system/large/ml/test_compose.py b/tests/system/large/ml/test_compose.py index 0c2744819d4..b65baa63ebd 100644 --- a/tests/system/large/ml/test_compose.py +++ b/tests/system/large/ml/test_compose.py @@ -21,7 +21,7 @@ import bigframes.ml.preprocessing -def test_columntransformer_standalone_fit_transform( +def test_columntransformer_standalone_fit_and_transform( penguins_df_default_index, new_penguins_df ): transformer = bigframes.ml.compose.ColumnTransformer( @@ -73,3 +73,54 @@ def test_columntransformer_standalone_fit_transform( ) pandas.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_columntransformer_standalone_fit_transform(new_penguins_df): + transformer = bigframes.ml.compose.ColumnTransformer( + [ + ( + "onehot", + bigframes.ml.preprocessing.OneHotEncoder(), + "species", + ), + ( + "scale", + bigframes.ml.preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ] + ) + + result = transformer.fit_transform( + new_penguins_df[["species", "culmen_length_mm", "flipper_length_mm"]] + ).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pandas.DataFrame( + { + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + "scaled_culmen_length_mm": [ + 1.313249, + -0.20198, + -1.111118, + ], + "scaled_flipper_length_mm": [1.251098, -1.196588, -0.054338], + }, + index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"), + ) + expected.scaled_culmen_length_mm = expected.scaled_culmen_length_mm.astype( + "Float64" + ) + expected.scaled_flipper_length_mm = expected.scaled_flipper_length_mm.astype( + "Float64" + ) + + pandas.testing.assert_frame_equal(result, expected, rtol=1e-3) diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 74356c81e1b..181678ebcb2 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -15,6 +15,7 @@ from unittest import TestCase import numpy as np +import pytest def test_create_text_generator_model(palm2_text_generator_model): @@ -22,6 +23,8 @@ def test_create_text_generator_model(palm2_text_generator_model): assert palm2_text_generator_model is not None +# Marked as flaky only because BQML LLM is in preview, the service only has limited capacity, not stable enough. +@pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_default_params_success( palm2_text_generator_model, llm_text_df ): @@ -32,6 +35,7 @@ def test_text_generator_predict_default_params_success( assert all(series.str.len() > 20) +@pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_series_default_params_success( palm2_text_generator_model, llm_text_df ): @@ -42,6 +46,7 @@ def test_text_generator_predict_series_default_params_success( assert all(series.str.len() > 20) +@pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_arbitrary_col_label_success( palm2_text_generator_model, llm_text_df ): @@ -53,6 +58,7 @@ def test_text_generator_predict_arbitrary_col_label_success( assert all(series.str.len() > 20) +@pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_with_params_success( palm2_text_generator_model, llm_text_df ): @@ -70,6 +76,7 @@ def test_create_embedding_generator_model(palm2_embedding_generator_model): assert palm2_embedding_generator_model is not None +@pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df ): @@ -82,6 +89,7 @@ def test_embedding_generator_predict_success( assert value.size == 768 +@pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_series_success( palm2_embedding_generator_model, llm_text_df ): diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index 420a80754fb..57b9900c489 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -24,13 +24,13 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): scaler = bigframes.ml.preprocessing.StandardScaler() scaler.fit( penguins_df_default_index[ - "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm" + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] ] ) result = scaler.transform( penguins_df_default_index[ - "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm" + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] ] ).to_pandas() @@ -58,6 +58,35 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): pd.testing.assert_frame_equal(result, expected, rtol=1e-3) +def test_standard_scaler_normalizeds_fit_transform(new_penguins_df): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + scaler = bigframes.ml.preprocessing.StandardScaler() + result = scaler.fit_transform( + new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] + ).to_pandas() + + # If standard-scaled correctly, mean should be 0.0 + for column in result.columns: + assert math.isclose(result[column].mean(), 0.0, abs_tol=1e-3) + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "scaled_culmen_depth_mm": [1.17072, -1.272416, 0.101848], + "scaled_culmen_length_mm": [1.313249, -0.20198, -1.111118], + "scaled_flipper_length_mm": [1.251089, -1.196588, -0.054338], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + def test_standard_scaler_series_normalizes(penguins_df_default_index, new_penguins_df): # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. scaler = bigframes.ml.preprocessing.StandardScaler() @@ -93,7 +122,7 @@ def test_standard_scaler_series_normalizes(penguins_df_default_index, new_pengui def test_one_hot_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() - encoder.fit(new_penguins_df["species", "sex"]) + encoder.fit(new_penguins_df[["species", "sex"]]) result = encoder.transform(new_penguins_df).to_pandas() @@ -121,6 +150,35 @@ def test_one_hot_encoder_default_params(new_penguins_df): pd.testing.assert_frame_equal(result, expected) +def test_one_hot_encoder_default_params_fit_transform(new_penguins_df): + encoder = bigframes.ml.preprocessing.OneHotEncoder() + + result = encoder.fit_transform(new_penguins_df[["species", "sex"]]).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "onehotencoded_sex": [ + [{"index": 2, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + ], + "onehotencoded_species": [ + [{"index": 1, "value": 1.0}], + [{"index": 1, "value": 1.0}], + [{"index": 2, "value": 1.0}], + ], + }, + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + def test_one_hot_encoder_series_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() encoder.fit(new_penguins_df["species"]) @@ -148,7 +206,7 @@ def test_one_hot_encoder_series_default_params(new_penguins_df): def test_one_hot_encoder_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder("most_frequent", 100, 2) - encoder.fit(new_penguins_df["species", "sex"]) + encoder.fit(new_penguins_df[["species", "sex"]]) result = encoder.transform(new_penguins_df).to_pandas() @@ -178,7 +236,7 @@ def test_one_hot_encoder_params(new_penguins_df): def test_one_hot_encoder_different_data(penguins_df_default_index, new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() - encoder.fit(penguins_df_default_index["species", "sex"]) + encoder.fit(penguins_df_default_index[["species", "sex"]]) result = encoder.transform(new_penguins_df).to_pandas() diff --git a/tests/system/small/operations/test_strings.py b/tests/system/small/operations/test_strings.py index 31b64f4314d..241cbd576b1 100644 --- a/tests/system/small/operations/test_strings.py +++ b/tests/system/small/operations/test_strings.py @@ -254,31 +254,93 @@ def test_upper(scalars_dfs): ) -def test_isnumeric(session): - pandas_df = pd.DataFrame( - { - "numeric_string_col": [ - "٠١٢٣٤٥٦٧٨٩", - "", - "0", - "字", - "五", - "0123456789", - pd.NA, - "abc 123 mixed letters and numbers", - "no numbers here", - "123a", - "23!", - " 45", - "a45", - ] - } - ) - - df = session.read_pandas(pandas_df) - - pd_result = pandas_df.numeric_string_col.str.isnumeric() - bf_result = df.numeric_string_col.str.isnumeric().to_pandas() +def test_isnumeric(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isnumeric() + bf_result = weird_strings.str.isnumeric().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isalpha(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isalpha() + bf_result = weird_strings.str.isalpha().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isdigit(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isdigit() + bf_result = weird_strings.str.isdigit().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isdecimal(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isdecimal() + bf_result = weird_strings.str.isdecimal().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isalnum(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isalnum() + bf_result = weird_strings.str.isalnum().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isspace(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isspace() + bf_result = weird_strings.str.isspace().to_pandas() + + pd.testing.assert_series_equal( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_islower(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.islower() + bf_result = weird_strings.str.islower().to_pandas() + + assert_series_equal_ignoring_order( + bf_result, + pd_result.astype(pd.BooleanDtype()) + # the dtype here is a case of intentional diversion from pandas + # see go/bigframes-dtypes + ) + + +def test_isupper(weird_strings, weird_strings_pd): + pd_result = weird_strings_pd.str.isupper() + bf_result = weird_strings.str.isupper().to_pandas() assert_series_equal_ignoring_order( bf_result, @@ -394,9 +456,6 @@ def test_str_get(scalars_dfs): bf_result = bf_series.str.get(8).to_pandas() pd_result = scalars_pandas_df[col_name].str.get(8) - print(pd_result) - print(bf_result) - assert_series_equal_ignoring_order( pd_result, bf_result, @@ -416,6 +475,16 @@ def test_str_pad(scalars_dfs): ) +def test_str_zfill(weird_strings, weird_strings_pd): + bf_result = weird_strings.str.zfill(5).to_pandas() + pd_result = weird_strings_pd.str.zfill(5) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + def test_str_ljust(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 2c44dd80670..5b4f9ebcccf 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -840,6 +840,14 @@ def test_set_index(scalars_dfs, index_column, drop, append): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_set_index_key_error(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + with pytest.raises(KeyError): + scalars_pandas_df.set_index(["not_a_col"]) + with pytest.raises(KeyError): + scalars_df.set_index(["not_a_col"]) + + @pytest.mark.parametrize( ("ascending",), ((True,), (False,)), @@ -1321,6 +1329,41 @@ def test_df_describe(scalars_dfs): ).all() +@pytest.mark.parametrize( + ("values", "index", "columns"), + [ + ("int64_col", "int64_too", ["string_col"]), + (["int64_col"], "int64_too", ["string_col"]), + (["int64_col", "float64_col"], "int64_too", ["string_col"]), + ], +) +def test_df_pivot(scalars_dfs, values, index, columns): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df.pivot( + values=values, index=index, columns=columns + ).to_pandas() + pd_result = scalars_pandas_df.pivot(values=values, index=index, columns=columns) + + # Pandas produces NaN, where bq dataframes produces pd.NA + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("values", "index", "columns"), + [ + (["goals", "assists"], ["team_name", "season"], ["position"]), + (["goals", "assists"], ["season"], ["team_name", "position"]), + ], +) +def test_df_pivot_hockey(hockey_df, hockey_pandas_df, values, index, columns): + bf_result = hockey_df.pivot(values=values, index=index, columns=columns).to_pandas() + pd_result = hockey_pandas_df.pivot(values=values, index=index, columns=columns) + + # Pandas produces NaN, where bq dataframes produces pd.NA + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + def test_ipython_key_completions_with_drop(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_names = "string_col" @@ -1621,6 +1664,7 @@ def test_sample_raises_value_error(scalars_dfs): @pytest.mark.parametrize( ("axis",), [ + (None,), (0,), (1,), ], diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index d5dd4e357bc..5a2562bfb20 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -109,8 +109,6 @@ def test_dataframe_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): ) bf_result_computed = bf_result.to_pandas() - # Pandas produces multi-index which isn't supported in bq df yet - pd_result = pd_result.set_axis(bf_result.columns, axis=1) pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) @@ -128,8 +126,6 @@ def test_dataframe_groupby_agg_dict(scalars_df_index, scalars_pandas_df_index): ) bf_result_computed = bf_result.to_pandas() - # Pandas produces multi-index which isn't supported in bq df yet - pd_result = pd_result.set_axis(bf_result.columns, axis=1) pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) diff --git a/tests/system/small/test_ipython.py b/tests/system/small/test_ipython.py new file mode 100644 index 00000000000..6725805d9a7 --- /dev/null +++ b/tests/system/small/test_ipython.py @@ -0,0 +1,28 @@ +# 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 + +IPython = pytest.importorskip("IPython") + + +def test_repr_cache(scalars_df_index): + display_formatter = IPython.core.formatters.DisplayFormatter() + # Make sure the df has a new block that the method return value + # is not already cached. + test_df = scalars_df_index.head() + results = display_formatter.format(test_df) + assert results[0].keys() == {"text/plain", "text/html"} + assert test_df._block.retrieve_repr_request_results.cache_info().misses == 1 + assert test_df._block.retrieve_repr_request_results.cache_info().hits == 1 diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 914be6dae4a..b2937d7da9c 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -19,6 +19,7 @@ from tests.system.utils import assert_pandas_df_equal_ignore_ordering +# Row Multi-index tests def test_set_multi_index(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.set_index(["bool_col", "int64_too"]).to_pandas() pd_result = scalars_pandas_df_index.set_index(["bool_col", "int64_too"]) @@ -443,3 +444,179 @@ def test_multi_index_series_rename_dict_same_type( pandas.testing.assert_series_equal( bf_result, pd_result, check_dtype=False, check_index_type=False ) + + +# Column Multi-index tests + + +def test_column_multi_index_getitem(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], columns)) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_a = bf_df["a"].to_pandas() + pd_a = pd_df["a"] + pandas.testing.assert_frame_equal(bf_a, pd_a) + + bf_b = bf_df["b"].to_pandas() + pd_b = pd_df["b"] + pandas.testing.assert_frame_equal(bf_b, pd_b) + + bf_fullkey = bf_df[("a", "int64_too")].to_pandas() + pd_fullkey = pd_df[("a", "int64_too")] + pandas.testing.assert_series_equal(bf_fullkey, pd_fullkey) + + +def test_column_multi_index_concat(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col", "int64_col"] + multi_columns1 = pandas.MultiIndex.from_tuples( + zip(["a", "b", "a", "b"], [1, 1, 2, 2]) + ) + multi_columns2 = pandas.MultiIndex.from_tuples( + zip(["a", "b", "a", "c"], [3, 1, 2, 1]) + ) + + bf_df1 = scalars_df_index[columns].copy() + bf_df1.columns = multi_columns1 + bf_df2 = scalars_df_index[columns].copy() + bf_df2.columns = multi_columns2 + + pd_df1 = scalars_pandas_df_index[columns].copy() + pd_df1.columns = multi_columns1 + pd_df2 = scalars_pandas_df_index[columns].copy() + pd_df2.columns = multi_columns2 + + bf_result = bpd.concat([bf_df1, bf_df2, bf_df1]).to_pandas() + pd_result = pandas.concat([pd_df1, pd_df2, pd_df1]) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_column_multi_index_drop(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], columns)) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_a = bf_df.drop(("a", "int64_too"), axis=1).to_pandas() + pd_a = pd_df.drop(("a", "int64_too"), axis=1) + pandas.testing.assert_frame_equal(bf_a, pd_a) + + +@pytest.mark.parametrize( + ("key",), + [ + ("a",), + ("b",), + ("c",), + ], +) +def test_column_multi_index_assign(scalars_df_index, scalars_pandas_df_index, key): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], columns)) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + kwargs = {key: 42} + bf_result = bf_df.assign(**kwargs).to_pandas() + pd_result = pd_df.assign(**kwargs) + + # Pandas assign results in non-nullable dtype + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +def test_column_multi_index_rename(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.rename(columns={"b": "c"}).to_pandas() + pd_result = pd_df.rename(columns={"b": "c"}) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_column_multi_index_reset_index(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.reset_index().to_pandas() + pd_result = pd_df.reset_index() + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pandas.Int64Dtype()) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_column_multi_index_binary_op(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = (bf_df[("a", "a")] + 3).to_pandas() + pd_result = pd_df[("a", "a")] + 3 + + pandas.testing.assert_series_equal(bf_result, pd_result) + + +def test_column_multi_index_agg(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.agg(["sum", "mean"]).to_pandas() + pd_result = pd_df.agg(["sum", "mean"]) + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + pandas.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) + + +def test_column_multi_index_prefix_suffix(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.add_prefix("prefixed_").add_suffix("_suffixed").to_pandas() + pd_result = pd_df.add_prefix("prefixed_").add_suffix("_suffixed") + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_column_multi_index_cumsum(scalars_df_index, scalars_pandas_df_index): + if pandas.__version__.startswith("1."): + pytest.skip("pandas 1.x. does not handle nullable ints properly in cumsum") + columns = ["int64_too", "int64_col", "float64_col"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.cumsum().to_pandas() + pd_result = pd_df.cumsum() + + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index 96697dbcabe..9a3d55aed23 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -20,6 +20,7 @@ import google.auth.exceptions import pytest +import bigframes.core.global_session import bigframes.pandas as bpd @@ -289,7 +290,7 @@ def test_reset_session_after_bq_session_ended(): # Now try to reset session and verify that it works bpd.reset_session() - assert bpd._global_session is None + assert bigframes.core.global_session._global_session is None # Now verify that use is able to start over df = bpd.read_gbq(test_query) @@ -332,7 +333,7 @@ def test_reset_session_after_credentials_need_reauthentication(monkeypatch): # Now verify that resetting the session works bpd.reset_session() - assert bpd._global_session is None + assert bigframes.core.global_session._global_session is None # Now verify that use is able to start over df = bpd.read_gbq(test_query) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 47d758763ba..e40addc4ebd 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from google.cloud import bigquery +from ibis.backends.bigquery import datatypes as bq_types +from ibis.expr import datatypes as ibis_types import pandas as pd import pytest import bigframes -from bigframes.remote_function import read_gbq_function, remote_function +from bigframes import remote_function as rf from tests.system.utils import assert_pandas_df_equal_ignore_ordering @@ -98,6 +101,16 @@ def session_with_bq_connection_location_project_specified( ) +def test_supported_types_correspond(): + # The same types should be representable by the supported Python and BigQuery types. + ibis_types_from_python = {ibis_types.dtype(t) for t in rf.SUPPORTED_IO_PYTHON_TYPES} + ibis_types_from_bigquery = { + bq_types.BigQueryType.to_ibis(tk) for tk in rf.SUPPORTED_IO_BIGQUERY_TYPEKINDS + } + + assert ibis_types_from_python == ibis_types_from_bigquery + + @pytest.mark.flaky(retries=2, delay=120) def test_remote_function_direct_no_session_param( bigquery_client, @@ -107,7 +120,7 @@ def test_remote_function_direct_no_session_param( dataset_id_permanent, bq_cf_connection, ): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -157,7 +170,7 @@ def test_remote_function_direct_no_session_param_location_specified( dataset_id_permanent, bq_cf_connection_location, ): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -205,7 +218,7 @@ def test_remote_function_direct_no_session_param_location_mismatched( ): with pytest.raises(ValueError): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -229,7 +242,7 @@ def test_remote_function_direct_no_session_param_location_project_specified( dataset_id_permanent, bq_cf_connection_location_project, ): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -277,7 +290,7 @@ def test_remote_function_direct_no_session_param_project_mismatched( ): with pytest.raises(ValueError): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -294,7 +307,7 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) def test_remote_function_direct_session_param(session_with_bq_connection, scalars_dfs): - @remote_function( + @rf.remote_function( [int], int, session=session_with_bq_connection, @@ -500,6 +513,18 @@ def add_one(x): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.flaky(retries=2, delay=120) +def test_read_gbq_function_detects_invalid_function(bigquery_client, dataset_id): + dataset_ref = bigquery.DatasetReference.from_string(dataset_id) + with pytest.raises(ValueError) as e: + rf.read_gbq_function( + str(dataset_ref.routine("not_a_function")), + bigquery_client=bigquery_client, + ) + + assert "Unknown function" in str(e.value) + + @pytest.mark.flaky(retries=2, delay=120) def test_read_gbq_function_like_original( bigquery_client, @@ -509,7 +534,7 @@ def test_read_gbq_function_like_original( dataset_id_permanent, bq_cf_connection, ): - @remote_function( + @rf.remote_function( [int], int, bigquery_client=bigquery_client, @@ -522,7 +547,7 @@ def test_read_gbq_function_like_original( def square1(x): return x * x - square2 = read_gbq_function( + square2 = rf.read_gbq_function( function_name=square1.bigframes_remote_function, bigquery_client=bigquery_client, ) @@ -551,3 +576,111 @@ def square1(x): s2_result = int64_col_filtered.to_frame().assign(result=s2_result_col) assert_pandas_df_equal_ignore_ordering(s1_result.to_pandas(), s2_result.to_pandas()) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_read_gbq_function_reads_udfs(bigquery_client, scalars_dfs, dataset_id): + dataset_ref = bigquery.DatasetReference.from_string(dataset_id) + arg = bigquery.RoutineArgument( + name="x", + data_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + ) + sql_routine = bigquery.Routine( + dataset_ref.routine("square_sql"), + body="x * x", + arguments=[arg], + return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + js_routine = bigquery.Routine( + dataset_ref.routine("square_js"), + body="return x * x", + language="JAVASCRIPT", + arguments=[arg], + return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + + for routine in (sql_routine, js_routine): + # Create the routine in BigQuery and read it back using read_gbq_function. + bigquery_client.create_routine(routine, exists_ok=True) + square = rf.read_gbq_function( + str(routine.reference), bigquery_client=bigquery_client + ) + + # It should point to the named routine and yield the expected results. + assert square.bigframes_remote_function == str(routine.reference) + + src = {"x": [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5]} + + routine_ref_str = rf.routine_ref_to_string_for_query(routine.reference) + direct_sql = " UNION ALL ".join( + [f"SELECT {x} AS x, {routine_ref_str}({x}) AS y" for x in src["x"]] + ) + direct_df = bigquery_client.query(direct_sql).to_dataframe() + + indirect_df = bigframes.dataframe.DataFrame(src) + indirect_df = indirect_df.assign(y=indirect_df.x.apply(square)) + indirect_df = indirect_df.to_pandas() + + assert_pandas_df_equal_ignore_ordering(direct_df, indirect_df) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_read_gbq_function_enforces_explicit_types(bigquery_client, dataset_id): + dataset_ref = bigquery.DatasetReference.from_string(dataset_id) + typed_arg = bigquery.RoutineArgument( + name="x", + data_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + ) + untyped_arg = bigquery.RoutineArgument( + name="x", + kind="ANY_TYPE", # With this kind, data_type not required for SQL functions. + ) + + both_types_specified = bigquery.Routine( + dataset_ref.routine("both_types_specified"), + body="x * x", + arguments=[typed_arg], + return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + only_return_type_specified = bigquery.Routine( + dataset_ref.routine("only_return_type_specified"), + body="x * x", + arguments=[untyped_arg], + return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + only_arg_type_specified = bigquery.Routine( + dataset_ref.routine("only_arg_type_specified"), + body="x * x", + arguments=[typed_arg], + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + neither_type_specified = bigquery.Routine( + dataset_ref.routine("neither_type_specified"), + body="x * x", + arguments=[untyped_arg], + type_=bigquery.RoutineType.SCALAR_FUNCTION, + ) + + bigquery_client.create_routine(both_types_specified, exists_ok=True) + bigquery_client.create_routine(only_return_type_specified, exists_ok=True) + bigquery_client.create_routine(only_arg_type_specified, exists_ok=True) + bigquery_client.create_routine(neither_type_specified, exists_ok=True) + + rf.read_gbq_function( + str(both_types_specified.reference), bigquery_client=bigquery_client + ) + rf.read_gbq_function( + str(only_return_type_specified.reference), bigquery_client=bigquery_client + ) + with pytest.raises(ValueError): + rf.read_gbq_function( + str(only_arg_type_specified.reference), bigquery_client=bigquery_client + ) + with pytest.raises(ValueError): + rf.read_gbq_function( + str(neither_type_specified.reference), bigquery_client=bigquery_client + ) diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 2fc34f9bae9..d825c625617 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import io import random import tempfile import textwrap @@ -683,6 +684,43 @@ def test_read_csv_local_w_encoding(session, penguins_pandas_df_default_index, en assert df.shape[0] == penguins_pandas_df_default_index.shape[0] +def test_read_pickle_local(session, penguins_pandas_df_default_index, tmp_path): + path = tmp_path / "test_read_csv_local_w_encoding.pkl" + + penguins_pandas_df_default_index.to_pickle(path) + df = session.read_pickle(path) + + pd.testing.assert_frame_equal(penguins_pandas_df_default_index, df.to_pandas()) + + +def test_read_pickle_buffer(session, penguins_pandas_df_default_index): + buffer = io.BytesIO() + penguins_pandas_df_default_index.to_pickle(buffer) + buffer.seek(0) + df = session.read_pickle(buffer) + + pd.testing.assert_frame_equal(penguins_pandas_df_default_index, df.to_pandas()) + + +def test_read_pickle_series_buffer(session): + pd_series = pd.Series([1, 2, 3, 4, 5], dtype="Int64") + buffer = io.BytesIO() + pd_series.to_pickle(buffer) + buffer.seek(0) + bf_series = session.read_pickle(buffer).to_pandas() + pd_series.index = pd_series.index.astype("Int64") + + assert (pd_series == bf_series).all() + + +def test_read_pickle_gcs(session, penguins_pandas_df_default_index, gcs_folder): + path = gcs_folder + "test_read_pickle_gcs.pkl" + penguins_pandas_df_default_index.to_pickle(path) + df = session.read_pickle(path) + + pd.testing.assert_frame_equal(penguins_pandas_df_default_index, df.to_pandas()) + + def test_read_parquet_gcs(session: bigframes.Session, scalars_dfs, gcs_folder): scalars_df, _ = scalars_dfs # Include wildcard so that multiple files can be written/read if > 1 GB. diff --git a/tests/unit/test_dtypes.py b/tests/unit/test_dtypes.py index dafed08980c..bb8ae570dca 100644 --- a/tests/unit/test_dtypes.py +++ b/tests/unit/test_dtypes.py @@ -175,9 +175,12 @@ def test_literal_to_ibis_scalar_throws_on_incompatible_literal(): def test_remote_function_io_types_are_supported_bigframes_types(): + from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type + from bigframes.remote_function import ( - _supported_io_ibis_types as rf_supported_io_ibis_types, + SUPPORTED_IO_PYTHON_TYPES as rf_supported_io_types, ) - for ibis_type in rf_supported_io_ibis_types: + for python_type in rf_supported_io_types: + ibis_type = python_type_to_bigquery_type(python_type) assert ibis_type in bigframes.dtypes.IBIS_TO_BIGFRAMES diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 653b65c8346..6762afc61fd 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -44,6 +44,7 @@ def axes(self) -> list: They are returned in that order. Examples + .. code-block:: df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) @@ -175,8 +176,8 @@ def assign(self, **kwargs) -> DataFrame: are simply assigned to the column. Returns: - DataFrame: A new DataFrame with the new columns in addition to - all the existing columns. + bigframes.dataframe.DataFrame: A new DataFrame with the new columns + in addition to all the existing columns. """ raise NotImplementedError("abstract method") @@ -205,7 +206,7 @@ def drop( level: For MultiIndex, level from which the labels will be removed. Returns: - DataFrame: DataFrame without the removed column labels. + bigframes.dataframe.DataFrame: DataFrame without the removed column labels. Raises: KeyError: If any of the labels is not found in the selected axis. @@ -227,7 +228,7 @@ def rename( Dict-like from old column labels to new column labels. Returns: - DataFrame: DataFrame with the renamed axis labels. + bigframes.dataframe.DataFrame: DataFrame with the renamed axis labels. Raises: KeyError: If any of the labels is not found. @@ -238,7 +239,7 @@ def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: """ Set the name of the axis for the index. - .. Note:: + .. note:: Currently only accepts a single string parameter (the new name of the index). @@ -247,7 +248,7 @@ def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: Value to set the axis name attribute. Returns: - DataFrame: DataFrame with the new index name + bigframes.dataframe.DataFrame: DataFrame with the new index name """ raise NotImplementedError("abstract method") @@ -317,7 +318,7 @@ def reset_index( the index to the default integer index. Returns: - DataFrame: DataFrame with the new index. + bigframes.dataframe.DataFrame: DataFrame with the new index. """ raise NotImplementedError("abstract method") @@ -344,7 +345,7 @@ def drop_duplicates( - ``False`` : Drop all duplicates. Returns: - DataFrame: DataFrame with duplicates removed + bigframes.dataframe.DataFrame: DataFrame with duplicates removed """ raise NotImplementedError("abstract method") @@ -366,7 +367,7 @@ def duplicated(self, subset=None, keep="first"): - False : Mark all duplicates as ``True``. Returns: - Boolean series for each duplicated rows. + bigframes.series.Series: Boolean series for each duplicated rows. """ raise NotImplementedError("abstract method") @@ -379,7 +380,7 @@ def dropna( """Remove missing values. Returns: - DataFrame: DataFrame with NA entries dropped from it. + bigframes.dataframe.DataFrame: DataFrame with NA entries dropped from it. """ raise NotImplementedError("abstract method") @@ -844,7 +845,7 @@ def groupby( values will also be treated as the key in groups. Returns: - A groupby object that contains information about the groups. + bigframes.core.groupby.SeriesGroupBy: A groupby object that contains information about the groups. """ raise NotImplementedError("abstract method") @@ -871,7 +872,7 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: values, without passing them to func. Returns: - DataFrame: Transformed DataFrame. + bigframes.dataframe.DataFrame: Transformed DataFrame. """ raise NotImplementedError("abstract method") @@ -899,7 +900,7 @@ def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame: index, preserving the order of the calling's one. Returns: - DataFrame: A dataframe containing columns from both the caller and `other`. + bigframes.dataframe.DataFrame: A dataframe containing columns from both the caller and `other`. """ raise NotImplementedError("abstract method") @@ -969,7 +970,7 @@ def merge( no suffix. At least one of the values must not be None. Returns: - DataFrame: A DataFrame of the two merged objects. + bigframes.dataframe.DataFrame: A DataFrame of the two merged objects. """ raise NotImplementedError("abstract method") @@ -1006,7 +1007,7 @@ def all(self, *, bool_only: bool = False): Include only boolean columns. Returns: - Series + bigframes.series.Series: Series if all elements are True. """ raise NotImplementedError("abstract method") @@ -1019,7 +1020,7 @@ def prod(self, *, numeric_only: bool = False): Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with the product of the values. """ raise NotImplementedError("abstract method") @@ -1034,7 +1035,7 @@ def min(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with the minimum of the values. """ raise NotImplementedError("abstract method") @@ -1049,7 +1050,7 @@ def max(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series after the maximum of values. """ raise NotImplementedError("abstract method") @@ -1063,7 +1064,7 @@ def sum(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with the sum of values. """ raise NotImplementedError("abstract method") @@ -1075,7 +1076,7 @@ def mean(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with the mean of values. """ raise NotImplementedError("abstract method") @@ -1090,7 +1091,7 @@ def median(self, *, numeric_only: bool = False, exact: bool = False): one. Note: ``exact=True`` not yet supported. Returns: - Series + bigframes.series.Series: Series with the median of values. """ raise NotImplementedError("abstract method") @@ -1104,7 +1105,7 @@ def var(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with unbiased variance over requested axis. """ raise NotImplementedError("abstract method") @@ -1118,7 +1119,7 @@ def std(self, *, numeric_only: bool = False): Default False. Include only float, int, boolean columns. Returns: - Series + bigframes.series.Series: Series with sample standard deviation. """ raise NotImplementedError("abstract method") @@ -1134,8 +1135,8 @@ def count(self, *, numeric_only: bool = False): Include only `float`, `int` or `boolean` data. Returns: - For each column/row the number of non-NA/null entries. - If `level` is specified returns a `DataFrame`. + bigframes.series.Series: For each column/row the number of + non-NA/null entries. If `level` is specified returns a `DataFrame`. """ raise NotImplementedError("abstract method") @@ -1143,10 +1144,8 @@ def nunique(self): """ Count number of distinct elements in specified axis. - Return Series with number of distinct elements. - Returns: - Series + bigframes.series.Series: Series with number of distinct elements. """ raise NotImplementedError("abstract method") @@ -1156,7 +1155,7 @@ def cummin(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative minimum. Returns: - DataFrame: Return cumulative minimum of DataFrame. + bigframes.dataframe.DataFrame: Return cumulative minimum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1166,7 +1165,7 @@ def cummax(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative maximum. Returns: - DataFrame: Return cumulative maximum of DataFrame. + bigframes.dataframe.DataFrame: Return cumulative maximum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1176,7 +1175,7 @@ def cumsum(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative sum. Returns: - DataFrame: Return cumulative sum of DataFrame. + bigframes.dataframe.DataFrame: Return cumulative sum of DataFrame. """ raise NotImplementedError("abstract method") @@ -1186,7 +1185,7 @@ def cumprod(self) -> DataFrame: Returns a DataFrame of the same size containing the cumulative product. Returns: - DataFrame: Return cumulative product of DataFrame. + bigframes.dataframe.DataFrame: Return cumulative product of DataFrame. """ raise NotImplementedError("abstract method") @@ -1201,7 +1200,7 @@ def agg(self, func): function names, e.g. ``['sum', 'mean']``. Returns: - Series or DataFrame: Aggregated results + DataFrame or bigframes.series.Series: Aggregated results. """ raise NotImplementedError("abstract method") @@ -1218,17 +1217,51 @@ def describe(self): .. note:: Percentile values are approximates only. + .. note:: + For numeric data, the result's index will include ``count``, + ``mean``, ``std``, ``min``, ``max`` as well as lower, ``50`` and + upper percentiles. By default the lower percentile is ``25`` and the + upper percentile is ``75``. The ``50`` percentile is the + same as the median. + Returns: - Summary statistics of the Series or Dataframe provided. + bigframes.dataframe.DataFrame: Summary statistics of the Series or Dataframe provided. + """ + raise NotImplementedError("abstract method") + + def pivot(self, *, columns, index=None, values=None): + """ + Return reshaped DataFrame organized by given index / column values. + + Reshape data (produce a "pivot" table) based on column values. Uses + unique values from specified `index` / `columns` to form axes of the + resulting DataFrame. This function does not support data + aggregation, multiple values will result in a MultiIndex in the + columns. + + .. note:: + BigQuery supports up to 10000 columns. Pivot operations on columns + with too many unique values will fail if they would exceed this limit. + + .. note:: + The validity of the pivot operation is not checked. If columns and index + do not together uniquely identify input rows, the output will be + silently non-deterministic. + Args: + columns (str or object or a list of str): + Column to use to make new frame's columns. - Notes - ----- - For numeric data, the result's index will include ``count``, - ``mean``, ``std``, ``min``, ``max`` as well as lower, ``50`` and - upper percentiles. By default the lower percentile is ``25`` and the - upper percentile is ``75``. The ``50`` percentile is the - same as the median. + index (str or object or a list of str, optional): + Column to use to make new frame's index. If not given, uses existing index. + + values (str, object or a list of the previous, optional): + Column(s) to use for populating new frame's values. If not + specified, all remaining columns will be used and the result will + have hierarchically indexed columns. + + Returns: + Returns reshaped DataFrame. """ raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 039dc1eae0f..79eb4026968 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -211,7 +211,7 @@ def to_frame(self) -> DataFrame: Convert Series to DataFrame. Returns: - DataFrame: DataFrame representation of Series. + bigframes.dataframe.DataFrame: DataFrame representation of Series. """ raise NotImplementedError("abstract method") @@ -419,7 +419,7 @@ def mode(self) -> Series: Always returns Series even if only one value is returned. Returns: - Series: Modes of the Series in sorted order. + bigframes.series.Series: Modes of the Series in sorted order. """ raise NotImplementedError("abstract method") @@ -440,7 +440,7 @@ def drop_duplicates( ``False`` : Drop all duplicates. Returns: - Series: Series with duplicates dropped or None if ``inplace=True``. + bigframes.series.Series: Series with duplicates dropped or None if ``inplace=True``. """ raise NotImplementedError("abstract method") @@ -463,8 +463,8 @@ def duplicated(self, keep="first") -> Series: ``False`` : Mark all duplicates as ``True``. Returns: - Series: Series indicating whether each value has occurred in the - preceding values. + bigframes.series.Series: Series indicating whether each value has occurred in the + preceding values. """ raise NotImplementedError("abstract method") @@ -478,7 +478,7 @@ def round(self, decimals: int = 0) -> Series: it specifies the number of positions to the left of the decimal point. Returns: - Series: Rounded values of the Series. + bigframes.series.Series: Rounded values of the Series. """ raise NotImplementedError("abstract method") @@ -569,7 +569,7 @@ def sort_values( the end. Returns: - Series or None: Series ordered by values or None if ``inplace=True``. + bigframes.series.Series: Series ordered by values or None if ``inplace=True``. """ raise NotImplementedError("abstract method") @@ -597,7 +597,7 @@ def sort_index( Not implemented for MultiIndex. Returns: - Series or None: The original Series sorted by the labels or None if + bigframes.series.Series: The original Series sorted by the labels or None if ``inplace=True``. """ @@ -624,7 +624,7 @@ def nlargest( size larger than `n`. Returns: - Series: The `n` largest values in the Series, sorted in decreasing order. + bigframes.series.Series: The `n` largest values in the Series, sorted in decreasing order. """ raise NotImplementedError("abstract method") @@ -647,7 +647,7 @@ def nsmallest(self, n: int = 5, keep: str = "first") -> Series: size larger than `n`. Returns: - Series: The `n` smallest values in the Series, sorted in increasing order. + bigframes.series.Series: The `n` smallest values in the Series, sorted in increasing order. """ raise NotImplementedError("abstract method") @@ -669,7 +669,7 @@ def apply( Python function or NumPy ufunc to apply. Returns: - Series or DataFrame: If func returns a Series object the result + bigframes.series.Series: If func returns a Series object the result will be a DataFrame. """ raise NotImplementedError("abstract method") @@ -723,7 +723,8 @@ def groupby( If False, NA values will also be treated as the key in groups. Returns: - SeriesGroupBy: Returns a groupby object that contains information about the groups. + bigframes.core.groupby.SeriesGroupBy: Returns a groupby object that contains + information about the groups. """ raise NotImplementedError("abstract method") @@ -750,15 +751,12 @@ def drop( level: For MultiIndex, level for which the labels will be removed. - Returns - ------- - Series or None - Series with specified index labels removed or None if ``inplace=True``. + Returns: + bigframes.series.Series: Series with specified index labels removed + or None if ``inplace=True``. - Raises - ------ - KeyError - If none of the labels are found in the index. + Raises: + KeyError: If none of the labels are found in the index. """ raise NotImplementedError("abstract method") @@ -844,7 +842,7 @@ def cumprod(self): product. Returns: - Return cumulative sum of scalar or Series. + bigframes.series.Series: Return cumulative sum of scalar or Series. """ raise NotImplementedError("abstract method") @@ -878,7 +876,7 @@ def cummax(self): For `Series` this parameter is unused and defaults to 0. Returns: - scalar or Series: Return cumulative maximum of scalar or Series. + bigframes.series.Series: Return cumulative maximum of scalar or Series. """ raise NotImplementedError("abstract method") @@ -901,7 +899,7 @@ def cummin(self): compatibility with NumPy. Returns: - scalar or Series: Return cumulative minimum of scalar or Series. + bigframes.series.Series: Return cumulative minimum of scalar or Series. """ raise NotImplementedError("abstract method") @@ -930,7 +928,7 @@ def ne(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -945,7 +943,7 @@ def le(self, other) -> Series: other: Series, or scalar value Returns: - Series. The result of the comparison. + bigframes.series.Series. The result of the comparison. """ raise NotImplementedError("abstract method") @@ -960,7 +958,7 @@ def lt(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -975,7 +973,7 @@ def ge(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -990,7 +988,7 @@ def gt(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1005,7 +1003,7 @@ def add(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1020,7 +1018,7 @@ def radd(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1038,7 +1036,7 @@ def sub( other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1053,7 +1051,7 @@ def rsub(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1068,7 +1066,7 @@ def mul(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1097,7 +1095,7 @@ def truediv(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1112,7 +1110,7 @@ def rtruediv(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1127,7 +1125,7 @@ def floordiv(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1142,7 +1140,7 @@ def rfloordiv(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1157,7 +1155,7 @@ def mod(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1172,7 +1170,7 @@ def rmod(self, other) -> Series: other (Series, or scalar value): Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError("abstract method") @@ -1373,7 +1371,7 @@ def where(self, cond, other): extension dtypes). Returns: - Series + bigframes.series.Series: Series after the replacement. """ raise NotImplementedError("abstract method") @@ -1397,7 +1395,7 @@ def mask(self, cond, other): extension dtypes). Returns: - Series + bigframes.series.Series: Series after the replacement. """ raise NotImplementedError("abstract method") @@ -1460,7 +1458,7 @@ def rename(self, index, **kwargs) -> Series | None: attribute. Returns: - Series: Series with index labels + bigframes.series.Series: Series with index labels. """ raise NotImplementedError("abstract method") @@ -1474,7 +1472,7 @@ def rename_axis(self, mapper, **kwargs): Value to set the axis name attribute. Returns: - Series: Series with the name of the axis set. + bigframes.series.Series: Series with the name of the axis set. """ raise NotImplementedError("abstract method") @@ -1514,7 +1512,8 @@ def rolling( to the size of the window. Returns: - ``Window`` subclass if a ``win_type`` is passed.``Rolling`` subclass if ``win_type`` is not passed + bigframes.core.window.Window: ``Window`` subclass if a ``win_type`` is passed. + ``Rolling`` subclass if ``win_type`` is not passed. """ raise NotImplementedError("abstract method") @@ -1528,7 +1527,7 @@ def expanding(self, min_periods=1): otherwise, result is ``np.nan``. Returns: - ``Expanding`` subclass + bigframes.core.window.Window: ``Expanding`` subclass. """ raise NotImplementedError("abstract method") @@ -1591,7 +1590,7 @@ def isin(self, values): TypeError. Instead, turn a single string into a list of one element. Returns: - bigframes.series.Series: Series of booleans indicating if each element is in values. + bigframes.series.Series: Series of booleans indicating if each element is in values. Raises: TypeError: If input is not list-like. diff --git a/third_party/bigframes_vendored/pandas/core/strings/accessor.py b/third_party/bigframes_vendored/pandas/core/strings/accessor.py index a27093b552c..5e3d0b047ff 100644 --- a/third_party/bigframes_vendored/pandas/core/strings/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/strings/accessor.py @@ -136,6 +136,102 @@ def isnumeric(self): raise NotImplementedError("abstract method") + def isalpha(self): + """Check whether all characters in each string are alphabetic. + + This is equivalent to running the Python string method + :meth:`str.isalpha` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series with the same length as the originalSeries/Index. + """ + + raise NotImplementedError("abstract method") + + def isdigit(self): + """Check whether all characters in each string are digits. + + This is equivalent to running the Python string method + :meth:`str.isdigit` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series with the same length as the originalSeries/Index. + """ + + raise NotImplementedError("abstract method") + + def isalnum(self): + """Check whether all characters in each string are alphanumeric. + + This is equivalent to running the Python string method + :meth:`str.isalnum` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. + """ + + raise NotImplementedError("abstract method") + + def isspace(self): + """Check whether all characters in each string are whitespace. + + This is equivalent to running the Python string method + :meth:`str.isspace` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. + """ + + raise NotImplementedError("abstract method") + + def islower(self): + """Check whether all characters in each string are lowercase. + + This is equivalent to running the Python string method + :meth:`str.islower` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. + """ + + raise NotImplementedError("abstract method") + + def isupper(self): + """Check whether all characters in each string are uppercase. + + This is equivalent to running the Python string method + :meth:`str.isupper` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. + """ + + raise NotImplementedError("abstract method") + + def isdecimal(self): + """Check whether all characters in each string are decimal. + + This is equivalent to running the Python string method + :meth:`str.isdecimal` for each element of the Series/Index. If a string + has zero characters, ``False`` is returned for that check. + + Returns: + bigframes.series.Series: Series or Index of boolean values with the + same length as the original Series/Index. + """ + + raise NotImplementedError("abstract method") + def rstrip(self): """Remove trailing characters. @@ -427,3 +523,47 @@ def rjust( bigframes.series.Series: Returns Series or Index with minimum number of char in object. """ raise NotImplementedError("abstract method") + + def zfill( + self, + width: int, + ): + """ + Pad strings in the Series/Index by prepending '0' characters. + + Strings in the Series/Index are padded with '0' characters on the + left of the string to reach a total string length `width`. Strings + in the Series/Index with length greater or equal to `width` are + unchanged. + + Args: + width (int): + Minimum length of resulting string; strings with length less + than `width` be prepended with '0' characters. + + Returns: + bigframes.series.Series: Series of objects. + """ + raise NotImplementedError("abstract method") + + def center( + self, + width: int, + fillchar: str = " ", + ): + """ + Pad left and right side of strings in the Series/Index. + + Equivalent to :meth:`str.center`. + + Args: + width (int): + Minimum width of resulting string; additional characters will be filled + with character defined in `fillchar`. + fillchar (str, default ' '): + Additional character for filling, default is whitespace. + + Returns: + bigframes.series.Series: Returns Series or Index with minimum number of char in object. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 9425ead0e34..730872034dc 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -15,7 +15,56 @@ def read_gbq( col_order: Iterable[str] = (), max_results: Optional[int] = None, ): - """Loads DataFrame from BigQuery. + """Loads a DataFrame from BigQuery. + + BigQuery tables are an unordered, unindexed data source. By default, + the DataFrame will have an arbitrary index and ordering. + + Set the `index_col` argument to one or more columns to choose an + index. The resulting DataFrame is sorted by the index columns. For the + best performance, ensure the index columns don't contain duplicate + values. + + .. note:: + By default, even SQL query inputs with an ORDER BY clause create a + DataFrame with an arbitrary ordering. Use ``row_number() OVER + (ORDER BY ...) AS rowindex`` in your SQL query and set + ``index_col='rowindex'`` to preserve the desired ordering. + + If your query doesn't have an ordering, select ``GENERATE_UUID() AS + rowindex`` in your SQL and set ``index_col='rowindex'`` for the + best performance. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Preserve ordering in a query input. + + >>> bpd.read_gbq(''' + ... SELECT + ... -- Instead of an ORDER BY clause on the query, use + ... -- ROW_NUMBER() to create an ordered DataFrame. + ... ROW_NUMBER() OVER (ORDER BY AVG(pitchSpeed) DESC) + ... AS rowindex, + ... + ... pitcherFirstName, + ... pitcherLastName, + ... AVG(pitchSpeed) AS averagePitchSpeed + ... FROM `bigquery-public-data.baseball.games_wide` + ... WHERE year = 2016 + ... GROUP BY pitcherFirstName, pitcherLastName + ... ''', index_col="rowindex").head(n=5) + pitcherFirstName pitcherLastName averagePitchSpeed + rowindex + 1 Albertin Chapman 96.514113 + 2 Zachary Britton 94.591039 + 3 Trevor Rosenthal 94.213953 + 4 Jose Torres 94.103448 + 5 Tayron Guerrero 93.863636 + + [5 rows x 3 columns] Args: query (str): diff --git a/third_party/bigframes_vendored/pandas/io/pickle.py b/third_party/bigframes_vendored/pandas/io/pickle.py new file mode 100644 index 00000000000..a160ef0c4ea --- /dev/null +++ b/third_party/bigframes_vendored/pandas/io/pickle.py @@ -0,0 +1,55 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/io/pickle.py +""" pickle compat """ +from __future__ import annotations + +from pandas._typing import ( + CompressionOptions, + FilePath, + ReadPickleBuffer, + StorageOptions, +) + + +class PickleIOMixin: + def read_pickle( + self, + filepath_or_buffer: FilePath | ReadPickleBuffer, + compression: CompressionOptions = "infer", + storage_options: StorageOptions = None, + ): + """Load pickled BigFrames object (or any object) from file. + + .. note:: + If the content of the pickle file is a Series and its name attribute is None, + the name will be set to '0' by default. + + Args: + filepath_or_buffer (str, path object, or file-like object): + String, path object (implementing os.PathLike[str]), or file-like object + implementing a binary readlines() function. Also accepts URL. URL is not + limited to S3 and GCS. + compression (str or dict, default 'infer'): + For on-the-fly decompression of on-disk data. If 'infer' and + 'filepath_or_buffer' is path-like, then detect compression from the following + extensions: '.gz', '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' + or '.tar.bz2' (otherwise no compression). If using 'zip' or 'tar', the ZIP + file must contain only one data file to be read in. Set to None for no + decompression. Can also be a dict with key 'method' set to one of {'zip', + 'gzip', 'bz2', 'zstd', 'tar'} and other key-value pairs are forwarded to + zipfile.ZipFile, gzip.GzipFile, bz2.BZ2File, zstandard.ZstdDecompressor or + tarfile.TarFile, respectively. As an example, the following could be passed + for Zstandard decompression using a custom compression dictionary + compression={'method': 'zstd', 'dict_data': my_compression_dict}. + storage_options (dict, default None): + Extra options that make sense for a particular storage connection, e.g. host, + port, username, password, etc. For HTTP(S) URLs the key-value pairs are + forwarded to urllib.request.Request as header options. For other URLs (e.g. + starting with “s3://”, and “gcs://”) the key-value pairs are forwarded to + fsspec.open. Please see fsspec and urllib for more details, and for more + examples on storage options refer here. + + Returns: + bigframes.dataframe.DataFrame or bigframes.series.Series: same type as object + stored in file. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py index 03958f75950..fc48cde85ba 100644 --- a/third_party/bigframes_vendored/sklearn/base.py +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -113,7 +113,7 @@ def score(self, X, y): ``(n_samples, n_samples_fitted)``, where ``n_samples_fitted`` is the number of samples used in the fitting for the estimator. - y (bigframes.dataframe.DataFrame or bigframes.series.Series: + y (bigframes.dataframe.DataFrame or bigframes.series.Series): Series or DataFrame of shape (n_samples,) or (n_samples, n_outputs). True values for `X`. @@ -123,6 +123,27 @@ def score(self, X, y): raise NotImplementedError("abstract method") +class TransformerMixin: + """Mixin class for all transformers.""" + + def fit_transform(self, X, y=None): + """Fit to data, then transform it. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples, n_features). + Input samples. + + y (bigframes.dataframe.DataFrame or bigframes.series.Series): + Series or DataFrame of shape (n_samples,) or (n_samples, n_outputs). Default None. + Target values (None for unsupervised transformations). + + Returns: + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_features_new) + Transformed DataFrame. + """ + + class MetaEstimatorMixin: _required_parameters = ["estimator"] """Mixin class for all meta estimators in scikit-learn.""" diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index 068aa4d2908..ea4df0dc023 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -76,8 +76,6 @@ def predict( Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples, n_features). New data to predict. - y: (default None) - Not used, present here for API consistency by convention. Returns: bigframes.dataframe.DataFrame: DataFrame of the cluster each sample belongs to. From 89b95033d6b449bfc21249057d7c024d096c80d0 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 29 Aug 2023 11:37:20 -0500 Subject: [PATCH 07/80] chore: sync changes from internal repo (#10) feat: support `Series.corr` fix: raise AttributeError for unimplemented pandas methods feat: support `DataFrame.stack` feat: support `np.arcsin`, `np.arccos`, `np.arctan`, `np.sinh`, `np.cosh`, `np.tanh`, `np.arcsinh`, `np.arccosh`, `np.arctanh`, `np.exp` with Series argument fix: align column names with pandas in `DataFrame.agg` results docs: set `options.bigquery.project` in sample code chore: unit test internal `get_standardized_ids` method fix: include survey link in abstract `NotImplementedError` exception messages perf: lazily instantiate client library objects fix: allow (but still not recommended) `ORDER BY` in `read_gbq` input when an `index_col` is defined feat: support `read_json` with `engine=bigquery` for newline-delimited JSON files chore: remove unneeded `types-retry` reference feat: support `np.sin`, `np.cos`, `np.tan`, `np.log`, `np.log10`, `np.sqrt`, `np.abs` with Series argument fix: label temp table creation jobs with `source=bigquery-dataframes-temp` label fix: support spaces in column names in `DataFrame` initializater chore: fix permissions on publish docs script feat: support `df[my_column] = [a python list]` feat: add `components_`, `explained_variance_`, and `explained_variance_ratio_` properties to `bigframes.ml.decomposition.PCA` chore: add execute permissions on publish docs script docs: fix link to GitHub chore: fix docs build fix: check for IAM role on the BigQuery connection when initializing a `remote_function` chore: revert pin to maximum pytest-retry plugin version in tests --- .kokoro/docs/common.cfg | 2 +- .kokoro/publish-docs.sh | 0 README.rst | 9 +- bigframes/constants.py | 2 + bigframes/core/__init__.py | 102 +- bigframes/core/block_transforms.py | 32 + bigframes/core/blocks.py | 159 ++- bigframes/core/groupby/__init__.py | 23 +- bigframes/core/indexes/__init__.py | 3 - bigframes/core/utils.py | 53 + bigframes/dataframe.py | 167 ++- bigframes/ml/core.py | 32 +- bigframes/ml/decomposition.py | 29 + bigframes/ml/sql.py | 10 + bigframes/operations/__init__.py | 163 +++ bigframes/operations/base.py | 5 + bigframes/pandas/__init__.py | 31 +- bigframes/remote_function.py | 109 +- bigframes/series.py | 55 +- bigframes/session.py | 467 ++++-- mypy.ini | 3 + notebooks/dataframes/dataframe.ipynb | 1249 +++++++++++++++-- .../bq_dataframes_ml_linear_regression.ipynb | 6 +- .../getting_started_bq_dataframes.ipynb | 6 +- noxfile.py | 8 +- setup.py | 2 + testing/constraints-3.9.txt | 2 + tests/system/conftest.py | 77 + tests/system/small/ml/conftest.py | 43 +- tests/system/small/ml/test_core.py | 94 ++ tests/system/small/ml/test_decomposition.py | 109 +- tests/system/small/test_dataframe.py | 80 +- tests/system/small/test_dataframe_io.py | 92 +- tests/system/small/test_groupby.py | 21 +- tests/system/small/test_multiindex.py | 36 + tests/system/small/test_numpy.py | 69 + tests/system/small/test_remote_function.py | 12 + tests/system/small/test_series.py | 20 +- tests/system/small/test_session.py | 125 +- tests/unit/core/test_utils.py | 56 + tests/unit/ml/test_sql.py | 15 + .../bigframes_vendored/pandas/core/frame.py | 144 +- .../bigframes_vendored/pandas/core/generic.py | 37 +- .../pandas/core/groupby/__init__.py | 46 +- .../pandas/core/indexes/accessor.py | 23 +- .../pandas/core/indexes/base.py | 8 +- .../pandas/core/indexing.py | 6 +- .../pandas/core/reshape/concat.py | 4 +- .../pandas/core/reshape/tile.py | 4 +- .../bigframes_vendored/pandas/core/series.py | 206 +-- .../pandas/core/strings/accessor.py | 66 +- .../pandas/core/window/rolling.py | 16 +- .../bigframes_vendored/pandas/io/gbq.py | 4 +- .../bigframes_vendored/pandas/io/parquet.py | 4 +- .../pandas/io/parsers/readers.py | 82 +- .../bigframes_vendored/pandas/io/pickle.py | 4 +- .../bigframes_vendored/sklearn/base.py | 6 +- .../sklearn/cluster/_kmeans.py | 11 +- .../sklearn/compose/_column_transformer.py | 5 +- .../sklearn/decomposition/_pca.py | 54 +- .../sklearn/ensemble/_forest.py | 8 +- .../sklearn/linear_model/_base.py | 7 +- .../sklearn/linear_model/_logistic.py | 3 +- .../sklearn/metrics/_classification.py | 12 +- .../sklearn/metrics/_ranking.py | 8 +- .../sklearn/metrics/_regression.py | 4 +- .../bigframes_vendored/sklearn/pipeline.py | 7 +- .../sklearn/preprocessing/_data.py | 5 +- .../sklearn/preprocessing/_encoder.py | 5 +- .../bigframes_vendored/xgboost/sklearn.py | 6 +- 70 files changed, 3599 insertions(+), 774 deletions(-) mode change 100644 => 100755 .kokoro/publish-docs.sh create mode 100644 tests/system/small/test_numpy.py create mode 100644 tests/unit/core/test_utils.py diff --git a/.kokoro/docs/common.cfg b/.kokoro/docs/common.cfg index ce84d7ec49a..bd73988540c 100644 --- a/.kokoro/docs/common.cfg +++ b/.kokoro/docs/common.cfg @@ -20,7 +20,7 @@ env_vars: { } env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: "git/bigframes/.kokoro/publish-docs.sh" + value: ".kokoro/publish-docs.sh" } env_vars: { diff --git a/.kokoro/publish-docs.sh b/.kokoro/publish-docs.sh old mode 100644 new mode 100755 diff --git a/README.rst b/README.rst index 6ae3753eed4..935c54cc8b5 100644 --- a/README.rst +++ b/README.rst @@ -41,6 +41,7 @@ method accepts either a fully-qualified table ID or a SQL query. import bigframes.pandas as bpd + bpd.options.bigquery.project = your_gcp_project_id df1 = bpd.read_gbq("project.dataset.table") df2 = bpd.read_gbq("SELECT a, b, c, FROM `project.dataset.table`") @@ -260,7 +261,7 @@ To view and manage Cloud Functions functions, use the `Functions `_ page and use the project picker to select the project in which you created the function. For easy identification, the names of the functions -created by BigQuery DataFrames are prefixed by ``bigframes-``. +created by BigQuery DataFrames are prefixed by ``bigframes``. **Requirements** @@ -283,7 +284,9 @@ following IAM roles: * BigQuery Data Editor (roles/bigquery.dataEditor) * BigQuery Connection Admin (roles/bigquery.connectionAdmin) * Cloud Functions Developer (roles/cloudfunctions.developer) -* Service Account User (roles/iam.serviceAccountUser) +* Service Account User (roles/iam.serviceAccountUser) on the + `service account ` + ``PROJECT_NUMBER-compute@developer.gserviceaccount.com`` * Storage Object Viewer (roles/storage.objectViewer) * Project IAM Admin (roles/resourcemanager.projectIamAdmin) @@ -330,7 +333,7 @@ Data processing location BigQuery DataFrames is designed for scale, which it achieves by keeping data and processing on the BigQuery service. However, you can bring data into the -memory of your client machine by calling ``.execute()`` on a DataFrame or Series +memory of your client machine by calling ``.to_pandas()`` on a DataFrame or Series object. If you choose to do this, the memory limitation of your client machine applies. diff --git a/bigframes/constants.py b/bigframes/constants.py index 3f3f1557336..90837c79eba 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -21,3 +21,5 @@ "Share your usecase with the BigQuery DataFrames team at the " "https://bit.ly/bigframes-feedback survey." ) + +ABSTRACT_METHOD_ERROR_MESSAGE = f"Abstract method. You have likely encountered a bug. Please share this stacktrace and how you reached it with the BigQuery DataFrames team. {FEEDBACK_LINK}" diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 9f392ce1492..d6509e4c0a9 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -35,6 +35,7 @@ reencode_order_string, StringEncoding, ) +import bigframes.core.utils as utils import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -562,6 +563,36 @@ def aggregate( ordering=ordering, ) + def corr_aggregate( + self, corr_aggregations: typing.Sequence[typing.Tuple[str, str, str]] + ) -> ArrayValue: + """ + Get correlations between each lef_column_id and right_column_id, stored in the respective output_column_id. + This uses BigQuery's CORR under the hood, and thus only Pearson's method is used. + Arguments: + corr_aggregations: left_column_id, right_column_id, output_column_id tuples + """ + table = self.to_ibis_expr(ordering_mode="unordered") + stats = { + col_out: table[col_left].corr(table[col_right], how="pop") + for col_left, col_right, col_out in corr_aggregations + } + aggregates = {**stats, ORDER_ID_COLUMN: ibis_types.literal(0)} + result = table.aggregate(**aggregates) + # Ordering is irrelevant for single-row output, but set ordering id regardless as other ops(join etc.) expect it. + ordering = ExpressionOrdering( + ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ) + return ArrayValue( + self._session, + result, + columns=[result[col_id] for col_id in [*stats.keys()]], + hidden_ordering_columns=[result[ORDER_ID_COLUMN]], + ordering=ordering, + ) + def project_window_op( self, column_name: str, @@ -852,38 +883,75 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal group_by=group_by, ) - def unpivot_single_row( + def unpivot( self, row_labels: typing.Sequence[typing.Hashable], - unpivot_columns: typing.Sequence[typing.Tuple[str, typing.Sequence[str]]], + unpivot_columns: typing.Sequence[ + typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + ], *, + passthrough_columns: typing.Sequence[str] = (), index_col_id: str = "index", - dtype=pandas.Float64Dtype(), + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + ] = pandas.Float64Dtype(), ) -> ArrayValue: - """Unpivot a single row.""" - # TODO: Generalize to multiple row input - table = self.to_ibis_expr(ordering_mode="unordered") + """ + Unpivot ArrayValue columns. + + Args: + row_labels: Identifies the source of the row. Must be equal to length to source column list in unpivot_columns argument. + unpivot_columns: Mapping of column id to list of input column ids. Lists of input columns may use None. + passthrough_columns: Columns that will not be unpivoted. Column id will be preserved. + index_col_id (str): The column id to be used for the row labels. + dtype (dtype or list of dtype): Dtype to use for the unpivot columns. If list, must be equal in number to unpivot_columns. + + Returns: + ArrayValue: The unpivoted ArrayValue + """ + table = self.to_ibis_expr(ordering_mode="offset_col") sub_expressions = [] - # TODO: validate all columns are equal length, as well as row labels + # Use ibis memtable to infer type of rowlabels (if possible) + # TODO: Allow caller to specify dtype + labels_ibis_type = ibis.memtable({"col": row_labels})["col"].type() + labels_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(labels_ibis_type) + row_n = len(row_labels) if not all( len(source_columns) == row_n for _, source_columns in unpivot_columns ): raise ValueError("Columns and row labels must all be same length.") - # Select each column for i in range(row_n): values = [] - for result_col, source_cols in unpivot_columns: - values.append( - ops.AsTypeOp(dtype)._as_ibis(table[source_cols[i]]).name(result_col) - ) - + for j in range(len(unpivot_columns)): + result_col, source_cols = unpivot_columns[j] + col_dtype = dtype[j] if utils.is_list_like(dtype) else dtype + if source_cols[i] is not None: + values.append( + ops.AsTypeOp(col_dtype) + ._as_ibis(table[source_cols[i]]) + .name(result_col) + ) + else: + values.append( + bigframes.dtypes.literal_to_ibis_scalar( + None, force_dtype=col_dtype + ).name(result_col) + ) + offsets_value = ( + ((table[ORDER_ID_COLUMN] * row_n) + i) + .cast(ibis_dtypes.int64) + .name(ORDER_ID_COLUMN), + ) sub_expr = table.select( - ibis_types.literal(row_labels[i]).name(index_col_id), + passthrough_columns, + bigframes.dtypes.literal_to_ibis_scalar( + row_labels[i], force_dtype=labels_dtype # type:ignore + ).name(index_col_id), *values, - ibis_types.literal(i).name(ORDER_ID_COLUMN), + offsets_value, ) sub_expressions.append(sub_expr) rotated_table = ibis.union(*sub_expressions) @@ -891,13 +959,15 @@ def unpivot_single_row( value_columns = [ rotated_table[value_col_id] for value_col_id, _ in unpivot_columns ] + passthrough_values = [rotated_table[col] for col in passthrough_columns] return ArrayValue( session=self._session, table=rotated_table, - columns=[rotated_table[index_col_id], *value_columns], + columns=[rotated_table[index_col_id], *value_columns, *passthrough_values], hidden_ordering_columns=[rotated_table[ORDER_ID_COLUMN]], ordering=ExpressionOrdering( ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), total_ordering_columns=frozenset([ORDER_ID_COLUMN]), ), ) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index b13d7bf2d3a..abf8b887d82 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -197,3 +197,35 @@ def rank( ) return block.select_columns(rownum_col_ids).with_column_labels(labels) + + +def dropna(block: blocks.Block, how: typing.Literal["all", "any"] = "any"): + """ + Drop na entries from block + """ + if how == "any": + filtered_block = block + for column in block.value_columns: + filtered_block, result_id = filtered_block.apply_unary_op( + column, ops.notnull_op + ) + filtered_block = filtered_block.filter(result_id) + filtered_block = filtered_block.drop_columns([result_id]) + return filtered_block + else: # "all" + filtered_block = block + predicate = None + for column in block.value_columns: + filtered_block, partial_predicate = filtered_block.apply_unary_op( + column, ops.notnull_op + ) + if predicate: + filtered_block, predicate = filtered_block.apply_binary_op( + partial_predicate, predicate, ops.or_op + ) + else: + predicate = partial_predicate + if predicate: + filtered_block = filtered_block.filter(predicate) + filtered_block = filtered_block.select_columns(block.value_columns) + return filtered_block diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 2731990febe..f23a4d0b5c9 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -80,15 +80,18 @@ def __init__( self, expr: core.ArrayValue, index_columns: Iterable[str], - column_labels: typing.Union[pd.Index, typing.Sequence[Label]], - index_labels: typing.Union[pd.Index, typing.Sequence[Label], None] = None, + column_labels: typing.Union[pd.Index, typing.Iterable[Label]], + index_labels: typing.Union[pd.Index, typing.Iterable[Label], None] = None, ): """Construct a block object, will create default index if no index columns specified.""" - if index_labels and (len(index_labels) != len(list(index_columns))): - raise ValueError( - "'index_columns' and 'index_labels' must have equal length" - ) - if len(list(index_columns)) == 0: + index_columns = list(index_columns) + if index_labels: + index_labels = list(index_labels) + if len(index_labels) != len(index_columns): + raise ValueError( + "'index_columns' and 'index_labels' must have equal length" + ) + if len(index_columns) == 0: expr, new_index_col_id = expr.promote_offsets() index_columns = [new_index_col_id] self._index_columns = tuple(index_columns) @@ -114,6 +117,7 @@ def __init__( self._stats_cache: dict[str, dict[str, typing.Any]] = { col_id: {} for col_id in self.value_columns } + # TODO(kemppeterson) Add a cache for corr to parallel the single-column stats. @property def index(self) -> indexes.IndexValue: @@ -826,9 +830,7 @@ def aggregate_all_and_pivot( dtype=pd.Float64Dtype(), ) -> Block: aggregations = [(col_id, operation, col_id) for col_id in self.value_columns] - result_expr = self.expr.aggregate( - aggregations, dropna=dropna - ).unpivot_single_row( + result_expr = self.expr.aggregate(aggregations, dropna=dropna).unpivot( row_labels=self.column_labels.to_list(), index_col_id="index", unpivot_columns=[(value_col_id, self.value_columns)], @@ -966,6 +968,26 @@ def get_stat(self, column_id: str, stat: agg_ops.AggregateOp): self._stats_cache[column_id].update(stats_map) return stats_map[stat.name] + def get_corr_stat(self, column_id_left: str, column_id_right: str): + # TODO(kemppeterson): Clean up the column names for DataFrames.corr support + # TODO(kemppeterson): Add a cache here. + corr_aggregations = [ + ( + column_id_left, + column_id_right, + "corr_" + column_id_left + column_id_right, + ) + ] + expr = self.expr.corr_aggregate(corr_aggregations) + expr, offset_index_id = expr.promote_offsets() + block = Block( + expr, + index_columns=[offset_index_id], + column_labels=[a[2] for a in corr_aggregations], + ) + df, _ = block.to_pandas() + return df.loc[0, "corr_" + column_id_left + column_id_right] + def summarize( self, column_ids: typing.Sequence[str], @@ -983,7 +1005,7 @@ def summarize( (col_id, [f"{col_id}-{stat.name}" for stat in stats]) for col_id in column_ids ] - expr = self.expr.aggregate(aggregations).unpivot_single_row( + expr = self.expr.aggregate(aggregations).unpivot( labels, unpivot_columns=columns, index_col_id=label_col_id, @@ -1166,6 +1188,121 @@ def pivot( return result_block.with_column_labels(column_index) + def stack(self): + """Unpivot last column axis level into row axis""" + if isinstance(self.column_labels, pd.MultiIndex): + return self._stack_multi() + else: + return self._stack_mono() + + def _stack_mono(self): + if isinstance(self.column_labels, pd.MultiIndex): + raise ValueError("Expected single level index") + + # These are the values that will be turned into rows + stack_values = self.column_labels.drop_duplicates().sort_values() + + # Get matching columns + unpivot_columns: List[Tuple[str, List[str]]] = [] + dtypes: List[bigframes.dtypes.Dtype] = [] + col_id = guid.generate_guid("unpivot_") + dtype = None + input_columns: Sequence[Optional[str]] = [] + for uvalue in stack_values: + matching_ids = self.label_to_col_id.get(uvalue, []) + input_id = matching_ids[0] if len(matching_ids) > 0 else None + if input_id: + if dtype and dtype != self._column_type(input_id): + raise NotImplementedError( + "Cannot stack columns with non-matching dtypes." + ) + else: + dtype = self._column_type(input_id) + input_columns.append(input_id) + unpivot_columns.append((col_id, input_columns)) + if dtype: + dtypes.append(dtype or pd.Float64Dtype()) + + added_index_column = col_id = guid.generate_guid() + unpivot_expr = self._expr.unpivot( + row_labels=stack_values, + passthrough_columns=self.index_columns, + unpivot_columns=unpivot_columns, + index_col_id=added_index_column, + dtype=dtypes, + ) + block = Block( + unpivot_expr, + index_columns=[*self.index_columns, added_index_column], + column_labels=[None], + index_labels=[*self._index_labels, self.column_labels.names[-1]], + ) + return block + + def _stack_multi(self): + if not isinstance(self.column_labels, pd.MultiIndex): + raise ValueError("Expected multi-index") + + # These are the values that will be turned into rows + stack_values = ( + self.column_labels.get_level_values(-1).drop_duplicates().sort_values() + ) + + result_col_labels = ( + self.column_labels.droplevel(-1) + .drop_duplicates() + .sort_values() + .dropna(how="all") + ) + + # Get matching columns + unpivot_columns: List[Tuple[str, List[str]]] = [] + dtypes = [] + for val in result_col_labels: + col_id = guid.generate_guid("unpivot_") + dtype = None + input_columns: Sequence[Optional[str]] = [] + for uvalue in stack_values: + # Need to unpack if still a multi-index after dropping 1 level + label_to_match = ( + (val, uvalue) if result_col_labels.nlevels == 1 else (*val, uvalue) + ) + matching_ids = self.label_to_col_id.get(label_to_match, []) + input_id = matching_ids[0] if len(matching_ids) > 0 else None + if input_id: + if dtype and dtype != self._column_type(input_id): + raise NotImplementedError( + "Cannot stack columns with non-matching dtypes." + ) + else: + dtype = self._column_type(input_id) + input_columns.append(input_id) + # Input column i is the first one that + unpivot_columns.append((col_id, input_columns)) + if dtype: + dtypes.append(dtype or pd.Float64Dtype()) + + added_index_column = col_id = guid.generate_guid() + unpivot_expr = self._expr.unpivot( + row_labels=stack_values, + passthrough_columns=self.index_columns, + unpivot_columns=unpivot_columns, + index_col_id=added_index_column, + dtype=dtypes, + ) + block = Block( + unpivot_expr, + index_columns=[*self.index_columns, added_index_column], + column_labels=result_col_labels, + index_labels=[*self._index_labels, self.column_labels.names[-1]], + ) + return block + + def _column_type(self, col_id: str) -> bigframes.dtypes.Dtype: + col_offset = self.value_columns.index(col_id) + dtype = self.dtypes[col_offset] + return dtype + @staticmethod def _create_pivot_column_index( value_labels: Sequence[typing.Hashable], columns_values: pd.Index diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 5b217effdd1..589c5c251ca 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -197,8 +197,11 @@ def _agg_string(self, func: str) -> df.DataFrame: return df.DataFrame(agg_block) def _agg_dict(self, func: typing.Mapping) -> df.DataFrame: - aggregations = [] + aggregations: typing.List[typing.Tuple[str, agg_ops.AggregateOp]] = [] column_labels = [] + + want_aggfunc_level = any(utils.is_list_like(aggs) for aggs in func.values()) + for label, funcs_for_id in func.items(): col_id = self._resolve_label(label) func_list = ( @@ -206,16 +209,22 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame: ) for f in func_list: aggregations.append((col_id, agg_ops.lookup_agg_func(f))) - column_labels.append((col_id, f)) + column_labels.append(label) agg_block, _ = self._block.aggregate( by_column_ids=self._by_col_ids, aggregations=aggregations, as_index=self._as_index, dropna=self._dropna, ) - agg_block = agg_block.with_column_labels( - pd.MultiIndex.from_tuples(column_labels) - ) + if want_aggfunc_level: + agg_block = agg_block.with_column_labels( + utils.combine_indices( + pd.Index(column_labels), + pd.Index(agg[1].name for agg in aggregations), + ) + ) + else: + agg_block = agg_block.with_column_labels(pd.Index(column_labels)) return df.DataFrame(agg_block) def _agg_list(self, func: typing.Sequence) -> df.DataFrame: @@ -234,7 +243,9 @@ def _agg_list(self, func: typing.Sequence) -> df.DataFrame: dropna=self._dropna, ) agg_block = agg_block.with_column_labels( - pd.MultiIndex.from_tuples(column_labels) + pd.MultiIndex.from_tuples( + column_labels, names=[*self._block.column_labels.names, None] + ) ) return df.DataFrame(agg_block) diff --git a/bigframes/core/indexes/__init__.py b/bigframes/core/indexes/__init__.py index d797c57955a..184a9ce262f 100644 --- a/bigframes/core/indexes/__init__.py +++ b/bigframes/core/indexes/__init__.py @@ -14,10 +14,7 @@ from bigframes.core.indexes.index import Index, IndexValue -INDEX_COLUMN_ID = "bigframes_index_{}" - __all__ = [ "Index", "IndexValue", - "INDEX_COLUMN_ID", ] diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py index a3300029054..1c0a2a1a81e 100644 --- a/bigframes/core/utils.py +++ b/bigframes/core/utils.py @@ -12,10 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. import typing +from typing import Hashable, Iterable, List import pandas as pd import typing_extensions +import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common + +UNNAMED_COLUMN_ID = "bigframes_unnamed_column" +UNNAMED_INDEX_ID = "bigframes_unnamed_index" + def get_axis_number(axis: typing.Union[str, int, None]) -> typing.Literal[0, 1]: if axis in {0, "index", "rows", None}: @@ -31,3 +37,50 @@ def is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence def is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]: return pd.api.types.is_dict_like(obj) + + +def combine_indices(index1: pd.Index, index2: pd.Index) -> pd.MultiIndex: + """Combines indices into multi-index while preserving dtypes, names.""" + multi_index = pd.MultiIndex.from_frame( + pd.concat([index1.to_frame(index=False), index2.to_frame(index=False)], axis=1) + ) + # to_frame will produce numbered default names, we don't want these + multi_index.names = [*index1.names, *index2.names] + return multi_index + + +def get_standardized_ids( + col_labels: Iterable[Hashable], idx_labels: Iterable[Hashable] = () +) -> tuple[list[str], list[str]]: + """Get stardardized column ids as column_ids_list, index_ids_list. + The standardized_column_id must be valid BQ SQL schema column names, can only be string type and unique. + + Args: + col_labels: column labels + + idx_labels: index labels, optional. If empty, will only return column ids. + + Return: + Tuple of (standardized_column_ids, standardized_index_ids) + """ + col_ids = [ + UNNAMED_COLUMN_ID if col_label is None else str(col_label) + for col_label in col_labels + ] + idx_ids = [ + UNNAMED_INDEX_ID if idx_label is None else str(idx_label) + for idx_label in idx_labels + ] + + ids = idx_ids + col_ids + # Column values will be loaded as null if the column name has spaces. + # https://github.com/googleapis/python-bigquery/issues/1566 + ids = [id.replace(" ", "_") for id in ids] + + ids = typing.cast( + List[str], + vendored_pandas_io_common.dedup_names(ids, is_potential_multiindex=False), + ) + idx_ids, col_ids = ids[: len(idx_ids)], ids[len(idx_ids) :] + + return col_ids, idx_ids diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 5fbe5d1f9eb..ef443db0797 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -56,7 +56,6 @@ import bigframes.series import bigframes.series as bf_series import third_party.bigframes_vendored.pandas.core.frame as vendored_pandas_frame -import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing if typing.TYPE_CHECKING: @@ -291,84 +290,59 @@ def astype( self, dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], ) -> DataFrame: - return self._apply_to_rows(ops.AsTypeOp(dtype)) + return self._apply_unary_op(ops.AsTypeOp(dtype)) def _to_sql_query( - self, always_include_index: bool - ) -> Tuple[str, List[Tuple[str, bool]]]: + self, include_index: bool + ) -> Tuple[str, list[str], list[blocks.Label]]: """Compiles this DataFrame's expression tree to SQL, optionally - including unnamed index columns. + including index columns. Args: - always_include_index (bool): - whether to include unnamed index columns. If False, only named - indexes are included. + include_index (bool): + whether to include index columns. - Returns: a tuple of (sql_string, index_column_list) - Each entry in the index column list is a tuple of (column_name, named). - If named is false, then the column name exists only in SQL + Returns: + a tuple of (sql_string, index_column_id_list, index_column_label_list). + If include_index is set to False, index_column_id_list and index_column_label_list + return empty lists. """ # Has to be unordered as it is impossible to order the sql without # including metadata columns in selection with ibis. ibis_expr = self._block.expr.to_ibis_expr(ordering_mode="unordered") - column_labels = list(self._block.column_labels) + col_labels, idx_labels = list(self._block.column_labels), list( + self._block.index_labels + ) + old_col_ids, old_idx_ids = list(self._block.value_columns), list( + self._block.index_columns + ) - # TODO(swast): Need to have a better way of controlling when to include - # the index or not. - index_has_names = all([name is not None for name in self.index.names]) - if index_has_names: - column_labels = column_labels + list(self.index.names) - elif always_include_index: - # In this mode include the index even if it is a nameless generated - # column like 'bigframes_index_0' - index_labels = [] - unnamed_index_count = 0 - for index_label in self._block.index_labels: - if isinstance(index_label, str): - index_labels.append(index_label) - else: - index_labels.append( - indexes.INDEX_COLUMN_ID.format(unnamed_index_count), - ) - unnamed_index_count += 1 + if not include_index: + idx_labels, old_idx_ids = [], [] + ibis_expr = ibis_expr.drop(*self._block.index_columns) - column_labels = column_labels + typing.cast( - List[Optional[str]], index_labels - ) + old_ids = old_idx_ids + old_col_ids + + new_col_ids, new_idx_ids = utils.get_standardized_ids(col_labels, idx_labels) + new_ids = new_idx_ids + new_col_ids - column_labels_deduped = typing.cast( - List[str], - vendored_pandas_io_common.dedup_names( - column_labels, is_potential_multiindex=False - ), - ) - column_ids = self._block.value_columns substitutions = {} - for column_id, column_label in zip(column_ids, column_labels_deduped): + for old_id, new_id in zip(old_ids, new_ids): # TODO(swast): Do we need to further escape this, or can we rely on # the BigQuery unicode column name feature? - substitutions[column_id] = column_label - - index_cols: List[Tuple[str, bool]] = [] - first_index_offset = len(self._block.column_labels) - if index_has_names or always_include_index: - for i, index_col in enumerate(self._block.index_columns): - offset = first_index_offset + i - substitutions[index_col] = column_labels_deduped[offset] - index_cols = [ - (label, index_has_names) - for label in column_labels_deduped[first_index_offset:] - ] - else: - ibis_expr = ibis_expr.drop(*self._block.index_columns) + substitutions[old_id] = new_id ibis_expr = ibis_expr.relabel(substitutions) - return typing.cast(str, ibis_expr.compile()), index_cols + return ( + typing.cast(str, ibis_expr.compile()), + new_ids[: len(idx_labels)], + idx_labels, + ) @property def sql(self) -> str: """Compiles this DataFrame's expression tree to SQL.""" - sql, _ = self._to_sql_query(always_include_index=False) + sql, _, _ = self._to_sql_query(include_index=False) return sql @property @@ -469,12 +443,12 @@ def __getattr__(self, key: str): if key in self._block.column_labels: return self.__getitem__(key) elif hasattr(pandas.DataFrame, key): - raise NotImplementedError( + raise AttributeError( textwrap.dedent( f""" - BigQuery DataFrames has not yet implemented an equivalent to - 'pandas.DataFrame.{key}'. {constants.FEEDBACK_LINK} - """ + BigQuery DataFrames has not yet implemented an equivalent to + 'pandas.DataFrame.{key}'. {constants.FEEDBACK_LINK} + """ ) ) else: @@ -872,6 +846,32 @@ def _assign_single_item( copy = self.copy() copy[k] = v(copy) return copy + elif utils.is_list_like(v): + given_rows = len(v) + actual_rows = len(self) + if given_rows != actual_rows: + raise ValueError( + f"Length of values ({given_rows}) does not match length of index ({actual_rows})" + ) + + local_df = bigframes.dataframe.DataFrame( + {k: v}, session=self._get_block().expr._session + ) + # local_df is likely (but not guarunteed) to be cached locally + # since the original list came from memory and so is probably < MAX_INLINE_DF_SIZE + + this_expr, this_offsets_col_id = self._get_block()._expr.promote_offsets() + block = blocks.Block( + expr=this_expr, + index_labels=self.index.names, + index_columns=self._block.index_columns, + column_labels=[this_offsets_col_id] + list(self._block.value_columns), + ) # offsets are temporarily the first value column, label set to id + this_df_with_offsets = DataFrame(data=block) + join_result = this_df_with_offsets.join( + other=local_df, on=this_offsets_col_id, how="left" + ) + return join_result.drop(columns=[this_offsets_col_id]) else: return self._assign_scalar(k, v) @@ -1024,13 +1024,7 @@ def add_suffix(self, suffix: str, axis: int | str | None = None) -> DataFrame: return DataFrame(self._get_block().add_suffix(suffix, axis)) def dropna(self) -> DataFrame: - block = self._block - for column in self._block.value_columns: - block, result_id = block.apply_unary_op(column, ops.notnull_op) - block = block.filter(result_id) - block = block.drop_columns([result_id]) - - return DataFrame(block) + return DataFrame(block_ops.dropna(self._block, how="any")) def any( self, @@ -1209,6 +1203,14 @@ def pivot( ) return DataFrame(pivot_block) + def stack(self): + # TODO: support 'level' param by simply reordering levels such that selected level is last before passing to Block.stack. + # TODO: support 'dropna' param by executing dropna only conditionally + result_block = block_ops.dropna(self._block.stack(), how="all") + if not isinstance(self.columns, pandas.MultiIndex): + return bigframes.series.Series(result_block) + return DataFrame(result_block) + def _drop_non_numeric(self, keep_bool=True) -> DataFrame: types_to_keep = set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) if not keep_bool: @@ -1508,15 +1510,15 @@ def _groupby_series( ) def abs(self) -> DataFrame: - return self._apply_to_rows(ops.abs_op) + return self._apply_unary_op(ops.abs_op) def isna(self) -> DataFrame: - return self._apply_to_rows(ops.isnull_op) + return self._apply_unary_op(ops.isnull_op) isnull = isna def notna(self) -> DataFrame: - return self._apply_to_rows(ops.notnull_op) + return self._apply_unary_op(ops.notnull_op) notnull = notna @@ -1736,7 +1738,7 @@ def to_parquet(self, path: str, *, index: bool = True) -> None: _, query_job = self._block.expr._session._start_query(export_data_statement) self._set_internal_query_job(query_job) - def _apply_to_rows(self, operation: ops.UnaryOp): + def _apply_unary_op(self, operation: ops.UnaryOp) -> DataFrame: block = self._block.multi_apply_unary_op(self._block.value_columns, operation) return DataFrame(block) @@ -1813,7 +1815,7 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: # to be applied before passing data to remote function, protecting from bad # inputs causing errors. reprojected_df = DataFrame(self._block._force_reproject()) - return reprojected_df._apply_to_rows( + return reprojected_df._apply_unary_op( ops.RemoteFunctionOp(func, apply_on_null=(na_action is None)) ) @@ -1871,6 +1873,25 @@ def _slice( block = self._block.slice(start=start, stop=stop, step=step) return DataFrame(block) + def __array_ufunc__( + self, ufunc: numpy.ufunc, method: str, *inputs, **kwargs + ) -> DataFrame: + """Used to support numpy ufuncs. + See: https://numpy.org/doc/stable/reference/ufuncs.html + """ + if ( + inputs[0] is not self + or method != "__call__" + or len(inputs) > 1 + or len(kwargs) > 0 + ): + return NotImplemented + + if ufunc in ops.NUMPY_TO_OP: + return self._apply_unary_op(ops.NUMPY_TO_OP[ufunc]) + + return NotImplemented + def _set_block(self, block: blocks.Block): self._block = block diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 812bb08dc38..27727c9f81f 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -74,20 +74,18 @@ def _apply_sql( string from which to construct the output dataframe. It must include the index columns of the input SQL. """ - source_sql, tagged_index_cols = input_data._to_sql_query( - always_include_index=True + source_sql, index_col_ids, index_labels = input_data._to_sql_query( + include_index=True ) - if len(tagged_index_cols) != 1: + if len(index_col_ids) != 1: raise NotImplementedError( f"Only exactly one index column is supported. {constants.FEEDBACK_LINK}" ) - index_col_name, is_named_index = tagged_index_cols[0] sql = func(source_sql) - df = session.read_gbq(sql, index_col=[index_col_name]) - if not is_named_index: - df.index.name = None + df = session.read_gbq(sql, index_col=index_col_ids) + df.index.names = index_labels return df @@ -150,10 +148,10 @@ def forecast(self) -> bpd.DataFrame: def evaluate(self, input_data: Optional[bpd.DataFrame] = None): # TODO: validate input data schema # Note: don't need index as evaluate returns a new table - source_sql, _ = ( - input_data._to_sql_query(always_include_index=False) + source_sql, _, _ = ( + input_data._to_sql_query(include_index=False) if (input_data is not None) - else (None, None) + else (None, None, None) ) sql = ml_sql.ml_evaluate(self.model_name, source_sql) @@ -166,6 +164,20 @@ def centroids(self): return self._session.read_gbq(sql) + def principal_components(self): + assert self._model.model_type == "PCA" + + sql = ml_sql.ml_principal_components(self.model_name) + + return self._session.read_gbq(sql) + + def principal_component_info(self): + assert self._model.model_type == "PCA" + + sql = ml_sql.ml_principal_component_info(self.model_name) + + return self._session.read_gbq(sql) + def copy(self, new_model_name: str, replace: bool = False) -> BqmlModel: job_config = bigquery.job.CopyJobConfig() if replace: diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 76b4f9ced6a..16106d3a7b3 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -70,6 +70,35 @@ def fit( ) return self + @property + def components_(self) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("A model must be fitted before calling components_.") + + return self._bqml_model.principal_components() + + @property + def explained_variance_(self) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError( + "A model must be fitted before calling explained_variance_." + ) + + return self._bqml_model.principal_component_info()[ + ["principal_component_id", "eigenvalue"] + ].rename(columns={"eigenvalue": "explained_variance"}) + + @property + def explained_variance_ratio_(self) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError( + "A model must be fitted before calling explained_variance_ratio_." + ) + + return self._bqml_model.principal_component_info()[ + ["principal_component_id", "explained_variance_ratio"] + ] + def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 80054d40e19..bcd82435823 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -182,3 +182,13 @@ def ml_generate_text_embedding( def ml_forecast(model_name: str) -> str: """Encode ML.FORECAST for BQML""" return f"""SELECT * FROM ML.FORECAST(MODEL `{model_name}`)""" + + +def ml_principal_components(model_name: str) -> str: + """Encode ML.PRINCIPAL_COMPONENTS for BQML""" + return f"""SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `{model_name}`)""" + + +def ml_principal_component_info(model_name: str) -> str: + """Encode ML.PRINCIPAL_COMPONENT_INFO for BQML""" + return f"""SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `{model_name}`)""" diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 58f19ea8e7f..9305cf1ddab 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -32,6 +32,12 @@ _ZERO = typing.cast(ibis_types.NumericValue, ibis_types.literal(0)) _NAN = typing.cast(ibis_types.NumericValue, ibis_types.literal(np.nan)) _INF = typing.cast(ibis_types.NumericValue, ibis_types.literal(np.inf)) +_NEG_INF = typing.cast(ibis_types.NumericValue, ibis_types.literal(-np.inf)) + +# Approx Highest number you can pass in to EXP function and get a valid FLOAT64 result +# FLOAT64 has 11 exponent bits, so max values is about 2**(2**10) +# ln(2**(2**10)) == (2**10)*ln(2) ~= 709.78, so EXP(x) for x>709.78 will overflow. +_FLOAT64_EXP_BOUND = typing.cast(ibis_types.NumericValue, ibis_types.literal(709.78)) BinaryOp = typing.Callable[[ibis_types.Value, ibis_types.Value], ibis_types.Value] TernaryOp = typing.Callable[ @@ -51,11 +57,142 @@ def is_windowed(self): return False +# Trig Functions class AbsOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.NumericValue, x).abs() +class SinOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).sin() + + +class CosOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).cos() + + +class TanOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).tan() + + +# Inverse trig functions +class ArcsinOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value.abs() <= _ibis_num(1) + return (~domain).ifelse(_NAN, numeric_value.asin()) + + +class ArccosOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value.abs() <= _ibis_num(1) + return (~domain).ifelse(_NAN, numeric_value.acos()) + + +class ArctanOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).atan() + + +# Hyperbolic trig functions +# BQ has these functions, but Ibis doesn't +class SinhOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + sinh_result = ( + numeric_value.exp() - (numeric_value.negate()).exp() + ) / _ibis_num(2) + domain = numeric_value.abs() < _FLOAT64_EXP_BOUND + return (~domain).ifelse(_INF * numeric_value.sign(), sinh_result) + + +class CoshOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + cosh_result = ( + numeric_value.exp() + (numeric_value.negate()).exp() + ) / _ibis_num(2) + domain = numeric_value.abs() < _FLOAT64_EXP_BOUND + return (~domain).ifelse(_INF, cosh_result) + + +class TanhOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + tanh_result = (numeric_value.exp() - (numeric_value.negate()).exp()) / ( + numeric_value.exp() + (numeric_value.negate()).exp() + ) + # Beyond +-20, is effectively just the sign function + domain = numeric_value.abs() < _ibis_num(20) + return (~domain).ifelse(numeric_value.sign(), tanh_result) + + +class ArcsinhOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + sqrt_part = ((numeric_value * numeric_value) + _ibis_num(1)).sqrt() + return (numeric_value.abs() + sqrt_part).ln() * numeric_value.sign() + + +class ArccoshOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + sqrt_part = ((numeric_value * numeric_value) - _ibis_num(1)).sqrt() + acosh_result = (numeric_value + sqrt_part).ln() + domain = numeric_value >= _ibis_num(1) + return (~domain).ifelse(_NAN, acosh_result) + + +class ArctanhOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value.abs() < _ibis_num(1) + numerator = numeric_value + _ibis_num(1) + denominator = _ibis_num(1) - numeric_value + ln_input = typing.cast(ibis_types.NumericValue, numerator.div(denominator)) + atanh_result = ln_input.ln().div(2) + + out_of_domain = (numeric_value.abs() == _ibis_num(1)).ifelse( + _INF * numeric_value, _NAN + ) + + return (~domain).ifelse(out_of_domain, atanh_result) + + +class SqrtOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value >= _ZERO + return (~domain).ifelse(_NAN, numeric_value.sqrt()) + + +class Log10Op(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value > _ZERO + out_of_domain = (numeric_value == _ZERO).ifelse(_NEG_INF, _NAN) + return (~domain).ifelse(out_of_domain, numeric_value.log10()) + + +class LnOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value > _ZERO + out_of_domain = (numeric_value == _ZERO).ifelse(_NEG_INF, _NAN) + return (~domain).ifelse(out_of_domain, numeric_value.ln()) + + +class ExpOp(UnaryOp): + def _as_ibis(self, x: ibis_types.Value): + numeric_value = typing.cast(ibis_types.NumericValue, x) + domain = numeric_value < _FLOAT64_EXP_BOUND + return (~domain).ifelse(_INF, numeric_value.exp()) + + class InvertOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): return typing.cast(ibis_types.NumericValue, x).negate() @@ -484,6 +621,28 @@ def _as_ibis(self, x: ibis_types.Value): year_op = YearOp() capitalize_op = CapitalizeOp() +# Just parameterless unary ops for now +# TODO: Parameter mappings +NUMPY_TO_OP: typing.Final = { + np.sin: SinOp(), + np.cos: CosOp(), + np.tan: TanOp(), + np.arcsin: ArcsinOp(), + np.arccos: ArccosOp(), + np.arctan: ArctanOp(), + np.sinh: SinhOp(), + np.cosh: CoshOp(), + np.tanh: TanhOp(), + np.arcsinh: ArcsinhOp(), + np.arccosh: ArccoshOp(), + np.arctanh: ArctanhOp(), + np.exp: ExpOp(), + np.log: LnOp(), + np.log10: Log10Op(), + np.sqrt: SqrtOp(), + np.abs: AbsOp(), +} + ### Binary Ops def short_circuit_nulls(type_override: typing.Optional[ibis_dtypes.DataType] = None): @@ -785,3 +944,7 @@ def partial_arg3(op: TernaryOp, scalar: typing.Any) -> BinaryOp: def is_null(value) -> bool: # float NaN/inf should be treated as distinct from 'true' null values return typing.cast(bool, pd.isna(value)) and not isinstance(value, float) + + +def _ibis_num(number: float): + return typing.cast(ibis_types.NumericValue, ibis_types.literal(number)) diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index 361fdca0558..81a5bc4c412 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -167,6 +167,11 @@ def _apply_binary_op( partial_op = ops.BinopPartialRight(op, other) return self._apply_unary_op(partial_op) + def _apply_corr_aggregation(self, other: series.Series) -> float: + (left, right, block) = self._align(other, how="outer") + + return block.get_corr_stat(left, right) + def _align(self, other: series.Series, how="outer") -> tuple[str, str, blocks.Block]: # type: ignore """Aligns the series value with another scalar or series object. Returns new left column id, right column id and joined tabled expression.""" values, block = self._align_n( diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index b688c18723b..280fce11126 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -148,13 +148,15 @@ def _set_default_session_location_if_possible(query): ): return - bqclient, _, _, _ = bigframes.session._create_cloud_clients( + clients_provider = bigframes.session.ClientsProvider( project=options.bigquery.project, location=options.bigquery.location, use_regional_endpoints=options.bigquery.use_regional_endpoints, credentials=options.bigquery.credentials, ) + bqclient = clients_provider.bqclient + if bigframes.session._is_query(query): job = bqclient.query(query, bigquery.QueryJobConfig(dry_run=True)) options.bigquery.location = job.location @@ -226,6 +228,33 @@ def read_csv( read_csv.__doc__ = inspect.getdoc(bigframes.session.Session.read_csv) +def read_json( + path_or_buf: str | IO["bytes"], + *, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + dtype: Optional[Dict] = None, + encoding: Optional[str] = None, + lines: bool = False, + engine: Literal["ujson", "pyarrow", "bigquery"] = "ujson", + **kwargs, +) -> bigframes.dataframe.DataFrame: + return global_session.with_default_session( + bigframes.session.Session.read_json, + path_or_buf=path_or_buf, + orient=orient, + dtype=dtype, + encoding=encoding, + lines=lines, + engine=engine, + **kwargs, + ) + + +read_json.__doc__ = inspect.getdoc(bigframes.session.Session.read_json) + + def read_gbq( query: str, *, diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 7cf74d63111..2a4b919dab0 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -34,7 +34,14 @@ import cloudpickle import google.api_core.exceptions -from google.cloud import bigquery, bigquery_connection_v1, functions_v2 +import google.api_core.retry +from google.cloud import ( + bigquery, + bigquery_connection_v1, + functions_v2, + resourcemanager_v3, +) +import google.iam.v1 from ibis.backends.bigquery.compiler import compiles from ibis.backends.bigquery.datatypes import BigQueryType from ibis.expr.datatypes.core import DataType as IbisDataType @@ -152,6 +159,7 @@ def __init__( bq_client, bq_connection_client, bq_connection_id, + cloud_resource_manager_client, ): self._gcp_project_id = gcp_project_id self._cloud_function_region = cloud_function_region @@ -161,6 +169,7 @@ def __init__( self._bq_client = bq_client self._bq_connection_client = bq_connection_client self._bq_connection_id = bq_connection_id + self._cloud_resource_manager_client = cloud_resource_manager_client def create_bq_remote_function( self, input_args, input_types, output_type, endpoint, bq_function_name @@ -175,7 +184,8 @@ def create_bq_remote_function( # raise ValueError("Failed to enable BigQuery Connection API") # If the intended connection does not exist then create it - if self.check_bq_connection_exists(): + service_account_id = self.get_service_account_if_connection_exists() + if service_account_id: logger.info(f"Connector {self._bq_connection_id} already exists") else: connection_name, service_account_id = self.create_bq_connection() @@ -183,21 +193,9 @@ def create_bq_remote_function( f"Created BQ connection {connection_name} with service account id: {service_account_id}" ) - # Set up access on the newly created BQ connection - # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function - # We would explicitly wait for 60+ seconds for the IAM binding to take effect - command_iam = ( - f"gcloud projects add-iam-policy-binding {self._gcp_project_id}" - + f' --member="serviceAccount:{service_account_id}"' - + ' --role="roles/run.invoker"' - ) - logger.info(f"Setting up IAM binding on the BQ connection: {command_iam}") - _run_system_command(command_iam) - - logger.info( - f"Waiting {self._iam_wait_seconds} seconds for IAM to take effect.." - ) - time.sleep(self._iam_wait_seconds) + # Ensure IAM role on the BQ connection + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + self._ensure_iam_binding(service_account_id, "run.invoker") # Create BQ function # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2 @@ -239,6 +237,53 @@ def get_cloud_function_endpoint(self, name): pass return None + # Introduce retries to accommodate transient errors like etag mismatch, + # which can be caused by concurrent operation on the same resource, and + # manifests with message like: + # google.api_core.exceptions.Aborted: 409 There were concurrent policy + # changes. Please retry the whole read-modify-write with exponential + # backoff. The request's ETag '\007\006\003,\264\304\337\272' did not match + # the current policy's ETag '\007\006\003,\3750&\363'. + @google.api_core.retry.Retry( + predicate=google.api_core.retry.if_exception_type( + google.api_core.exceptions.Aborted + ), + initial=10, + maximum=20, + multiplier=2, + timeout=60, + ) + def _ensure_iam_binding(self, service_account: str, role: str): + """Ensure necessary IAM role is configured on a service account.""" + project = f"projects/{self._gcp_project_id}" + service_account = f"serviceAccount:{service_account}" + role = f"roles/{role}" + request = google.iam.v1.iam_policy_pb2.GetIamPolicyRequest(resource=project) + policy = self._cloud_resource_manager_client.get_iam_policy(request=request) + + # Check if the binding already exists, and if does, do nothing more + for binding in policy.bindings: + if binding.role == role: + if service_account in binding.members: + return + + # Create a new binding + new_binding = google.iam.v1.policy_pb2.Binding( + role=role, members=[service_account] + ) + policy.bindings.append(new_binding) + request = google.iam.v1.iam_policy_pb2.SetIamPolicyRequest( + resource=project, policy=policy + ) + self._cloud_resource_manager_client.set_iam_policy(request=request) + + # We would wait for the IAM policy change to take effect + # https://cloud.google.com/iam/docs/access-change-propagation + logger.info( + f"Waiting {self._iam_wait_seconds} seconds for IAM to take effect.." + ) + time.sleep(self._iam_wait_seconds) + def create_bq_connection(self): """Create the BigQuery Connection and returns corresponding service account id.""" client = self._bq_connection_client @@ -253,7 +298,7 @@ def create_bq_connection(self): connection = client.create_connection(request) return connection.name, connection.cloud_resource.service_account_id - def check_bq_connection_exists(self): + def get_service_account_if_connection_exists(self) -> Optional[str]: """Check if the BigQuery Connection exists.""" client = self._bq_connection_client request = bigquery_connection_v1.GetConnectionRequest( @@ -262,12 +307,15 @@ def check_bq_connection_exists(self): ) ) + service_account = None try: - client.get_connection(request=request) - return True + service_account = client.get_connection( + request=request + ).cloud_resource.service_account_id except google.api_core.exceptions.NotFound: pass - return False + + return service_account def generate_udf_code(self, def_, dir): """Generate serialized bytecode using cloudpickle given a udf.""" @@ -624,6 +672,7 @@ def remote_function( bigquery_connection_v1.ConnectionServiceClient ] = None, cloud_functions_client: Optional[functions_v2.FunctionServiceClient] = None, + resource_manager_client: Optional[resourcemanager_v3.ProjectsClient] = None, dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, @@ -688,6 +737,11 @@ def remote_function( Client to use for BigQuery connection operations. If this param is not provided then bigquery connection client from the session would be used. + resource_manager_client (google.cloud.resourcemanager_v3.ProjectsClient, Optional): + Client to use for cloud resource management operations, e.g. for + getting and setting IAM roles on cloud resources. If this param is + not provided then resource manager client from the session would be + used. dataset (str, Optional.): Dataset in which to create a BigQuery remote function. It should be in `.` or `` format. If this @@ -734,7 +788,17 @@ def remote_function( cloud_functions_client = session.cloudfunctionsclient if not cloud_functions_client: raise ValueError( - "A functions connection client must be provided, either directly or via session. " + "A cloud functions client must be provided, either directly or via session. " + f"{constants.FEEDBACK_LINK}" + ) + + # A resource manager client is required to get/set IAM operations + if not resource_manager_client: + if session: + resource_manager_client = session.resourcemanagerclient + if not resource_manager_client: + raise ValueError( + "A resource manager client must be provided, either directly or via session. " f"{constants.FEEDBACK_LINK}" ) @@ -819,6 +883,7 @@ def wrapper(f): bigquery_client, bigquery_connection_client, bigquery_connection, + resource_manager_client, ) rf_name, cf_name = remote_function_client.provision_bq_remote_function( f, ibis_signature.input_types, ibis_signature.output_type, uniq_suffix diff --git a/bigframes/series.py b/bigframes/series.py index a1da93dee38..49b0a5b1f08 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -609,6 +609,39 @@ def round_op(x: ibis_types.Value, y: ibis_types.Value): return self._apply_binary_op(decimals, round_op) + def corr(self, other: Series, method="pearson", min_periods=None) -> float: + """ + Compute the correlation with the other Series. Non-number values are ignored in the + computation. + + Uses the "Pearson" method of correlation. Numbers are converted to float before + calculation, so the result may be unstable. + + Args: + other (Series): + The series with which this is to be correlated. + method (string, default "pearson"): + Correlation method to use - currently only "pearson" is supported. + min_periods (int, default None): + The minimum number of observations needed to return a result. Non-default values + are not yet supported, so a result will be returned for at least two observations. + + Returns: + float; Will return NaN if there are fewer than two numeric pairs, either series has a + variance or covariance of zero, or any input value is infinite. + """ + # TODO(kemppeterson): Validate early that both are numeric + # TODO(kemppeterson): Handle partially-numeric columns + if method != "pearson": + raise NotImplementedError( + f"Only Pearson correlation is currently supported. {constants.FEEDBACK_LINK}" + ) + if min_periods: + raise NotImplementedError( + f"min_periods not yet supported. {constants.FEEDBACK_LINK}" + ) + return self._apply_corr_aggregation(other) + def all(self) -> bool: return typing.cast(bool, self._apply_aggregation(agg_ops.all_op)) @@ -851,7 +884,7 @@ def __getitem__(self, indexer): def __getattr__(self, key: str): if hasattr(pandas.Series, key): - raise NotImplementedError( + raise AttributeError( textwrap.dedent( f""" BigQuery DataFrames has not yet implemented an equivalent to @@ -1158,6 +1191,26 @@ def to_string( def to_xarray(self): return self.to_pandas().to_xarray() + def __array_ufunc__( + self, ufunc: numpy.ufunc, method: str, *inputs, **kwargs + ) -> Series: + """Used to support numpy ufuncs. + See: https://numpy.org/doc/stable/reference/ufuncs.html + """ + # Only __call__ supported with zero arguments + if ( + inputs[0] is not self + or method != "__call__" + or len(inputs) > 1 + or len(kwargs) > 0 + ): + return NotImplemented + + if ufunc in ops.NUMPY_TO_OP: + return self._apply_unary_op(ops.NUMPY_TO_OP[ufunc]) + + return NotImplemented + # Keep this at the bottom of the Series class to avoid # confusing type checker by overriding str @property diff --git a/bigframes/session.py b/bigframes/session.py index 3ef52507460..ac2f8fa53ab 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -47,6 +47,7 @@ import google.cloud.bigquery_connection_v1 import google.cloud.bigquery_storage_v1 import google.cloud.functions_v2 +import google.cloud.resourcemanager_v3 import google.cloud.storage as storage # type: ignore import ibis import ibis.backends.bigquery as ibis_bigquery @@ -69,6 +70,7 @@ import bigframes.core.guid as guid import bigframes.core.io as bigframes_io from bigframes.core.ordering import IntegerEncoding, OrderingColumnReference +import bigframes.core.utils as utils import bigframes.dataframe as dataframe import bigframes.formatting_helpers as formatting_helpers from bigframes.remote_function import read_gbq_function as bigframes_rgf @@ -99,6 +101,16 @@ # TODO(swast): Need to connect to regional endpoints when performing remote # functions operations (BQ Connection IAM, Cloud Run / Cloud Functions). +# Also see if resource manager client library supports regional endpoints. + +_VALID_ENCODINGS = { + "UTF-8", + "ISO-8859-1", + "UTF-16BE", + "UTF-16LE", + "UTF-32BE", + "UTF-32LE", +} logger = logging.getLogger(__name__) @@ -112,90 +124,143 @@ def _get_default_credentials_with_project(): return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False) -def _create_cloud_clients( - project: Optional[str], - location: Optional[str], - use_regional_endpoints: Optional[bool], - credentials: Optional[google.auth.credentials.Credentials], -) -> typing.Tuple[ - bigquery.Client, - google.cloud.bigquery_connection_v1.ConnectionServiceClient, - google.cloud.bigquery_storage_v1.BigQueryReadClient, - google.cloud.functions_v2.FunctionServiceClient, -]: - """Create and initialize BigQuery client objects.""" - - credentials_project = None - if credentials is None: - credentials, credentials_project = _get_default_credentials_with_project() - - # Prefer the project in this order: - # 1. Project explicitly specified by the user - # 2. Project set in the environment - # 3. Project associated with the default credentials - project = ( - project - or os.getenv(_ENV_DEFAULT_PROJECT) - or typing.cast(Optional[str], credentials_project) - ) +class ClientsProvider: + """Provides client instances necessary to perform cloud operations.""" - if not project: - raise ValueError( - "Project must be set to initialize BigQuery client. " - "Try setting `bigframes.options.bigquery.project` first." + def __init__( + self, + project: Optional[str], + location: Optional[str], + use_regional_endpoints: Optional[bool], + credentials: Optional[google.auth.credentials.Credentials], + ): + credentials_project = None + if credentials is None: + credentials, credentials_project = _get_default_credentials_with_project() + + # Prefer the project in this order: + # 1. Project explicitly specified by the user + # 2. Project set in the environment + # 3. Project associated with the default credentials + project = ( + project + or os.getenv(_ENV_DEFAULT_PROJECT) + or typing.cast(Optional[str], credentials_project) ) - if use_regional_endpoints: - bq_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERY_REGIONAL_ENDPOINT.format(location=location), - ) - bqstorage_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERYSTORAGE_REGIONAL_ENDPOINT.format(location=location) - ) - bqconnection_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERYCONNECTION_REGIONAL_ENDPOINT.format(location=location) - ) - else: - bq_options = None - bqstorage_options = None - bqconnection_options = None - - bq_info = google.api_core.client_info.ClientInfo(user_agent=_APPLICATION_NAME) - bqclient = bigquery.Client( - client_info=bq_info, - client_options=bq_options, - credentials=credentials, - project=project, - location=location, - ) + if not project: + raise ValueError( + "Project must be set to initialize BigQuery client. " + "Try setting `bigframes.options.bigquery.project` first." + ) - bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo( - user_agent=_APPLICATION_NAME - ) - bqconnectionclient = google.cloud.bigquery_connection_v1.ConnectionServiceClient( - client_info=bqconnection_info, - client_options=bqconnection_options, - credentials=credentials, - ) + self._project = project + self._location = location + self._use_regional_endpoints = use_regional_endpoints + self._credentials = credentials - bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo( - user_agent=_APPLICATION_NAME - ) - bqstorageclient = google.cloud.bigquery_storage_v1.BigQueryReadClient( - client_info=bqstorage_info, - client_options=bqstorage_options, - credentials=credentials, - ) + # cloud clients initialized for lazy load + self._bqclient = None + self._bqconnectionclient = None + self._bqstorageclient = None + self._cloudfunctionsclient = None + self._resourcemanagerclient = None - functions_info = google.api_core.gapic_v1.client_info.ClientInfo( - user_agent=_APPLICATION_NAME - ) - cloudfunctionsclient = google.cloud.functions_v2.FunctionServiceClient( - client_info=functions_info, - credentials=credentials, - ) + @property + def bqclient(self): + if not self._bqclient: + bq_options = None + if self._use_regional_endpoints: + bq_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERY_REGIONAL_ENDPOINT.format( + location=self._location + ), + ) + bq_info = google.api_core.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + self._bqclient = bigquery.Client( + client_info=bq_info, + client_options=bq_options, + credentials=self._credentials, + project=self._project, + location=self._location, + ) - return bqclient, bqconnectionclient, bqstorageclient, cloudfunctionsclient + return self._bqclient + + @property + def bqconnectionclient(self): + if not self._bqconnectionclient: + bqconnection_options = None + if self._use_regional_endpoints: + bqconnection_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERYCONNECTION_REGIONAL_ENDPOINT.format( + location=self._location + ) + ) + bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + self._bqconnectionclient = ( + google.cloud.bigquery_connection_v1.ConnectionServiceClient( + client_info=bqconnection_info, + client_options=bqconnection_options, + credentials=self._credentials, + ) + ) + + return self._bqconnectionclient + + @property + def bqstorageclient(self): + if not self._bqstorageclient: + bqstorage_options = None + if self._use_regional_endpoints: + bqstorage_options = google.api_core.client_options.ClientOptions( + api_endpoint=_BIGQUERYSTORAGE_REGIONAL_ENDPOINT.format( + location=self._location + ) + ) + bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + self._bqstorageclient = google.cloud.bigquery_storage_v1.BigQueryReadClient( + client_info=bqstorage_info, + client_options=bqstorage_options, + credentials=self._credentials, + ) + + return self._bqstorageclient + + @property + def cloudfunctionsclient(self): + if not self._cloudfunctionsclient: + functions_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + self._cloudfunctionsclient = ( + google.cloud.functions_v2.FunctionServiceClient( + client_info=functions_info, + credentials=self._credentials, + ) + ) + + return self._cloudfunctionsclient + + @property + def resourcemanagerclient(self): + if not self._resourcemanagerclient: + resourcemanager_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=_APPLICATION_NAME + ) + self._resourcemanagerclient = ( + google.cloud.resourcemanager_v3.ProjectsClient( + credentials=self._credentials, client_info=resourcemanager_info + ) + ) + + return self._resourcemanagerclient class Session( @@ -221,12 +286,9 @@ def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): else: self._location = context.location - ( - self.bqclient, - self.bqconnectionclient, - self.bqstorageclient, - self.cloudfunctionsclient, - ) = _create_cloud_clients( + # Instantiate a clients provider to help with cloud clients that will be + # used in the future operations in the session + self._clients_provider = ClientsProvider( project=context.project, location=self._location, use_regional_endpoints=context.use_regional_endpoints, @@ -249,6 +311,26 @@ def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): # changed. context._session_started = True + @property + def bqclient(self): + return self._clients_provider.bqclient + + @property + def bqconnectionclient(self): + return self._clients_provider.bqconnectionclient + + @property + def bqstorageclient(self): + return self._clients_provider.bqstorageclient + + @property + def cloudfunctionsclient(self): + return self._clients_provider.cloudfunctionsclient + + @property + def resourcemanagerclient(self): + return self._clients_provider.resourcemanagerclient + @property def _session_dataset_id(self): """A dataset for storing temporary objects local to the session @@ -343,6 +425,38 @@ def read_gbq( max_results=max_results, ) + def _query_to_destination( + self, query: str, index_cols: List[str] + ) -> Tuple[Optional[bigquery.TableReference], Optional[bigquery.QueryJob]]: + # If there are no index columns, then there's no reason to cache to a + # (clustered) session table, as we'll just have to query it again to + # create a default index & ordering. + if not index_cols: + _, query_job = self._start_query(query) + return query_job.destination, query_job + + # If a dry_run indicates this is not a query type job, then don't + # bother trying to do a CREATE TEMP TABLE ... AS SELECT ... statement. + dry_run_config = bigquery.QueryJobConfig() + dry_run_config.dry_run = True + _, dry_run_job = self._start_query(query, job_config=dry_run_config) + if dry_run_job.statement_type != "SELECT": + _, query_job = self._start_query(query) + return query_job.destination, query_job + + # Make sure we cluster by the index column(s) so that subsequent + # operations are as speedy as they can be. + try: + ibis_expr = self.ibis_client.sql(query) + return self._ibis_to_session_table(ibis_expr, index_cols), None + except google.api_core.exceptions.BadRequest: + # Some SELECT statements still aren't compatible with CREATE TEMP + # TABLE ... AS SELECT ... statements. For example, if the query has + # a top-level ORDER BY, this conflicts with our ability to cluster + # the table by the index column(s). + _, query_job = self._start_query(query) + return query_job.destination, query_job + def read_gbq_query( self, query: str, @@ -368,16 +482,7 @@ def read_gbq_query( else: index_cols = list(index_col) - # Make sure we cluster by the index column so that subsequent - # operations are as speedy as they can be. - if index_cols: - # Since index_cols are specified, assume that we have a normal SQL - # query. DDL or DML not supported. - ibis_expr = self.ibis_client.sql(query) - destination = self._ibis_to_session_table(ibis_expr, index_cols) - else: - _, query_job = self._start_query(query) - destination = query_job.destination + destination, query_job = self._query_to_destination(query, index_cols) # If there was no destination table, that means the query must have # been DDL or DML. Return some job metadata, instead. @@ -385,9 +490,11 @@ def read_gbq_query( return dataframe.DataFrame( data=pandas.DataFrame( { - "statement_type": [query_job.statement_type], - "job_id": [query_job.job_id], - "location": [query_job.location], + "statement_type": [ + query_job.statement_type if query_job else "unknown" + ], + "job_id": [query_job.job_id if query_job else "unknown"], + "location": [query_job.location if query_job else "unknown"], } ), session=self, @@ -551,9 +658,10 @@ def _read_gbq_with_ordering( table_expression: ibis_types.Table, *, col_order: Iterable[str] = (), - index_cols: Sequence[str] = (), - index_labels: Sequence[Optional[str]] = (), - hidden_cols: Sequence[str] = (), + col_labels: Iterable[Optional[str]] = (), + index_cols: Iterable[str] = (), + index_labels: Iterable[Optional[str]] = (), + hidden_cols: Iterable[str] = (), ordering: core.ExpressionOrdering, is_total_ordering: bool = False, ) -> dataframe.DataFrame: @@ -563,9 +671,13 @@ def _read_gbq_with_ordering( table_expression: an ibis table expression to be executed in BigQuery. col_order: - List of BigQuery column names in the desired order for results DataFrame. + List of BigQuery column ids in the desired order for results DataFrame. + col_labels: + List of column labels as the column names. index_cols: - List of column names to use as the index or multi-index. + List of index ids to use as the index or multi-index. + index_labels: + List of index labels as names of index. hidden_cols: Columns that should be hidden. Ordering columns may (not always) be hidden ordering: @@ -574,6 +686,7 @@ def _read_gbq_with_ordering( Returns: A DataFrame representing results of the query or table. """ + index_cols, index_labels = list(index_cols), list(index_labels) if len(index_cols) != len(index_labels): raise ValueError( "Needs same number of index labels are there are index columns. " @@ -597,11 +710,14 @@ def _read_gbq_with_ordering( table_expression, index_cols ) index_col_values = [table_expression[index_id] for index_id in index_cols] + if not col_labels: + col_labels = column_keys return self._read_ibis( table_expression, index_col_values, index_labels, column_keys, + col_labels, ordering=ordering, ) @@ -650,9 +766,10 @@ def _read_bigquery_load_job( def _read_ibis( self, table_expression: ibis_types.Table, - index_cols: Sequence[ibis_types.Value], - index_labels: Sequence[Optional[str]], - column_keys: Sequence[str], + index_cols: Iterable[ibis_types.Value], + index_labels: Iterable[blocks.Label], + column_keys: Iterable[str], + column_labels: Iterable[blocks.Label], ordering: core.ExpressionOrdering, ) -> dataframe.DataFrame: """Turns a table expression (plus index column) into a DataFrame.""" @@ -674,7 +791,7 @@ def _read_ibis( self, table_expression, columns, hidden_ordering_columns, ordering ), index_columns=[index_col.get_name() for index_col in index_cols], - column_labels=column_keys, + column_labels=column_labels, index_labels=index_labels, ) @@ -713,15 +830,23 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame Returns: bigframes.dataframe.DataFrame: The BigQuery DataFrame. """ + col_labels, idx_labels = ( + pandas_dataframe.columns.to_list(), + pandas_dataframe.index.names, + ) + new_col_ids, new_idx_ids = utils.get_standardized_ids(col_labels, idx_labels) + # Add order column to pandas DataFrame to preserve order in BigQuery ordering_col = "rowid" - columns = frozenset(pandas_dataframe.columns) + columns = frozenset(col_labels + idx_labels) suffix = 2 while ordering_col in columns: ordering_col = f"rowid_{suffix}" suffix += 1 pandas_dataframe_copy = pandas_dataframe.copy() + pandas_dataframe_copy.index.names = new_idx_ids + pandas_dataframe_copy.columns = pandas.Index(new_col_ids) pandas_dataframe_copy[ordering_col] = np.arange(pandas_dataframe_copy.shape[0]) # Specify the datetime dtypes, which is auto-detected as timestamp types. @@ -732,27 +857,12 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame bigquery.SchemaField(column, bigquery.enums.SqlTypeNames.DATETIME) ) - # Unnamed are not copied to BigQuery when load_table_from_dataframe - # executes. - index_cols = list( - filter(lambda name: name is not None, pandas_dataframe_copy.index.names) - ) - index_labels = typing.cast(List[Optional[str]], index_cols) - # Clustering probably not needed anyways as pandas tables are small cluster_cols = [ordering_col] - if len(index_cols) == 0: - # Block constructor will implicitly build default index - pass - job_config = bigquery.LoadJobConfig(schema=schema) job_config.clustering_fields = cluster_cols - # TODO(swast): Rename the unnamed index columns and restore them after - # the load job completes. - # Column values will be loaded as null if the column name has spaces. - # https://github.com/googleapis/python-bigquery/issues/1566 load_table_destination = self._create_session_table() load_job = self.bqclient.load_table_from_dataframe( pandas_dataframe_copy, @@ -770,14 +880,22 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame f"SELECT * FROM `{load_table_destination.table_id}`" ) - return self._read_gbq_with_ordering( + # b/297590178 Potentially a bug in bqclient.load_table_from_dataframe(), that only when the DF is empty, the index columns disappear in table_expression. + if any( + [new_idx_id not in table_expression.columns for new_idx_id in new_idx_ids] + ): + new_idx_ids, idx_labels = [], [] + + df = self._read_gbq_with_ordering( table_expression=table_expression, - index_cols=index_cols, - index_labels=index_labels, + col_labels=col_labels, + index_cols=new_idx_ids, + index_labels=idx_labels, hidden_cols=(ordering_col,), ordering=ordering, is_total_ordering=True, ) + return df def read_csv( self, @@ -844,10 +962,9 @@ def read_csv( f"{constants.FEEDBACK_LINK}" ) - valid_encodings = {"UTF-8", "ISO-8859-1"} - if encoding is not None and encoding not in valid_encodings: + if encoding is not None and encoding not in _VALID_ENCODINGS: raise NotImplementedError( - f"BigQuery engine only supports the following encodings: {valid_encodings}. " + f"BigQuery engine only supports the following encodings: {_VALID_ENCODINGS}. " f"{constants.FEEDBACK_LINK}" ) @@ -933,6 +1050,86 @@ def read_parquet( return self._read_bigquery_load_job(path, table, job_config=job_config) + def read_json( + self, + path_or_buf: str | IO["bytes"], + *, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + dtype: Optional[Dict] = None, + encoding: Optional[str] = None, + lines: bool = False, + engine: Literal["ujson", "pyarrow", "bigquery"] = "ujson", + **kwargs, + ) -> dataframe.DataFrame: + table = bigquery.Table(self._create_session_table()) + + if engine == "bigquery": + + if dtype is not None: + raise NotImplementedError( + "BigQuery engine does not support the dtype arguments." + ) + + if not lines: + raise NotImplementedError( + "Only newline delimited JSON format is supported." + ) + + if encoding is not None and encoding not in _VALID_ENCODINGS: + raise NotImplementedError( + f"BigQuery engine only supports the following encodings: {_VALID_ENCODINGS}" + ) + + if lines and orient != "records": + raise ValueError( + "'lines' keyword is only valid when 'orient' is 'records'." + ) + + job_config = bigquery.LoadJobConfig() + job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED + job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON + job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY + job_config.autodetect = True + job_config.encoding = encoding + + return self._read_bigquery_load_job( + path_or_buf, + table, + job_config=job_config, + ) + else: + if any(arg in kwargs for arg in ("chunksize", "iterator")): + raise NotImplementedError( + "'chunksize' and 'iterator' arguments are not supported." + ) + + if isinstance(path_or_buf, str): + self._check_file_size(path_or_buf) + + if engine == "ujson": + pandas_df = pandas.read_json( # type: ignore + path_or_buf, + orient=orient, + dtype=dtype, + encoding=encoding, + lines=lines, + **kwargs, + ) + + else: + pandas_df = pandas.read_json( # type: ignore + path_or_buf, + orient=orient, + dtype=dtype, + encoding=encoding, + lines=lines, + engine=engine, + **kwargs, + ) + return self.read_pandas(pandas_df) + def _check_file_size(self, filepath: str): max_size = 1024 * 1024 * 1024 # 1 GB in bytes if filepath.startswith("gs://"): # GCS file path @@ -1008,14 +1205,26 @@ def _query_to_session_table( table = self._create_session_table() cluster_cols_sql = ", ".join(f"`{cluster_col}`" for cluster_col in cluster_cols) - # TODO(swast): This might not support multi-statement SQL queries. + # TODO(swast): This might not support multi-statement SQL queries (scripts). ddl_text = f""" CREATE TEMP TABLE `_SESSION`.`{table.table_id}` CLUSTER BY {cluster_cols_sql} AS {query_text} """ + + job_config = bigquery.QueryJobConfig() + + # Include a label so that Dataplex Lineage can identify temporary + # tables that BigQuery DataFrames creates. Googlers: See internal issue + # 296779699. We're labeling the job instead of the table because + # otherwise we get `BadRequest: 400 OPTIONS on temporary tables are not + # supported`. + job_config.labels = {"source": "bigquery-dataframes-temp"} + try: - self._start_query(ddl_text) # Wait for the job to complete + self._start_query( + ddl_text, job_config=job_config + ) # Wait for the job to complete except google.api_core.exceptions.Conflict: # Allow query retry to succeed. pass diff --git a/mypy.ini b/mypy.ini index ce78c4686ed..901394813aa 100644 --- a/mypy.ini +++ b/mypy.ini @@ -15,6 +15,9 @@ ignore_missing_imports = True [mypy-google.colab] ignore_missing_imports = True +[mypy-google.iam.*] +ignore_missing_imports = True + [mypy-pytz] ignore_missing_imports = True diff --git a/notebooks/dataframes/dataframe.ipynb b/notebooks/dataframes/dataframe.ipynb index 241c767f572..85ea61d2810 100644 --- a/notebooks/dataframes/dataframe.ipynb +++ b/notebooks/dataframes/dataframe.ipynb @@ -31,7 +31,22 @@ "execution_count": 2, "id": "96757c59-fc22-420e-a42f-c6cb956110ec", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "944f0e4417154e81b6496302fe756465", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job ac4d1f2b-e9f3-4d95-b78d-57e40eee93fa is RUNNING. Cubs\n", " 175\n", " \n", + " \n", + " 20\n", + " 71ab82a4-6e07-430a-b695-1af3bc42ea61\n", + " 2016\n", + " Nationals\n", + " Cubs\n", + " 257\n", + " \n", + " \n", + " 21\n", + " d1a110c2-f6c8-4029-bcd8-2f8a01e1561c\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 178\n", + " \n", + " \n", + " 22\n", + " 6d111b57-fa0b-4f24-82df-ff33a26f0252\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 171\n", + " \n", + " \n", + " 23\n", + " a97e9539-bbbd-4e03-bf15-f25ea2c1d923\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 248\n", + " \n", + " \n", + " 24\n", + " dc0c9218-505c-4725-8c0c-40b72cca0956\n", + " 2016\n", + " Astros\n", + " Cubs\n", + " 174\n", + " \n", " \n", "\n", + "

25 rows × 5 columns

\n", "[2431 rows x 5 columns in total]" ], "text/plain": [ @@ -419,6 +503,34 @@ "tags": [] }, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "276760df4c904ced81cbaff3a65d026e", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 1943ae42-bcbd-4c2f-914f-209377b5c4d9 is DONE. 0 Bytes processed. Cubs\n", " 175\n", " \n", + " \n", + " 20\n", + " 71ab82a4-6e07-430a-b695-1af3bc42ea61\n", + " 2016\n", + " Nationals\n", + " Cubs\n", + " 257\n", + " \n", + " \n", + " 21\n", + " d1a110c2-f6c8-4029-bcd8-2f8a01e1561c\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 178\n", + " \n", + " \n", + " 22\n", + " 6d111b57-fa0b-4f24-82df-ff33a26f0252\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 171\n", + " \n", + " \n", + " 23\n", + " a97e9539-bbbd-4e03-bf15-f25ea2c1d923\n", + " 2016\n", + " Brewers\n", + " Cubs\n", + " 248\n", + " \n", + " \n", + " 24\n", + " dc0c9218-505c-4725-8c0c-40b72cca0956\n", + " 2016\n", + " Astros\n", + " Cubs\n", + " 174\n", + " \n", " \n", "\n", + "

25 rows × 5 columns

\n", "[2431 rows x 5 columns in total]" ], "text/plain": [ @@ -3063,6 +3768,34 @@ "id": "ac3ceabe-4317-453c-9418-826de5094454", "metadata": {}, "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "c73064d64afe41cea6738085b273e29d", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 08477df8-3e86-4f94-a905-8ac2f36e2b69 is DONE. 0 Bytes processed. core.BqmlModel: +def penguins_bqml_kmeans_model( + session: bigframes.Session, penguins_kmeans_model_name: str +) -> core.BqmlModel: model = session.bqclient.get_model(penguins_kmeans_model_name) return core.BqmlModel(session, model) +@pytest.fixture(scope="session") +def penguins_bqml_pca_model( + session: bigframes.Session, penguins_pca_model_name: str +) -> core.BqmlModel: + model = session.bqclient.get_model(penguins_pca_model_name) + return core.BqmlModel(session, model) + + @pytest.fixture(scope="session") def penguins_linear_model( session, penguins_linear_model_name: str @@ -140,32 +147,12 @@ def penguins_kmeans_model(session, penguins_kmeans_model_name: str) -> cluster.K @pytest.fixture(scope="session") def penguins_pca_model( - session: bigframes.Session, dataset_id_permanent, penguins_table_id + session: bigframes.Session, penguins_pca_model_name: str ) -> decomposition.PCA: - - # TODO(yunmengxie): Create a shared method to get different types of pretrained models. - sql = f""" -CREATE OR REPLACE MODEL `$model_name` -OPTIONS ( - model_type='pca', - num_principal_components=3 -) AS SELECT - * -FROM `{penguins_table_id}`""" - # We use the SQL hash as the name to ensure the model is regenerated if this fixture is edited - model_name = ( - f"{dataset_id_permanent}.penguins_pca_{hashlib.md5(sql.encode()).hexdigest()}" + return cast( + decomposition.PCA, + session.read_gbq_model(penguins_pca_model_name), ) - sql = sql.replace("$model_name", model_name) - - try: - return session.read_gbq_model(model_name) - except google.cloud.exceptions.NotFound: - logging.info( - "penguins_pca_model fixture was not found in the permanent dataset, regenerating it..." - ) - session.bqclient.query(sql).result() - return session.read_gbq_model(model_name) @pytest.fixture(scope="session") diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index 4b184b0d4c8..6c3e8e06f5b 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -140,6 +140,100 @@ def test_model_centroids(penguins_bqml_kmeans_model: core.BqmlModel): ) +def test_pca_model_principal_components(penguins_bqml_pca_model: core.BqmlModel): + result = penguins_bqml_pca_model.principal_components().to_pandas() + assert result.shape == (21, 4) + + # result is too long, only check the first principal component here. + result = result.head(7) + expected = pd.DataFrame( + { + "principal_component_id": [0] * 7, + "feature": [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + "sex", + ], + "numerical_value": [ + pd.NA, + pd.NA, + 0.401489, + -0.377482, + 0.524052, + 0.501174, + pd.NA, + ], + "categorical_value": [ + [ + { + "category": "Gentoo penguin (Pygoscelis papua)", + "value": 0.25068877125667804, + }, + { + "category": "Adelie Penguin (Pygoscelis adeliae)", + "value": -0.20622291900416198, + }, + { + "category": "Chinstrap penguin (Pygoscelis antarctica)", + "value": -0.030161149275185855, + }, + ], + [ + {"category": "Biscoe", "value": 0.19761120114410635}, + {"category": "Dream", "value": -0.11264736305259061}, + {"category": "Torgersen", "value": -0.07065913511418596}, + ], + [], + [], + [], + [], + [ + {"category": ".", "value": 0.0015916894448071784}, + {"category": "MALE", "value": 0.06869704739750442}, + {"category": "FEMALE", "value": -0.052521171596813174}, + {"category": "_null_filler", "value": -0.0034628622681684906}, + ], + ], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + check_dtype=False, + ) + + +def test_pca_model_principal_component_info(penguins_bqml_pca_model: core.BqmlModel): + result = penguins_bqml_pca_model.principal_component_info().to_pandas() + assert result.shape == (3, 4) + + expected = pd.DataFrame( + { + "principal_component_id": [0, 1, 2], + "eigenvalue": [3.278657, 1.270829, 1.125354], + "explained_variance_ratio": [0.469357, 0.181926, 0.1611], + "cumulative_explained_variance_ratio": [0.469357, 0.651283, 0.812383], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame + check_index_type=False, + check_dtype=False, + ) + + def test_model_predict(penguins_bqml_linear_model: core.BqmlModel, new_penguins_df): predictions = penguins_bqml_linear_model.predict(new_penguins_df).to_pandas() expected = pd.DataFrame( diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index 01d52077505..8df4145fcf0 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -55,7 +55,7 @@ def test_pca_predict(session, penguins_pca_model: decomposition.PCA): ) -def test_pca_score(session, penguins_pca_model: decomposition.PCA): +def test_pca_score(penguins_pca_model: decomposition.PCA): result = penguins_pca_model.score().to_pandas() expected = pd.DataFrame( {"total_explained_variance_ratio": [0.812383]}, @@ -68,3 +68,110 @@ def test_pca_score(session, penguins_pca_model: decomposition.PCA): rtol=0.1, check_index_type=False, ) + + +def test_pca_components_(penguins_pca_model: decomposition.PCA): + result = penguins_pca_model.components_.to_pandas() + + # result is too long, only check the first principal component here. + result = result.head(7) + expected = pd.DataFrame( + { + "principal_component_id": [0] * 7, + "feature": [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + "sex", + ], + "numerical_value": [ + pd.NA, + pd.NA, + 0.401489, + -0.377482, + 0.524052, + 0.501174, + pd.NA, + ], + "categorical_value": [ + [ + { + "category": "Gentoo penguin (Pygoscelis papua)", + "value": 0.25068877125667804, + }, + { + "category": "Adelie Penguin (Pygoscelis adeliae)", + "value": -0.20622291900416198, + }, + { + "category": "Chinstrap penguin (Pygoscelis antarctica)", + "value": -0.030161149275185855, + }, + ], + [ + {"category": "Biscoe", "value": 0.19761120114410635}, + {"category": "Dream", "value": -0.11264736305259061}, + {"category": "Torgersen", "value": -0.07065913511418596}, + ], + [], + [], + [], + [], + [ + {"category": ".", "value": 0.0015916894448071784}, + {"category": "MALE", "value": 0.06869704739750442}, + {"category": "FEMALE", "value": -0.052521171596813174}, + {"category": "_null_filler", "value": -0.0034628622681684906}, + ], + ], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + check_dtype=False, + ) + + +def test_pca_explained_variance_(penguins_pca_model: decomposition.PCA): + result = penguins_pca_model.explained_variance_.to_pandas() + + expected = pd.DataFrame( + { + "principal_component_id": [0, 1, 2], + "explained_variance": [3.278657, 1.270829, 1.125354], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + check_dtype=False, + ) + + +def test_pca_explained_variance_ratio_(penguins_pca_model: decomposition.PCA): + result = penguins_pca_model.explained_variance_ratio_.to_pandas() + + expected = pd.DataFrame( + { + "principal_component_id": [0, 1, 2], + "explained_variance_ratio": [0.469357, 0.181926, 0.1611], + }, + ) + pd.testing.assert_frame_equal( + result, + expected, + check_exact=False, + rtol=0.1, + check_index_type=False, + check_dtype=False, + ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 5b4f9ebcccf..01305adb204 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -77,6 +77,20 @@ def test_df_construct_from_series(scalars_dfs): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_df_construct_from_dict(): + input_dict = { + "Animal": ["Falcon", "Falcon", "Parrot", "Parrot"], + # With a space in column name. We use standardized SQL schema ids to solve the problem that BQ schema doesn't support column names with spaces. b/296751058 + "Max Speed": [380.0, 370.0, 24.0, 26.0], + } + bf_result = dataframe.DataFrame(input_dict).to_pandas() + pd_result = pd.DataFrame(input_dict) + + pandas.testing.assert_frame_equal( + bf_result, pd_result, check_dtype=False, check_index_type=False + ) + + def test_get_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_col" @@ -356,6 +370,52 @@ def test_assign_new_column_w_setitem(scalars_dfs): pd.testing.assert_frame_equal(bf_result, pd_result) +def test_assign_new_column_w_setitem_list(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df["new_col"] = [9, 8, 7, 6, 5, 4, 3, 2, 1] + pd_df["new_col"] = [9, 8, 7, 6, 5, 4, 3, 2, 1] + bf_result = bf_df.to_pandas() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_new_column_w_setitem_list_custom_index(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + + # set the custom index + pd_df = pd_df.set_index("string_col") + bf_df = bf_df.set_index("string_col") + + bf_df["new_col"] = [9, 8, 7, 6, 5, 4, 3, 2, 1] + pd_df["new_col"] = [9, 8, 7, 6, 5, 4, 3, 2, 1] + bf_result = bf_df.to_pandas() + pd_result = pd_df + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_assign_new_column_w_setitem_list_error(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + + with pytest.raises(ValueError): + pd_df["new_col"] = [1, 2, 3] # should be len 9, is 3 + with pytest.raises(ValueError): + bf_df["new_col"] = [1, 2, 3] + + def test_assign_existing_column(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs kwargs = {"int64_col": 2} @@ -1329,6 +1389,21 @@ def test_df_describe(scalars_dfs): ).all() +def test_df_stack(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + # To match bigquery dataframes + scalars_pandas_df = scalars_pandas_df.copy() + scalars_pandas_df.columns = scalars_pandas_df.columns.astype("string[pyarrow]") + # Can only stack identically-typed columns + columns = ["int64_col", "int64_too", "rowindex_2"] + + bf_result = scalars_df[columns].stack().to_pandas() + pd_result = scalars_pandas_df[columns].stack() + + # Pandas produces NaN, where bq dataframes produces pd.NA + pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + @pytest.mark.parametrize( ("values", "index", "columns"), [ @@ -1734,8 +1809,9 @@ def test_df___array__(scalars_df_index, scalars_pandas_df_index): ) -def test_getattr_not_implemented(scalars_df_index): - with pytest.raises(NotImplementedError): +def test_getattr_attribute_error_when_pandas_has(scalars_df_index): + # asof is implemented in pandas but not in bigframes + with pytest.raises(AttributeError): scalars_df_index.asof() diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index 1f5aa906c8c..3886b85f40c 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -32,6 +32,7 @@ import bigframes import bigframes.dataframe +import bigframes.pandas as bpd def test_to_pandas_w_correct_dtypes(scalars_df_default_index): @@ -339,51 +340,68 @@ def test_to_parquet_index(scalars_dfs, gcs_folder, index): pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df) +def test_to_sql_query_unnamed_index_included( + session: bigframes.Session, + scalars_df_default_index: bpd.DataFrame, + scalars_pandas_df_default_index: pd.DataFrame, +): + bf_df = scalars_df_default_index.reset_index(drop=True) + sql, idx_ids, idx_labels = bf_df._to_sql_query(include_index=True) + assert len(idx_labels) == 1 + assert len(idx_ids) == 1 + assert idx_labels[0] is None + assert idx_ids[0].startswith("bigframes") + + pd_df = scalars_pandas_df_default_index.reset_index(drop=True) + roundtrip = session.read_gbq(sql, index_col=idx_ids) + roundtrip.index.names = [None] + assert_pandas_df_equal_ignore_ordering(roundtrip.to_pandas(), pd_df) + + def test_to_sql_query_named_index_included( - session, scalars_df_index, scalars_pandas_df_index + session: bigframes.Session, + scalars_df_default_index: bpd.DataFrame, + scalars_pandas_df_default_index: pd.DataFrame, ): - sql, index_columns = scalars_df_index._to_sql_query(always_include_index=True) - assert len(index_columns) == 1 - index_column, is_named = index_columns[0] - assert index_column == "rowindex" - assert is_named - - roundtrip = session.read_gbq(sql, index_col=[index_column]) - assert_pandas_df_equal_ignore_ordering( - roundtrip.to_pandas(), scalars_pandas_df_index - ) + bf_df = scalars_df_default_index.set_index("rowindex_2", drop=True) + sql, idx_ids, idx_labels = bf_df._to_sql_query(include_index=True) + assert len(idx_labels) == 1 + assert len(idx_ids) == 1 + assert idx_labels[0] == "rowindex_2" + assert idx_ids[0] == "rowindex_2" + + pd_df = scalars_pandas_df_default_index.set_index("rowindex_2", drop=True) + roundtrip = session.read_gbq(sql, index_col=idx_ids) + assert_pandas_df_equal_ignore_ordering(roundtrip.to_pandas(), pd_df) def test_to_sql_query_unnamed_index_excluded( - session, scalars_df_default_index, scalars_pandas_df_default_index + session: bigframes.Session, + scalars_df_default_index: bpd.DataFrame, + scalars_pandas_df_default_index: pd.DataFrame, ): - # The .sql property should return SQL without the unnamed indexes - sql, index_columns = scalars_df_default_index._to_sql_query( - always_include_index=False - ) - assert len(index_columns) == 0 + bf_df = scalars_df_default_index.reset_index(drop=True) + sql, idx_ids, idx_labels = bf_df._to_sql_query(include_index=False) + assert len(idx_labels) == 0 + assert len(idx_ids) == 0 + pd_df = scalars_pandas_df_default_index.reset_index(drop=True) roundtrip = session.read_gbq(sql) - assert_pandas_df_equal_ignore_ordering( - roundtrip.to_pandas(), scalars_pandas_df_default_index - ) + assert_pandas_df_equal_ignore_ordering(roundtrip.to_pandas(), pd_df) -def test_to_sql_query_unnamed_index_always_include( - session, - scalars_df_default_index: bigframes.dataframe.DataFrame, - scalars_pandas_df_default_index, +def test_to_sql_query_named_index_excluded( + session: bigframes.Session, + scalars_df_default_index: bpd.DataFrame, + scalars_pandas_df_default_index: pd.DataFrame, ): - sql, index_columns = scalars_df_default_index._to_sql_query( - always_include_index=True - ) - assert len(index_columns) == 1 - index_column, is_named = index_columns[0] - assert index_column == "bigframes_index_0" - assert not is_named - - roundtrip = session.read_gbq(sql, index_col=[index_column]) - roundtrip.index.name = None - assert_pandas_df_equal_ignore_ordering( - roundtrip.to_pandas(), scalars_pandas_df_default_index - ) + bf_df = scalars_df_default_index.set_index("rowindex_2", drop=True) + sql, idx_ids, idx_labels = bf_df._to_sql_query(include_index=False) + assert len(idx_labels) == 0 + assert len(idx_ids) == 0 + + pd_df = scalars_pandas_df_default_index.set_index( + "rowindex_2", drop=True + ).reset_index(drop=True) + roundtrip = session.read_gbq(sql) + assert_pandas_df_equal_ignore_ordering(roundtrip.to_pandas(), pd_df) diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 5a2562bfb20..987368ce771 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -112,7 +112,9 @@ def test_dataframe_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) -def test_dataframe_groupby_agg_dict(scalars_df_index, scalars_pandas_df_index): +def test_dataframe_groupby_agg_dict_with_list( + scalars_df_index, scalars_pandas_df_index +): col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] bf_result = ( scalars_df_index[col_names] @@ -129,6 +131,23 @@ def test_dataframe_groupby_agg_dict(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) +def test_dataframe_groupby_agg_dict_no_lists(scalars_df_index, scalars_pandas_df_index): + col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] + bf_result = ( + scalars_df_index[col_names] + .groupby("string_col") + .agg({"int64_too": "mean", "string_col": "count"}) + ) + pd_result = ( + scalars_pandas_df_index[col_names] + .groupby("string_col") + .agg({"int64_too": "mean", "string_col": "count"}) + ) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + def test_dataframe_groupby_agg_named(scalars_df_index, scalars_pandas_df_index): col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"] bf_result = ( diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index b2937d7da9c..1baf3e66508 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -620,3 +620,39 @@ def test_column_multi_index_cumsum(scalars_df_index, scalars_pandas_df_index): pd_result = pd_df.cumsum() pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +def test_column_multi_index_stack(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "rowindex_2"] + level1 = pandas.Index(["b", "a", "b"]) + # Need resulting column to be pyarrow string rather than object dtype + level2 = pandas.Index(["a", "b", "b"], dtype="string[pyarrow]") + multi_columns = pandas.MultiIndex.from_arrays([level1, level2]) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.stack().to_pandas() + pd_result = pd_df.stack() + + # Pandas produces NaN, where bq dataframes produces pd.NA + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +def test_column_multi_index_w_na_stack(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "rowindex_2"] + level1 = pandas.Index(["b", pandas.NA, pandas.NA]) + # Need resulting column to be pyarrow string rather than object dtype + level2 = pandas.Index([pandas.NA, "b", "b"], dtype="string[pyarrow]") + multi_columns = pandas.MultiIndex.from_arrays([level1, level2]) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.stack().to_pandas() + pd_result = pd_df.stack() + + # Pandas produces NaN, where bq dataframes produces pd.NA + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) diff --git a/tests/system/small/test_numpy.py b/tests/system/small/test_numpy.py new file mode 100644 index 00000000000..fff689cabac --- /dev/null +++ b/tests/system/small/test_numpy.py @@ -0,0 +1,69 @@ +# 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 numpy as np +import pandas as pd +import pytest + + +@pytest.mark.parametrize( + ("opname",), + [ + ("sin",), + ("cos",), + ("tan",), + ("arcsin",), + ("arccos",), + ("arctan",), + ("sinh",), + ("cosh",), + ("tanh",), + ("arcsinh",), + ("arccosh",), + ("arctanh",), + ("exp",), + ("log",), + ("log10",), + ("sqrt",), + ("abs",), + ], +) +def test_series_ufuncs(floats_pd, floats_bf, opname): + bf_result = getattr(np, opname)(floats_bf).to_pandas() + pd_result = getattr(np, opname)(floats_pd) + + pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("opname",), + [ + ("sin",), + ("cos",), + ("tan",), + ("log",), + ("log10",), + ("sqrt",), + ("abs",), + ], +) +def test_df_ufuncs(scalars_dfs, opname): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = getattr(np, opname)( + scalars_df[["float64_col", "int64_col"]] + ).to_pandas() + pd_result = getattr(np, opname)(scalars_pandas_df[["float64_col", "int64_col"]]) + + pd.testing.assert_frame_equal(bf_result, pd_result) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index e40addc4ebd..c60d270fcaa 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -116,6 +116,7 @@ def test_remote_function_direct_no_session_param( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, scalars_dfs, dataset_id_permanent, bq_cf_connection, @@ -126,6 +127,7 @@ def test_remote_function_direct_no_session_param( bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection, # See e2e tests for tests that actually deploy the Cloud Function. @@ -166,6 +168,7 @@ def test_remote_function_direct_no_session_param_location_specified( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, scalars_dfs, dataset_id_permanent, bq_cf_connection_location, @@ -176,6 +179,7 @@ def test_remote_function_direct_no_session_param_location_specified( bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection_location, # See e2e tests for tests that actually deploy the Cloud Function. @@ -213,6 +217,7 @@ def test_remote_function_direct_no_session_param_location_mismatched( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, dataset_id_permanent, bq_cf_connection_location_mismatched, ): @@ -224,6 +229,7 @@ def test_remote_function_direct_no_session_param_location_mismatched( bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection_location_mismatched, # See e2e tests for tests that actually deploy the Cloud Function. @@ -238,6 +244,7 @@ def test_remote_function_direct_no_session_param_location_project_specified( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, scalars_dfs, dataset_id_permanent, bq_cf_connection_location_project, @@ -248,6 +255,7 @@ def test_remote_function_direct_no_session_param_location_project_specified( bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection_location_project, # See e2e tests for tests that actually deploy the Cloud Function. @@ -285,6 +293,7 @@ def test_remote_function_direct_no_session_param_project_mismatched( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, dataset_id_permanent, bq_cf_connection_location_project_mismatched, ): @@ -296,6 +305,7 @@ def test_remote_function_direct_no_session_param_project_mismatched( bigquery_client=bigquery_client, bigquery_connection_client=bigqueryconnection_client, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, dataset=dataset_id_permanent, bigquery_connection=bq_cf_connection_location_project_mismatched, # See e2e tests for tests that actually deploy the Cloud Function. @@ -530,6 +540,7 @@ def test_read_gbq_function_like_original( bigquery_client, bigqueryconnection_client, cloudfunctions_client, + resourcemanager_client, scalars_df_index, dataset_id_permanent, bq_cf_connection, @@ -541,6 +552,7 @@ def test_read_gbq_function_like_original( bigquery_connection_client=bigqueryconnection_client, dataset=dataset_id_permanent, cloud_functions_client=cloudfunctions_client, + resource_manager_client=resourcemanager_client, bigquery_connection=bq_cf_connection, reuse=True, ) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 70c56e5e137..88ad2245c9b 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -459,6 +459,19 @@ def test_mods(scalars_dfs, col_x, col_y, method): pd.testing.assert_series_equal(pd_result, bf_result) +# We work around a pandas bug that doesn't handle correlating nullable dtypes by doing this +# manually with dumb self-correlation instead of parameterized as test_mods is above. +def test_corr(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df["int64_too"].corr(scalars_df["int64_too"]) + pd_result = ( + scalars_pandas_df["int64_too"] + .astype("int64") + .corr(scalars_pandas_df["int64_too"].astype("int64")) + ) + assert math.isclose(pd_result, bf_result) + + @pytest.mark.parametrize( ("col_x",), [ @@ -900,7 +913,7 @@ def test_binop_repeated_application_does_row_identity_joins(scalars_dfs): pd_result, ) - bf_sql, _ = bf_series.to_frame()._to_sql_query(always_include_index=True) + bf_sql, _, _ = bf_series.to_frame()._to_sql_query(include_index=True) selects = re.findall("SELECT", bf_sql.upper()) assert 0 < len(selects) < (num_joins // 2) @@ -2222,8 +2235,9 @@ def test_argmax(scalars_df_index, scalars_pandas_df_index): assert bf_result == pd_result -def test_getattr_not_implemented(scalars_df_index): - with pytest.raises(NotImplementedError): +def test_getattr_attribute_error_when_pandas_has(scalars_df_index): + # asof is implemented in pandas but not in bigframes + with pytest.raises(AttributeError): scalars_df_index.string_col.asof() diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index d825c625617..b7bee16ffd4 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -111,6 +111,54 @@ def test_read_gbq_w_col_order( ["uuid"], id="unique_uuid_index_query", ), + pytest.param( + """ + SELECT my_index, my_value + FROM UNNEST( + [ + STRUCT(0, 12), + STRUCT(1, 12), + STRUCT(2, 24) + ] + ) + -- Can't normally cluster tables with ORDER BY clause. + ORDER BY my_index DESC + """, + ["my_index"], + id="unique_index_query_has_order_by", + ), + pytest.param( + """ + WITH my_table AS ( + SELECT * + FROM UNNEST( + [ + STRUCT(0, 12), + STRUCT(1, 12), + STRUCT(2, 24) + ] + ) + ) + SELECT my_index, my_value FROM my_table + """, + ["my_index"], + id="unique_index_query_with_named_table_expression", + ), + pytest.param( + """ + CREATE TEMP TABLE test_read_gbq_w_index_col_unique_index_query_with_script + AS SELECT * FROM UNNEST( + [ + STRUCT(0, 12), + STRUCT(1, 12), + STRUCT(2, 24) + ] + ); + SELECT my_index, my_value FROM test_read_gbq_w_index_col_unique_index_query_with_script + """, + ["my_index"], + id="unique_index_query_with_script", + ), pytest.param( "{scalars_table_id}", ["bool_col"], @@ -221,7 +269,7 @@ def test_read_gbq_w_max_results( assert bf_result.shape[0] == max_results -def test_read_gbq_w_script(session, dataset_id: str): +def test_read_gbq_w_script_no_select(session, dataset_id: str): ddl = f""" CREATE TABLE `{dataset_id}.test_read_gbq_w_ddl` ( `col_a` INT64, @@ -252,6 +300,20 @@ def test_read_pandas(session, scalars_dfs): pd.testing.assert_frame_equal(result, expected) +def test_read_pandas_col_label_w_space(session: bigframes.Session): + expected = pd.DataFrame( + { + "Animal": ["Falcon", "Falcon", "Parrot", "Parrot"], + "Max Speed": [380.0, 370.0, 24.0, 26.0], + } + ) + result = session.read_pandas(expected).to_pandas() + + pd.testing.assert_frame_equal( + result, expected, check_index_type=False, check_dtype=False + ) + + def test_read_pandas_multi_index(session, scalars_pandas_df_multi_index): df = session.read_pandas(scalars_pandas_df_multi_index) result = df.to_pandas() @@ -755,6 +817,67 @@ def test_read_parquet_gcs(session: bigframes.Session, scalars_dfs, gcs_folder): pd.testing.assert_frame_equal(pd_df_in, pd_df_out) +def test_read_json_gcs_bq_engine(session, scalars_dfs, gcs_folder): + scalars_df, _ = scalars_dfs + path = gcs_folder + "test_read_json_gcs_bq_engine_w_index*.json" + read_path = path.replace("*", FIRST_FILE) + scalars_df.to_json(path, index=False, lines=True, orient="records") + df = session.read_json(read_path, lines=True, orient="records", engine="bigquery") + + # The auto detects of BigQuery load job does not preserve any ordering of columns for json. + pd.testing.assert_index_equal( + df.columns.sort_values(), scalars_df.columns.sort_values() + ) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # datetime, numeric and geometry types, so they're skipped here. + df = df.drop(columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"]) + scalars_df = scalars_df.drop( + columns=["bytes_col", "datetime_col", "numeric_col", "geography_col"] + ) + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal( + df.dtypes.sort_index(), scalars_df.dtypes.sort_index() + ) + + +def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder): + scalars_df, _ = scalars_dfs + path = gcs_folder + "test_read_json_gcs_default_engine_w_index*.json" + read_path = path.replace("*", FIRST_FILE) + scalars_df.to_json( + path, + index=False, + lines=True, + orient="records", + ) + dtype = scalars_df.dtypes.to_dict() + dtype.pop("geography_col") + + df = session.read_json( + read_path, + # Convert default pandas dtypes to match BigQuery DataFrames dtypes. + dtype=dtype, + lines=True, + orient="records", + ) + + assert df._block._expr._ordering is not None + pd.testing.assert_index_equal(df.columns, scalars_df.columns) + + # The auto detects of BigQuery load job have restrictions to detect the bytes, + # numeric and geometry types, so they're skipped here. + df = df.drop(columns=["bytes_col", "numeric_col", "geography_col"]) + scalars_df = scalars_df.drop(columns=["bytes_col", "numeric_col", "geography_col"]) + + # pandas read_json does not respect the dtype overrides for these columns + df = df.drop(columns=["date_col", "datetime_col", "time_col"]) + scalars_df = scalars_df.drop(columns=["date_col", "datetime_col", "time_col"]) + + assert df.shape[0] == scalars_df.shape[0] + pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes) + + def test_session_id(session): assert session._session_id is not None diff --git a/tests/unit/core/test_utils.py b/tests/unit/core/test_utils.py new file mode 100644 index 00000000000..fc34f35d9c2 --- /dev/null +++ b/tests/unit/core/test_utils.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. + +from bigframes.core import utils + + +def test_get_standardized_ids_columns(): + col_labels = ["string", 0, None, "duplicate", "duplicate", "with space"] + + col_ids, idx_ids = utils.get_standardized_ids(col_labels) + + assert col_ids == [ + "string", + "0", + utils.UNNAMED_COLUMN_ID, + "duplicate", + "duplicate.1", + "with_space", + ] + assert idx_ids == [] + + +def test_get_standardized_ids_indexes(): + col_labels = ["duplicate"] + idx_labels = ["string", 0, None, "duplicate", "duplicate", "with space"] + + col_ids, idx_ids = utils.get_standardized_ids(col_labels, idx_labels) + + assert col_ids == ["duplicate.2"] + assert idx_ids == [ + "string", + "0", + utils.UNNAMED_INDEX_ID, + "duplicate", + "duplicate.1", + "with_space", + ] + + +def test_get_standardized_ids_tuple(): + col_labels = [("foo", 1), ("foo", 2), ("bar", 1)] + + col_ids, _ = utils.get_standardized_ids(col_labels) + + assert col_ids == ["('foo',_1)", "('foo',_2)", "('bar',_1)"] diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index d8c8a2d108e..c20a17f7d6e 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -162,3 +162,18 @@ def test_ml_generate_text_produces_correct_sql(): == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_dataset.my_model`, (SELECT * FROM my_table), STRUCT(value AS item))""" ) + + +def test_ml_principal_components_produces_correct_sql(): + sql = ml_sql.ml_principal_components(model_name="my_dataset.my_model") + assert ( + sql == """SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `my_dataset.my_model`)""" + ) + + +def test_ml_principal_component_info_produces_correct_sql(): + sql = ml_sql.ml_principal_component_info(model_name="my_dataset.my_model") + assert ( + sql + == """SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `my_dataset.my_model`)""" + ) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 6762afc61fd..5a812dae7e7 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -15,6 +15,7 @@ import numpy +from bigframes import constants from third_party.bigframes_vendored.pandas.core.generic import NDFrame # ----------------------------------------------------------------------- @@ -33,7 +34,7 @@ class DataFrame(NDFrame): @property def shape(self) -> tuple[int, int]: """Return a tuple representing the dimensionality of the DataFrame.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def axes(self) -> list: @@ -67,7 +68,7 @@ def values(self) -> numpy.ndarray: na_value (default None): The value to use for missing values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # IO methods (to / from other formats) @@ -90,7 +91,7 @@ def to_numpy( Returns: numpy.ndarray: The converted NumPy array. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_gbq( self, @@ -124,7 +125,7 @@ def to_gbq( If set, write the ordering of the DataFrame as a column in the result table with this name. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_parquet( self, @@ -151,7 +152,7 @@ def to_parquet( Returns: None. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Unsorted @@ -179,7 +180,7 @@ def assign(self, **kwargs) -> DataFrame: bigframes.dataframe.DataFrame: A new DataFrame with the new columns in addition to all the existing columns. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Reindexing and alignment @@ -211,7 +212,7 @@ def drop( Raises: KeyError: If any of the labels is not found in the selected axis. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rename( self, @@ -233,7 +234,7 @@ def rename( Raises: KeyError: If any of the labels is not found. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: """ @@ -250,7 +251,7 @@ def rename_axis(self, mapper: Optional[str], **kwargs) -> DataFrame: Returns: bigframes.dataframe.DataFrame: DataFrame with the new index name """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def set_index( self, @@ -273,7 +274,7 @@ def set_index( Returns: DataFrame: Changed row labels. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: """ @@ -287,7 +288,7 @@ def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: Returns: DataFrame: DataFrame of rearranged index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def droplevel(self, level): """ @@ -301,7 +302,7 @@ def droplevel(self, level): Returns: DataFrame: DataFrame with requested index / column level(s) removed. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def reset_index( self, @@ -320,7 +321,7 @@ def reset_index( Returns: bigframes.dataframe.DataFrame: DataFrame with the new index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def drop_duplicates( self, @@ -347,7 +348,7 @@ def drop_duplicates( Returns: bigframes.dataframe.DataFrame: DataFrame with duplicates removed """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def duplicated(self, subset=None, keep="first"): """ @@ -369,7 +370,7 @@ def duplicated(self, subset=None, keep="first"): Returns: bigframes.series.Series: Boolean series for each duplicated rows. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Reindex-based selection methods @@ -382,7 +383,7 @@ def dropna( Returns: bigframes.dataframe.DataFrame: DataFrame with NA entries dropped from it. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Sorting @@ -415,7 +416,7 @@ def sort_values( Returns: DataFrame with sorted values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sort_index( self, @@ -425,7 +426,7 @@ def sort_index( Returns: The original DataFrame sorted by the labels. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Arithmetic Methods @@ -450,7 +451,7 @@ def eq(self, other, axis: str | int = "columns") -> DataFrame: Returns: Result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def ne(self, other, axis: str | int = "columns") -> DataFrame: """ @@ -471,7 +472,7 @@ def ne(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: Result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def le(self, other, axis: str | int = "columns") -> DataFrame: """Get 'less than or equal to' of dataframe and other, element-wise (binary operator `<=`). @@ -497,7 +498,7 @@ def le(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame of bool. The result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def lt(self, other, axis: str | int = "columns") -> DataFrame: """Get 'less than' of DataFrame and other, element-wise (binary operator `<`). @@ -523,7 +524,7 @@ def lt(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame of bool. The result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def ge(self, other, axis: str | int = "columns") -> DataFrame: """Get 'greater than or equal to' of DataFrame and other, element-wise (binary operator `>=`). @@ -549,7 +550,7 @@ def ge(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame of bool. The result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def gt(self, other, axis: str | int = "columns") -> DataFrame: """Get 'greater than' of DataFrame and other, element-wise (binary operator `>`). @@ -575,7 +576,7 @@ def gt(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame of bool: The result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def add(self, other, axis: str | int = "columns") -> DataFrame: """Get addition of DataFrame and other, element-wise (binary operator `+`). @@ -598,7 +599,7 @@ def add(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sub(self, other, axis: str | int = "columns") -> DataFrame: """Get subtraction of DataFrame and other, element-wise (binary operator `-`). @@ -621,7 +622,7 @@ def sub(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rsub(self, other, axis: str | int = "columns") -> DataFrame: """Get subtraction of DataFrame and other, element-wise (binary operator `-`). @@ -644,7 +645,7 @@ def rsub(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mul(self, other, axis: str | int = "columns") -> DataFrame: """Get multiplication of DataFrame and other, element-wise (binary operator `*`). @@ -667,7 +668,7 @@ def mul(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def truediv(self, other, axis: str | int = "columns") -> DataFrame: """Get floating division of DataFrame and other, element-wise (binary operator `/`). @@ -690,7 +691,7 @@ def truediv(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: """Get floating division of DataFrame and other, element-wise (binary operator `/`). @@ -713,7 +714,7 @@ def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def floordiv(self, other, axis: str | int = "columns") -> DataFrame: """Get integer division of DataFrame and other, element-wise (binary operator `//`). @@ -736,7 +737,7 @@ def floordiv(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rfloordiv(self, other, axis: str | int = "columns") -> DataFrame: """Get integer division of DataFrame and other, element-wise (binary operator `//`). @@ -759,7 +760,7 @@ def rfloordiv(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mod(self, other, axis: str | int = "columns") -> DataFrame: """Get modulo of DataFrame and other, element-wise (binary operator `%`). @@ -782,7 +783,7 @@ def mod(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rmod(self, other, axis: str | int = "columns") -> DataFrame: """Get modulo of DataFrame and other, element-wise (binary operator `%`). @@ -805,7 +806,7 @@ def rmod(self, other, axis: str | int = "columns") -> DataFrame: Returns: DataFrame: DataFrame result of the arithmetic operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Data reshaping @@ -847,7 +848,7 @@ def groupby( Returns: bigframes.core.groupby.SeriesGroupBy: A groupby object that contains information about the groups. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Function application @@ -874,7 +875,7 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: Returns: bigframes.dataframe.DataFrame: Transformed DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Merging / joining methods @@ -902,7 +903,7 @@ def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame: Returns: bigframes.dataframe.DataFrame: A dataframe containing columns from both the caller and `other`. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def merge( self, @@ -972,7 +973,7 @@ def merge( Returns: bigframes.dataframe.DataFrame: A DataFrame of the two merged objects. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # ndarray-like stats methods @@ -992,7 +993,7 @@ def any(self, *, bool_only: bool = False): Returns: Series """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def all(self, *, bool_only: bool = False): """ @@ -1009,7 +1010,7 @@ def all(self, *, bool_only: bool = False): Returns: bigframes.series.Series: Series if all elements are True. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def prod(self, *, numeric_only: bool = False): """ @@ -1022,7 +1023,7 @@ def prod(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with the product of the values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def min(self, *, numeric_only: bool = False): """Return the minimum of the values over the requested axis. @@ -1037,7 +1038,7 @@ def min(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with the minimum of the values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def max(self, *, numeric_only: bool = False): """Return the maximum of the values over the requested axis. @@ -1052,7 +1053,7 @@ def max(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series after the maximum of values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sum(self, *, numeric_only: bool = False): """Return the sum of the values over the requested axis. @@ -1066,7 +1067,7 @@ def sum(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with the sum of values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mean(self, *, numeric_only: bool = False): """Return the mean of the values over the requested axis. @@ -1078,7 +1079,7 @@ def mean(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with the mean of values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def median(self, *, numeric_only: bool = False, exact: bool = False): """Return the median of the values over the requested axis. @@ -1093,7 +1094,7 @@ def median(self, *, numeric_only: bool = False, exact: bool = False): Returns: bigframes.series.Series: Series with the median of values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def var(self, *, numeric_only: bool = False): """Return unbiased variance over requested axis. @@ -1107,7 +1108,7 @@ def var(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with unbiased variance over requested axis. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def std(self, *, numeric_only: bool = False): """Return sample standard deviation over requested axis. @@ -1121,7 +1122,7 @@ def std(self, *, numeric_only: bool = False): Returns: bigframes.series.Series: Series with sample standard deviation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def count(self, *, numeric_only: bool = False): """ @@ -1138,7 +1139,7 @@ def count(self, *, numeric_only: bool = False): bigframes.series.Series: For each column/row the number of non-NA/null entries. If `level` is specified returns a `DataFrame`. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def nunique(self): """ @@ -1147,7 +1148,7 @@ def nunique(self): Returns: bigframes.series.Series: Series with number of distinct elements. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummin(self) -> DataFrame: """Return cumulative minimum over a DataFrame axis. @@ -1157,7 +1158,7 @@ def cummin(self) -> DataFrame: Returns: bigframes.dataframe.DataFrame: Return cumulative minimum of DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummax(self) -> DataFrame: """Return cumulative maximum over a DataFrame axis. @@ -1167,7 +1168,7 @@ def cummax(self) -> DataFrame: Returns: bigframes.dataframe.DataFrame: Return cumulative maximum of DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumsum(self) -> DataFrame: """Return cumulative sum over a DataFrame axis. @@ -1177,7 +1178,7 @@ def cumsum(self) -> DataFrame: Returns: bigframes.dataframe.DataFrame: Return cumulative sum of DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumprod(self) -> DataFrame: """Return cumulative product over a DataFrame axis. @@ -1187,7 +1188,7 @@ def cumprod(self) -> DataFrame: Returns: bigframes.dataframe.DataFrame: Return cumulative product of DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def agg(self, func): """ @@ -1202,7 +1203,7 @@ def agg(self, func): Returns: DataFrame or bigframes.series.Series: Aggregated results. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def describe(self): """ @@ -1227,7 +1228,7 @@ def describe(self): Returns: bigframes.dataframe.DataFrame: Summary statistics of the Series or Dataframe provided. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def pivot(self, *, columns, index=None, values=None): """ @@ -1263,7 +1264,30 @@ def pivot(self, *, columns, index=None, values=None): Returns: Returns reshaped DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def stack(self): + """ + Stack the prescribed level(s) from columns to index. + + Return a reshaped DataFrame or Series having a multi-level + index with one or more new inner-most levels compared to the current + DataFrame. The new inner-most levels are created by pivoting the + columns of the current dataframe: + + - if the columns have a single level, the output is a Series; + - if the columns have multiple levels, the new index + level(s) is (are) taken from the prescribed level(s) and + the output is a DataFrame. + + .. note:: + BigQuery DataFrames does not support stack operations that would + combine columns of different dtypes. + + Returns: + DataFrame or Series: Stacked dataframe or series. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Add index and columns @@ -1280,12 +1304,12 @@ def index(self): Returns: The index labels of the DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def columns(self): "The column labels of the DataFrame." - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def value_counts( self, @@ -1313,4 +1337,4 @@ def value_counts( Returns: Series: Series containing counts of unique rows in the DataFrame """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 4843c971da2..56d3b2434f0 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -3,6 +3,7 @@ from typing import Literal, Optional +from bigframes import constants from third_party.bigframes_vendored.pandas.core import indexing @@ -22,7 +23,7 @@ def ndim(self) -> int: Returns: int: Return 1 if Series. Otherwise return 2 if DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def size(self) -> int: @@ -32,7 +33,7 @@ def size(self) -> int: int: Return the number of rows if Series. Otherwise return the number of rows times number of columns if DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ------------------------------------------------------------------------- # Unary Methods @@ -46,7 +47,7 @@ def abs(self): Series/DataFrame containing the absolute value of each element. Returns a Series/DataFrame containing the absolute value of each element. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def astype(self, dtype): """ @@ -66,7 +67,7 @@ def astype(self, dtype): same type as caller """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Iteration @@ -85,7 +86,7 @@ def empty(self) -> bool: Returns: bool: If Series/DataFrame is empty, return True, if not return False. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # I/O Methods @@ -155,7 +156,7 @@ def to_json( Returns: None: String output not yet supported. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: """Write object to a comma-separated values (csv) file on Cloud Storage. @@ -177,7 +178,7 @@ def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: Returns: None: String output not yet supported. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Unsorted @@ -215,7 +216,7 @@ def add_prefix(self, prefix: str, axis: int | str | None = None): Returns: New Series or DataFrame with updated labels. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def add_suffix(self, suffix: str, axis: int | str | None = None): """Suffix labels with string `suffix`. @@ -233,7 +234,7 @@ def add_suffix(self, suffix: str, axis: int | str | None = None): Returns: New Series or DataFrame with updated labels. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def head(self, n: int = 5): """Return the first `n` rows. @@ -254,7 +255,7 @@ def head(self, n: int = 5): Returns: The first `n` rows of the caller object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def tail(self, n: int = 5): """Return the last `n` rows. @@ -275,7 +276,7 @@ def tail(self, n: int = 5): Returns: The last `n` rows of the caller object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sample( self, @@ -301,7 +302,7 @@ def sample( A new object of same type as caller containing `n` items randomly sampled from the caller object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Internal Interface Methods @@ -317,7 +318,7 @@ def dtypes(self): Returns: A *pandas* Series with the data type of each column. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def copy(self): """Make a copy of this object's indices and data. @@ -329,7 +330,7 @@ def copy(self): Returns: Object type matches caller. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # Action Methods @@ -346,7 +347,7 @@ def isna(self) -> NDFrame: Mask of bool values for each element that indicates whether an element is an NA value. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) isnull = isna @@ -362,7 +363,7 @@ def notna(self) -> NDFrame: NDFrame: Mask of bool values for each element that indicates whether an element is not an NA value. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) notnull = notna @@ -381,7 +382,7 @@ def shift( Returns: NDFrame: Copy of input object, shifted. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rank( self, @@ -419,7 +420,7 @@ def rank( Returns: same type as caller: Return a Series or DataFrame with data ranks as values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def __nonzero__(self): raise ValueError( diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index 95822718c3a..9271da8a5ea 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -9,6 +9,8 @@ class providing the base-class of operations. """ from __future__ import annotations +from bigframes import constants + class GroupBy: """ @@ -24,7 +26,7 @@ def any(self): where a value is True if any element is True within its respective group, False otherwise. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def all(self): """ @@ -35,7 +37,7 @@ def all(self): where a value is True if all elements are True within its respective group, False otherwise. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def count(self): """ @@ -44,7 +46,7 @@ def count(self): Returns: Series or DataFrame: Count of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mean( self, @@ -60,7 +62,7 @@ def mean( Returns: pandas.Series or pandas.DataFrame: Mean of groups. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def median( self, @@ -81,7 +83,7 @@ def median( Returns: pandas.Series or pandas.DataFrame: Median of groups. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def std( self, @@ -100,7 +102,7 @@ def std( Returns: Series or DataFrame: Standard deviation of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def var( self, @@ -120,7 +122,7 @@ def var( Series or DataFrame Variance of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sum( self, @@ -140,7 +142,7 @@ def sum( Returns: Series or DataFrame: Computed sum of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def prod(self, numeric_only: bool = False, min_count: int = 0): """ @@ -156,7 +158,7 @@ def prod(self, numeric_only: bool = False, min_count: int = 0): Returns: Series or DataFrame: Computed prod of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def min( self, @@ -176,7 +178,7 @@ def min( Returns: Series or DataFrame: Computed min of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def max( self, @@ -196,7 +198,7 @@ def max( Returns: Series or DataFrame: Computed max of values within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumcount(self, ascending: bool = True): """ @@ -209,7 +211,7 @@ def cumcount(self, ascending: bool = True): Returns: Series: Sequence number of each element within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumprod(self, *args, **kwargs): """ @@ -218,7 +220,7 @@ def cumprod(self, *args, **kwargs): Returns: Series or DataFrame: Cumulative product for each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumsum(self, *args, **kwargs): """ @@ -227,7 +229,7 @@ def cumsum(self, *args, **kwargs): Returns: Series or DataFrame: Cumulative sum for each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummin(self, *args, numeric_only: bool = False, **kwargs): """ @@ -236,7 +238,7 @@ def cummin(self, *args, numeric_only: bool = False, **kwargs): Returns: Series or DataFrame: Cumulative min for each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummax(self, *args, numeric_only: bool = False, **kwargs): """ @@ -245,7 +247,7 @@ def cummax(self, *args, numeric_only: bool = False, **kwargs): Returns: Series or DataFrame: Cumulative max for each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def diff(self): """ @@ -256,7 +258,7 @@ def diff(self): Returns: Series or DataFrame: First differences. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def shift(self, periods: int = 1): """ @@ -269,7 +271,7 @@ def shift(self, periods: int = 1): Returns: Series or DataFrame: Object shifted within each group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rolling(self, *args, **kwargs): """ @@ -289,7 +291,7 @@ def rolling(self, *args, **kwargs): Returns: Series or DataFrame: Return a new grouper with our rolling appended. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def expanding(self, *args, **kwargs): """ @@ -298,7 +300,7 @@ def expanding(self, *args, **kwargs): Returns: Series or DataFrame: A expanding grouper, providing expanding functionality per group. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class SeriesGroupBy(GroupBy): @@ -318,7 +320,7 @@ def agg(self, func): Returns: Series or DataFrame """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class DataFrameGroupBy(GroupBy): @@ -347,4 +349,4 @@ def agg(self, func, **kwargs): Returns: DataFrame """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/indexes/accessor.py b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py index d59886e8aad..2b4a3263175 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/accessor.py @@ -1,3 +1,6 @@ +from bigframes import constants + + class DatetimeProperties: """ Accessor object for datetime-like properties of the Series values. @@ -7,7 +10,7 @@ class DatetimeProperties: def day(self): """The day of the datetime.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def dayofweek(self): @@ -22,7 +25,7 @@ def dayofweek(self): Series or Index: Containing integers indicating the day number. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def date(self): @@ -36,31 +39,31 @@ def date(self): a numpy array. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def hour(self): """The hours of the datetime.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def minute(self): """The minutes of the datetime.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def month(self): """The month as January=1, December=12.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def second(self): """The seconds of the datetime.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def time(self): @@ -73,7 +76,7 @@ def time(self): a numpy array. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def quarter(self): @@ -84,10 +87,10 @@ def quarter(self): a numpy array. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def year(self): """The year of the datetime.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py index ebad5eb9181..404a99809c0 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/base.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -1,5 +1,7 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexes/base.py +from bigframes import constants + class Index: """Immutable sequence used for indexing and alignment. @@ -10,14 +12,14 @@ class Index: @property def name(self): """Returns Index name.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def shape(self): """ Return a tuple of the shape of the underlying data. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_numpy(self, dtype): """ @@ -33,4 +35,4 @@ def to_numpy(self, dtype): Returns: numpy.ndarray """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/indexing.py b/third_party/bigframes_vendored/pandas/core/indexing.py index d5b9f3c079f..fae5d6261f2 100644 --- a/third_party/bigframes_vendored/pandas/core/indexing.py +++ b/third_party/bigframes_vendored/pandas/core/indexing.py @@ -1,5 +1,7 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexing.py +from bigframes import constants + class IndexingMixin: """ @@ -32,7 +34,7 @@ def iloc(self): out-of-bounds, except *slice* indexers which allow out-of-bounds indexing (this conforms with python/numpy *slice* semantics). """ - raise NotImplementedError("abstract methdod") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def loc(self): @@ -63,4 +65,4 @@ def loc(self): NotImplementError: if the inputs are not supported. """ - raise NotImplementedError("abstract methdod") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/reshape/concat.py b/third_party/bigframes_vendored/pandas/core/reshape/concat.py index 6a5a9fdde92..6e6d2d8b5c9 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/concat.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/concat.py @@ -4,6 +4,8 @@ """ from __future__ import annotations +from bigframes import constants + def concat( objs, @@ -135,4 +137,4 @@ def concat( [4 rows x 2 columns] """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/reshape/tile.py b/third_party/bigframes_vendored/pandas/core/reshape/tile.py index 9381ad4552a..4f5f2efef0a 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/tile.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/tile.py @@ -4,6 +4,8 @@ """ from __future__ import annotations +from bigframes import constants + def cut( x, @@ -62,4 +64,4 @@ def cut( are whatever the type in the sequence is. False : returns an ndarray of integers. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 79eb4026968..8d505c1eadb 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -10,6 +10,7 @@ from pandas._typing import Axis, FilePath, NaPosition, WriteBuffer import pandas.io.formats.format as fmt +from bigframes import constants from third_party.bigframes_vendored.pandas.core.generic import NDFrame if TYPE_CHECKING: @@ -23,31 +24,31 @@ def dt(self): """ Accessor object for datetime-like properties of the Series values. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def index(self): """The index (axis labels) of the Series.""" - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def shape(self): """Return a tuple of the shape of the underlying data.""" - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def dtype(self): """ Return the dtype object of the underlying data. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def dtypes(self): """ Return the dtype object of the underlying data. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def name(self) -> Hashable: @@ -62,7 +63,7 @@ def name(self) -> Hashable: hashable object: The name of the Series, also the column name if part of a DataFrame. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def reset_index( self, @@ -94,13 +95,13 @@ def reset_index( In either case, if ``inplace=True``, no value is returned. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def __repr__(self) -> str: """ Return a string representation for a particular Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # IO methods (to / from other formats) @@ -165,7 +166,7 @@ def to_string( result = formatter.to_string() # catch contract violations - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_markdown( self, @@ -188,7 +189,7 @@ def to_markdown( Returns: str: {klass} in Markdown-friendly format. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_dict(self, into: type[dict] = dict) -> Mapping: """ @@ -204,7 +205,7 @@ def to_dict(self, into: type[dict] = dict) -> Mapping: Returns: collections.abc.Mapping: Key-value representation of Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_frame(self) -> DataFrame: """ @@ -213,7 +214,7 @@ def to_frame(self) -> DataFrame: Returns: bigframes.dataframe.DataFrame: DataFrame representation of Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_excel(self, excel_writer, sheet_name): """ @@ -235,7 +236,7 @@ def to_excel(self, excel_writer, sheet_name): sheet_name (str, default 'Sheet1'): Name of sheet to contain Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_latex(self, buf=None, columns=None, header=True, index=True, **kwargs): """ @@ -256,7 +257,7 @@ def to_latex(self, buf=None, columns=None, header=True, index=True, **kwargs): str or None: If buf is None, returns the result as a string. Otherwise returns None. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def tolist(self) -> list: """ @@ -269,7 +270,7 @@ def tolist(self) -> list: Returns: list: list of the values """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) to_list = tolist @@ -296,7 +297,7 @@ def to_numpy(self, dtype, copy=False, na_value=None): numpy.ndarray: A NumPy ndarray representing the values in this Series or Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_pickle(self, path, **kwargs): """ @@ -308,7 +309,7 @@ def to_pickle(self, path, **kwargs): object implementing a binary ``write()`` function. File path where the pickled object will be stored. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_xarray(self): """ @@ -319,7 +320,7 @@ def to_xarray(self): converted to Dataset if the object is a DataFrame, or a DataArray if the object is a Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_json( self, @@ -354,7 +355,7 @@ def to_json( None or str: If path_or_buf is None, returns the resulting json format as a string. Otherwise returns None. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: """ @@ -372,7 +373,7 @@ def to_csv(self, path_or_buf: str, *, index: bool = True) -> str | None: None or str: If path_or_buf is None, returns the resulting csv format as a string. Otherwise returns None. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def agg(self, func): """ @@ -387,7 +388,7 @@ def agg(self, func): Returns: scalar or Series: Aggregated results """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def count(self): """ @@ -397,7 +398,7 @@ def count(self): int or Series (if level specified): Number of non-null values in the Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def nunique(self) -> int: """ @@ -408,7 +409,7 @@ def nunique(self) -> int: Returns: int: number of unique elements in the object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mode(self) -> Series: """ @@ -421,7 +422,7 @@ def mode(self) -> Series: Returns: bigframes.series.Series: Modes of the Series in sorted order. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def drop_duplicates( self, @@ -442,7 +443,7 @@ def drop_duplicates( Returns: bigframes.series.Series: Series with duplicates dropped or None if ``inplace=True``. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def duplicated(self, keep="first") -> Series: """ @@ -466,7 +467,7 @@ def duplicated(self, keep="first") -> Series: bigframes.series.Series: Series indicating whether each value has occurred in the preceding values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def round(self, decimals: int = 0) -> Series: """ @@ -480,6 +481,29 @@ def round(self, decimals: int = 0) -> Series: Returns: bigframes.series.Series: Rounded values of the Series. """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def corr(self, other, method="pearson", min_periods=None) -> float: + """ + Compute the correlation with the other Series. Non-number values are ignored in the + computation. + + Uses the "Pearson" method of correlation. Numbers are converted to float before + calculation, so the result may be unstable. + + Args: + other (Series): + The series with which this is to be correlated. + method (string, default "pearson"): + Correlation method to use - currently only "pearson" is supported. + min_periods (int, default None): + The minimum number of observations needed to return a result. Non-default values + are not yet supported, so a result will be returned for at least two observations. + + Returns: + float; Will return NaN if there are fewer than two numeric pairs, either series has a + variance or covariance of zero, or any input value is infinite. + """ raise NotImplementedError("abstract method") def diff(self) -> Series: @@ -497,7 +521,7 @@ def diff(self) -> Series: Returns: {klass}: First differences of the Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def dot(self, other) -> Series | np.ndarray: """ @@ -527,19 +551,19 @@ def dot(self, other) -> Series | np.ndarray: """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def __matmul__(self, other): """ Matrix multiplication using binary `@` operator in Python>=3.5. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def __rmatmul__(self, other): """ Matrix multiplication using binary `@` operator in Python>=3.5. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sort_values( self, @@ -571,7 +595,7 @@ def sort_values( Returns: bigframes.series.Series: Series ordered by values or None if ``inplace=True``. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sort_index( self, @@ -602,7 +626,7 @@ def sort_index( """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def nlargest( self, n: int = 5, keep: Literal["first", "last", "all"] = "first" @@ -626,7 +650,7 @@ def nlargest( Returns: bigframes.series.Series: The `n` largest values in the Series, sorted in decreasing order. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def nsmallest(self, n: int = 5, keep: str = "first") -> Series: """ @@ -649,7 +673,7 @@ def nsmallest(self, n: int = 5, keep: str = "first") -> Series: Returns: bigframes.series.Series: The `n` smallest values in the Series, sorted in increasing order. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) # ---------------------------------------------------------------------- # function application @@ -672,7 +696,7 @@ def apply( bigframes.series.Series: If func returns a Series object the result will be a DataFrame. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def groupby( self, @@ -726,7 +750,7 @@ def groupby( bigframes.core.groupby.SeriesGroupBy: Returns a groupby object that contains information about the groups. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def drop( self, labels=None, *, axis=0, index=None, columns=None, level=None @@ -758,7 +782,7 @@ def drop( Raises: KeyError: If none of the labels are found in the index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def reorder_levels(self, order: Sequence) -> Series: """ @@ -773,7 +797,7 @@ def reorder_levels(self, order: Sequence) -> Series: Returns: type of caller (new object) """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def droplevel(self, level): """ @@ -788,7 +812,7 @@ def droplevel(self, level): Returns: Series with requested index / column level(s) removed. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def fillna( self, @@ -804,7 +828,7 @@ def fillna( Returns: Series or None: Object with missing values filled or None. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def between( self, @@ -832,7 +856,7 @@ def between( right (inclusive). """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumprod(self): """ @@ -844,7 +868,7 @@ def cumprod(self): Returns: bigframes.series.Series: Return cumulative sum of scalar or Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumsum(self): """ @@ -861,7 +885,7 @@ def cumsum(self): Returns: scalar or Series: Return cumulative sum of scalar or Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummax(self): """ @@ -878,7 +902,7 @@ def cummax(self): Returns: bigframes.series.Series: Return cumulative maximum of scalar or Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummin(self): """ @@ -901,7 +925,7 @@ def cummin(self): Returns: bigframes.series.Series: Return cumulative minimum of scalar or Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def eq(self, other) -> Series: """Return equal of Series and other, element-wise (binary operator eq). @@ -916,7 +940,7 @@ def eq(self, other) -> Series: Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def ne(self, other) -> Series: """Return not equal of Series and other, element-wise (binary operator ne). @@ -931,7 +955,7 @@ def ne(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def le(self, other) -> Series: """Get 'less than or equal to' of Series and other, element-wise (binary operator `<=`). @@ -946,7 +970,7 @@ def le(self, other) -> Series: bigframes.series.Series. The result of the comparison. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def lt(self, other) -> Series: """Get 'less than' of Series and other, element-wise (binary operator `<`). @@ -961,7 +985,7 @@ def lt(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def ge(self, other) -> Series: """Get 'greater than or equal to' of Series and other, element-wise (binary operator `>=`). @@ -976,7 +1000,7 @@ def ge(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def gt(self, other) -> Series: """Get 'less than or equal to' of Series and other, element-wise (binary operator `<=`). @@ -991,7 +1015,7 @@ def gt(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def add(self, other) -> Series: """Return addition of Series and other, element-wise (binary operator add). @@ -1006,7 +1030,7 @@ def add(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def radd(self, other) -> Series: """Return addition of Series and other, element-wise (binary operator radd). @@ -1021,7 +1045,7 @@ def radd(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sub( self, @@ -1039,7 +1063,7 @@ def sub( bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rsub(self, other) -> Series: """Return subtraction of Series and other, element-wise (binary operator rsub). @@ -1054,7 +1078,7 @@ def rsub(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mul(self, other) -> Series: """Return multiplication of Series and other, element-wise (binary operator mul). @@ -1069,7 +1093,7 @@ def mul(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rmul(self, other) -> Series: """Return multiplication of Series and other, element-wise (binary operator mul). @@ -1083,7 +1107,7 @@ def rmul(self, other) -> Series: Returns: Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def truediv(self, other) -> Series: """Return floating division of Series and other, element-wise (binary operator truediv). @@ -1098,7 +1122,7 @@ def truediv(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rtruediv(self, other) -> Series: """Return floating division of Series and other, element-wise (binary operator rtruediv). @@ -1113,7 +1137,7 @@ def rtruediv(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def floordiv(self, other) -> Series: """Return integer division of Series and other, element-wise (binary operator floordiv). @@ -1128,7 +1152,7 @@ def floordiv(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rfloordiv(self, other) -> Series: """Return integer division of Series and other, element-wise (binary operator rfloordiv). @@ -1143,7 +1167,7 @@ def rfloordiv(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mod(self, other) -> Series: """Return modulo of Series and other, element-wise (binary operator mod). @@ -1158,7 +1182,7 @@ def mod(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rmod(self, other) -> Series: """Get modulo of Series and other, element-wise (binary operator `rmod`). @@ -1173,7 +1197,7 @@ def rmod(self, other) -> Series: bigframes.series.Series: The result of the operation. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def divmod(self, other) -> Series: """Return integer division and modulo of Series and other, element-wise (binary operator divmod). @@ -1188,7 +1212,7 @@ def divmod(self, other) -> Series: consistent with (floordiv, mod) (though pandas may not). """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rdivmod(self, other) -> Series: """Return integer division and modulo of Series and other, element-wise (binary operator rdivmod). @@ -1203,7 +1227,7 @@ def rdivmod(self, other) -> Series: consistent with (rfloordiv, rmod) (though pandas may not). """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def all( self, @@ -1218,7 +1242,7 @@ def all( scalar or Series: If level is specified, then, Series is returned; otherwise, scalar is returned. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def any( self, @@ -1233,7 +1257,7 @@ def any( scalar or Series: If level is specified, then, Series is returned; otherwise, scalar is returned. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def max( self, @@ -1248,7 +1272,7 @@ def max( Returns: scalar or scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def min( self, @@ -1262,7 +1286,7 @@ def min( Returns: scalar or scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def std( self, @@ -1277,7 +1301,7 @@ def std( ------- scalar or Series (if level specified) """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def var( self, @@ -1290,7 +1314,7 @@ def var( Returns: scalar or Series (if level specified) """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sum(self): """Return the sum of the values over the requested axis. @@ -1300,7 +1324,7 @@ def sum(self): Returns: scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mean(self): """Return the mean of the values over the requested axis. @@ -1308,7 +1332,7 @@ def mean(self): Returns: scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def median(self, *, exact: bool = False): """Return the median of the values over the requested axis. @@ -1321,7 +1345,7 @@ def median(self, *, exact: bool = False): Returns: scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def prod(self): """Return the product of the values over the requested axis. @@ -1329,7 +1353,7 @@ def prod(self): Returns: scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def skew(self): """Return unbiased skew over requested axis. @@ -1339,7 +1363,7 @@ def skew(self): Returns: scalar """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def kurt(self): """Return unbiased kurtosis over requested axis. @@ -1349,7 +1373,7 @@ def kurt(self): Returns: scalar or scalar: Unbiased kurtosis over requested axis. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def where(self, cond, other): """Replace values where the condition is False. @@ -1373,7 +1397,7 @@ def where(self, cond, other): Returns: bigframes.series.Series: Series after the replacement. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mask(self, cond, other): """Replace values where the condition is True. @@ -1397,7 +1421,7 @@ def mask(self, cond, other): Returns: bigframes.series.Series: Series after the replacement. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def clip(self): """Trim values at input threshold(s). @@ -1416,7 +1440,7 @@ def clip(self): Returns: Series. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def argmax(self): """ @@ -1427,7 +1451,7 @@ def argmax(self): Returns: Series: Row position of the maximum value. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def argmin(self): """ @@ -1438,7 +1462,7 @@ def argmin(self): Returns: Series: Row position of the minimum value. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rename(self, index, **kwargs) -> Series | None: """ @@ -1461,7 +1485,7 @@ def rename(self, index, **kwargs) -> Series | None: bigframes.series.Series: Series with index labels. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rename_axis(self, mapper, **kwargs): """ @@ -1474,7 +1498,7 @@ def rename_axis(self, mapper, **kwargs): Returns: bigframes.series.Series: Series with the name of the axis set. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rolling( self, @@ -1515,7 +1539,7 @@ def rolling( bigframes.core.window.Window: ``Window`` subclass if a ``win_type`` is passed. ``Rolling`` subclass if ``win_type`` is not passed. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def expanding(self, min_periods=1): """ @@ -1529,7 +1553,7 @@ def expanding(self, min_periods=1): Returns: bigframes.core.window.Window: ``Expanding`` subclass. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def value_counts( self, @@ -1560,7 +1584,7 @@ def value_counts( Returns: Series: Series containing counts of unique values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def str(self): @@ -1570,7 +1594,7 @@ def str(self): NAs stay NA unless handled otherwise by a particular method. Patterned after Python’s string methods, with some inspiration from R’s stringr package. """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isin(self, values): """ @@ -1595,7 +1619,7 @@ def isin(self, values): Raises: TypeError: If input is not list-like. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def is_monotonic_increasing(self) -> bool: @@ -1605,7 +1629,7 @@ def is_monotonic_increasing(self) -> bool: Returns: bool """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def is_monotonic_decreasing(self) -> bool: @@ -1615,4 +1639,4 @@ def is_monotonic_decreasing(self) -> bool: Returns: bool """ - raise NotImplementedError("abstract property") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/strings/accessor.py b/third_party/bigframes_vendored/pandas/core/strings/accessor.py index 5e3d0b047ff..ecdd9547d54 100644 --- a/third_party/bigframes_vendored/pandas/core/strings/accessor.py +++ b/third_party/bigframes_vendored/pandas/core/strings/accessor.py @@ -1,6 +1,8 @@ import re import typing +from bigframes import constants + class StringMethods: """ @@ -32,7 +34,7 @@ def extract(self, pat: str, flags: int = 0): expression pat will be used for column names; otherwise capture group numbers will be used. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def find(self, sub, start: int = 0, end=None): """Return lowest indexes in each strings in the Series/Index. @@ -52,7 +54,7 @@ def find(self, sub, start: int = 0, end=None): Returns: bigframes.series.Series: Series with lowest indexes in each strings. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def len(self): """Compute the length of each element in the Series/Index. @@ -65,7 +67,7 @@ def len(self): the length of each element in the Series or Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def lower(self): """Convert strings in the Series/Index to lowercase. @@ -76,7 +78,7 @@ def lower(self): bigframes.series.Series: Series with lowercase. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def slice(self, start=None, stop=None): """Slice substrings from each element in the Series or Index. @@ -94,7 +96,7 @@ def slice(self, start=None, stop=None): substring from original string object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def strip(self): """Remove leading and trailing characters. @@ -109,7 +111,7 @@ def strip(self): and trailing characters. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def upper(self): """Convert strings in the Series/Index to uppercase. @@ -120,7 +122,7 @@ def upper(self): bigframes.series.Series: Series with uppercase strings. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isnumeric(self): """Check whether all characters in each string are numeric. @@ -134,7 +136,7 @@ def isnumeric(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isalpha(self): """Check whether all characters in each string are alphabetic. @@ -147,7 +149,7 @@ def isalpha(self): bigframes.series.Series: Series with the same length as the originalSeries/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isdigit(self): """Check whether all characters in each string are digits. @@ -160,7 +162,7 @@ def isdigit(self): bigframes.series.Series: Series with the same length as the originalSeries/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isalnum(self): """Check whether all characters in each string are alphanumeric. @@ -174,7 +176,7 @@ def isalnum(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isspace(self): """Check whether all characters in each string are whitespace. @@ -188,7 +190,7 @@ def isspace(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def islower(self): """Check whether all characters in each string are lowercase. @@ -202,7 +204,7 @@ def islower(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isupper(self): """Check whether all characters in each string are uppercase. @@ -216,7 +218,7 @@ def isupper(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def isdecimal(self): """Check whether all characters in each string are decimal. @@ -230,7 +232,7 @@ def isdecimal(self): same length as the original Series/Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rstrip(self): """Remove trailing characters. @@ -244,7 +246,7 @@ def rstrip(self): bigframes.series.Series: Series without trailing characters. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def lstrip(self): """Remove leading characters. @@ -258,7 +260,7 @@ def lstrip(self): bigframes.series.Series: Series without leading characters. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def repeat(self, repeats: int): """Duplicate each string in the Series or Index. @@ -272,7 +274,7 @@ def repeat(self, repeats: int): objects specified by input parameter repeats. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def capitalize(self): """Convert strings in the Series/Index to be capitalized. @@ -283,7 +285,7 @@ def capitalize(self): bigframes.series.Series: Series with captitalized strings. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cat(self, others, *, join): """Concatenate strings in the Series/Index with given separator. @@ -304,7 +306,7 @@ def cat(self, others, *, join): bigframes.series.Series: Series with concatenated strings. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def contains(self, pat, case: bool = True, flags: int = 0, *, regex: bool = True): """ @@ -329,7 +331,7 @@ def contains(self, pat, case: bool = True, flags: int = 0, *, regex: bool = True whether the given pattern is contained within the string of each element of the Series or Index. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def replace( self, @@ -373,7 +375,7 @@ def replace( of `pat` replaced by `repl`. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def startswith( self, @@ -391,7 +393,7 @@ def startswith( bigframes.series.Series: A Series of booleans indicating whether the given pattern matches the start of each string element. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def endswith( self, @@ -409,7 +411,7 @@ def endswith( bigframes.series.Series: A Series of booleans indicating whether the given pattern matches the end of each string element. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def match(self, pat: str, case: bool = True, flags: int = 0): """ @@ -426,7 +428,7 @@ def match(self, pat: str, case: bool = True, flags: int = 0): Returns: bigframes.series.Series: Series of boolean values """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def fullmatch(self, pat: str, case: bool = True, flags: int = 0): """ @@ -443,7 +445,7 @@ def fullmatch(self, pat: str, case: bool = True, flags: int = 0): Returns: bigframes.series.Series: Series of boolean values """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def get(self, i: int): """ @@ -459,7 +461,7 @@ def get(self, i: int): Returns: bigframes.series.Series: Series """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def pad( self, @@ -482,7 +484,7 @@ def pad( Returns: bigframes.series.Series: Returns Series or Index with minimum number of char in object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def ljust( self, @@ -502,7 +504,7 @@ def ljust( Returns: bigframes.series.Series: Returns Series or Index with minimum number of char in object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rjust( self, @@ -522,7 +524,7 @@ def rjust( Returns: bigframes.series.Series: Returns Series or Index with minimum number of char in object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def zfill( self, @@ -544,7 +546,7 @@ def zfill( Returns: bigframes.series.Series: Series of objects. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def center( self, @@ -566,4 +568,4 @@ def center( Returns: bigframes.series.Series: Returns Series or Index with minimum number of char in object. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/window/rolling.py b/third_party/bigframes_vendored/pandas/core/window/rolling.py index 7a9239b70cf..a869c86e72a 100644 --- a/third_party/bigframes_vendored/pandas/core/window/rolling.py +++ b/third_party/bigframes_vendored/pandas/core/window/rolling.py @@ -4,34 +4,36 @@ similar to how we have a Groupby object. """ +from bigframes import constants + class Window: """Provide window calculations.""" def count(self): """Calculate the window count of non-NULL observations.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def sum(self): """Calculate the weighted window sum.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def mean(self): """Calculate the weighted window mean.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def var(self): """Calculate the weighted window variance.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def std(self): """Calculate the weighted window standard deviation.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def max(self): """Calculate the weighted window maximum.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def min(self): """Calculate the weighted window minimum.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 730872034dc..95531ff5e82 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -5,6 +5,8 @@ from typing import Iterable, Optional +from bigframes import constants + class GBQIOMixin: def read_gbq( @@ -83,4 +85,4 @@ def read_gbq( Returns: bigframes.dataframe.DataFrame: A DataFrame representing results of the query or table. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/io/parquet.py b/third_party/bigframes_vendored/pandas/io/parquet.py index 6f0a2b3cb46..9aed9af5a82 100644 --- a/third_party/bigframes_vendored/pandas/io/parquet.py +++ b/third_party/bigframes_vendored/pandas/io/parquet.py @@ -2,6 +2,8 @@ """ parquet compat """ from __future__ import annotations +from bigframes import constants + class ParquetIOMixin: def read_parquet( @@ -22,4 +24,4 @@ def read_parquet( Returns: bigframes.dataframe.DataFrame: A BigQuery DataFrames. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/io/parsers/readers.py b/third_party/bigframes_vendored/pandas/io/parsers/readers.py index e01eb734fb8..d19a92ecdf8 100644 --- a/third_party/bigframes_vendored/pandas/io/parsers/readers.py +++ b/third_party/bigframes_vendored/pandas/io/parsers/readers.py @@ -6,10 +6,22 @@ """ from __future__ import annotations -from typing import Any, Dict, Literal, MutableSequence, Optional, Sequence, Tuple, Union +from typing import ( + Any, + Dict, + IO, + Literal, + MutableSequence, + Optional, + Sequence, + Tuple, + Union, +) import numpy as np +from bigframes import constants + class ReaderIOMixin: def read_csv( @@ -45,7 +57,8 @@ def read_csv( Args: filepath_or_buffer (str): - a string path including Cloud Storage and local file. + A local or Google Cloud Storage (`gs://`) path with `engine="bigquery"` + otherwise passed to pandas.read_csv. sep (Optional[str], default ","): the separator for fields in a CSV file. For the BigQuery engine, the separator can be any ISO-8859-1 single-byte character. To use a character in the range @@ -104,10 +117,71 @@ def read_csv( https://docs.python.org/3/library/codecs.html#standard-encodings The BigQuery engine only supports `UTF-8` and `ISO-8859-1`. **kwargs: - keyword arguments. + keyword arguments for `pandas.read_csv` when not using the BigQuery engine. Returns: bigframes.dataframe.DataFrame: A BigQuery DataFrames. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def read_json( + self, + path_or_buf: str | IO["bytes"], + *, + orient: Literal[ + "split", "records", "index", "columns", "values", "table" + ] = "columns", + dtype: Optional[Dict] = None, + encoding: Optional[str] = None, + lines: bool = False, + engine: Literal["ujson", "pyarrow", "bigquery"] = "ujson", + **kwargs, + ): + """ + Convert a JSON string to DataFrame object. + + .. note:: + using `engine="bigquery"` will not guarantee the same ordering as the + file. Instead, set a serialized index column as the index and sort by + that in the resulting DataFrame. + + Args: + path_or_buf (a valid JSON str, path object or file-like object): + A local or Google Cloud Storage (`gs://`) path with `engine="bigquery"` + otherwise passed to pandas.read_json. + orient (str, optional): + If `engine="bigquery"` orient only supports "records". + Indication of expected JSON string format. + Compatible JSON strings can be produced by ``to_json()`` with a + corresponding orient value. + The set of possible orients is: + + - ``'split'`` : dict like + ``{{index -> [index], columns -> [columns], data -> [values]}}`` + - ``'records'`` : list like + ``[{{column -> value}}, ... , {{column -> value}}]`` + - ``'index'`` : dict like ``{{index -> {{column -> value}}}}`` + - ``'columns'`` : dict like ``{{column -> {{index -> value}}}}`` + - ``'values'`` : just the values array + + dtype (bool or dict, default None): + If True, infer dtypes; if a dict of column to dtype, then use those; + if False, then don't infer dtypes at all, applies only to the data. + + For all ``orient`` values except ``'table'``, default is True. + encoding (str, default is 'utf-8'): + The encoding to use to decode py3 bytes. + lines (bool, default False): + Read the file as a json object per line. If using `engine="bigquery"` lines only supports True. + engine ({{"ujson", "pyarrow", "bigquery"}}, default "ujson"): + Type of engine to use. If `engine="bigquery"` is specified, then BigQuery's load API will be used. + Otherwise, the engine will be passed to `pandas.read_json`. + **kwargs: + keyword arguments for `pandas.read_json` when not using the BigQuery engine. + + Returns: + bigframes.dataframe.DataFrame: + The DataFrame representing JSON contents. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/io/pickle.py b/third_party/bigframes_vendored/pandas/io/pickle.py index a160ef0c4ea..71b31956a07 100644 --- a/third_party/bigframes_vendored/pandas/io/pickle.py +++ b/third_party/bigframes_vendored/pandas/io/pickle.py @@ -9,6 +9,8 @@ StorageOptions, ) +from bigframes import constants + class PickleIOMixin: def read_pickle( @@ -52,4 +54,4 @@ def read_pickle( bigframes.dataframe.DataFrame or bigframes.series.Series: same type as object stored in file. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py index fc48cde85ba..847ad06f750 100644 --- a/third_party/bigframes_vendored/sklearn/base.py +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -9,6 +9,8 @@ import inspect from typing import Any, Dict, List +from bigframes import constants + class BaseEstimator: """Base class for all estimators. @@ -94,7 +96,7 @@ def score(self, X, y): Returns: bigframes.dataframe.DataFrame: A DataFrame of the evaluation result. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class RegressorMixin: @@ -120,7 +122,7 @@ def score(self, X, y): Returns: bigframes.dataframe.DataFrame: A DataFrame of the evaluation result. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class TransformerMixin: diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index ea4df0dc023..ff1c04edbe3 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -14,6 +14,7 @@ from abc import ABC from typing import List, Optional +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -32,7 +33,7 @@ def predict(self, X): bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing the class labels for each sample. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class KMeans(_BaseKMeans): @@ -65,7 +66,7 @@ def fit( Returns: KMeans: Fitted Estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def predict( self, @@ -80,7 +81,7 @@ def predict( Returns: bigframes.dataframe.DataFrame: DataFrame of the cluster each sample belongs to. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def score( self, @@ -98,7 +99,7 @@ def score( Returns: bigframes.dataframe.DataFrame: DataFrame of the metrics. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property def cluster_centers_(self): @@ -119,4 +120,4 @@ def cluster_centers_(self): The output contains one row per feature per centroid. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py index bc8bc3980ac..dead173b2df 100644 --- a/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py +++ b/third_party/bigframes_vendored/sklearn/compose/_column_transformer.py @@ -6,6 +6,7 @@ from abc import ABCMeta +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -43,7 +44,7 @@ def fit( Returns: ColumnTransformer: Fitted estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def transform( self, @@ -58,4 +59,4 @@ def transform( Returns: bigframes.dataframe.DataFrame: Transformed result. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py index 619c13f35d7..85feab00245 100644 --- a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py +++ b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py @@ -12,6 +12,7 @@ from abc import ABCMeta +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -55,7 +56,7 @@ def fit(self, X, y=None): Returns: PCA: Fitted estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def score(self, X=None, y=None): """Return the metrics of the model. @@ -69,7 +70,7 @@ def score(self, X=None, y=None): Returns: bigframes.dataframe.DataFrame: DataFrame that represents model metrics. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def predict(self, X): """Predict the closest cluster for each sample in X. @@ -80,4 +81,51 @@ def predict(self, X): Returns: bigframes.dataframe.DataFrame: predicted DataFrames.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def components_(self): + """Principal axes in feature space, representing the directions of maximum variance in the data. + + Returns: + bigframes.dataframe.DataFrame: DataFrame of principal components, containing following columns: + principal_component_id: An integer that identifies the principal component. + + feature: The column name that contains the feature. + + numerical_value: If feature is numeric, the value of feature for the principal component that principal_component_id identifies. If feature isn't numeric, the value is NULL. + + categorical_value: An list of mappings containing information about categorical features. Each mapping contains the following fields: + categorical_value.category: The name of each category. + + categorical_value.value: The value of categorical_value.category for the centroid that centroid_id identifies. + + The output contains one row per feature per component. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def explained_variance_(self): + """The amount of variance explained by each of the selected components. + + Returns: + bigframes.dataframe.DataFrame: DataFrame containing following columns: + principal_component_id: An integer that identifies the principal component. + + explained_variance: The factor by which the eigenvector is scaled. Eigenvalue and explained variance are the same concepts in PCA. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def explained_variance_ratio_(self): + """Percentage of variance explained by each of the selected components. + + Returns: + bigframes.dataframe.DataFrame: DataFrame containing following columns: + principal_component_id: An integer that identifies the principal component. + + explained_variance_ratio: the total variance is the sum of variances, also known as eigenvalues, of all + of the individual principal components. The explained variance ratio by a principal component is + the ratio between the variance, also known as eigenvalue, of that principal component and the total variance. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py index 73f4684dc30..79224a772d6 100644 --- a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py +++ b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py @@ -33,6 +33,8 @@ class calls the ``fit`` method of each sub-estimator on random samples from abc import ABCMeta +from bigframes import constants + from ..base import BaseEstimator, ClassifierMixin, MetaEstimatorMixin, RegressorMixin @@ -60,7 +62,7 @@ def fit(self, X, y): Returns: Fitted Estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class ForestRegressor(RegressorMixin, BaseForest, metaclass=ABCMeta): @@ -82,7 +84,7 @@ def predict(self, X): Returns: The predicted values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class RandomForestRegressor(ForestRegressor): @@ -148,7 +150,7 @@ def predict(self, X): Returns: The predicted values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class RandomForestClassifier(ForestClassifier): diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py index 65e895298df..8141da4e3bc 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_base.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -18,6 +18,7 @@ from abc import ABCMeta from typing import List, Optional +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import ( BaseEstimator, ClassifierMixin, @@ -36,7 +37,7 @@ def predict(self, X): Returns: bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,). Returns predicted values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class LinearClassifierMixin(ClassifierMixin): @@ -52,7 +53,7 @@ def predict(self, X): bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing the class labels for each sample. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class LinearRegression(RegressorMixin, LinearModel): @@ -92,4 +93,4 @@ def fit( Returns: LinearRegression: Fitted Estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py index 8525e570689..a06035eef66 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py @@ -14,6 +14,7 @@ from typing import List, Optional +from bigframes import constants from third_party.bigframes_vendored.sklearn.linear_model._base import ( BaseEstimator, LinearClassifierMixin, @@ -57,4 +58,4 @@ def fit( Returns: LogisticRegression: Fitted Estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/metrics/_classification.py b/third_party/bigframes_vendored/sklearn/metrics/_classification.py index 6d9692ac8d6..a9d8038e59b 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_classification.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_classification.py @@ -20,6 +20,8 @@ # Michal Karbownik # License: BSD 3 clause +from bigframes import constants + def accuracy_score(y_true, y_pred, normalize=True) -> float: """Accuracy classification score. @@ -39,7 +41,7 @@ def accuracy_score(y_true, y_pred, normalize=True) -> float: classified samples (float), else returns the number of correctly classified samples (int). """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def confusion_matrix( @@ -68,7 +70,7 @@ def confusion_matrix( samples with true label being i-th class and predicted label being j-th class. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def recall_score( @@ -99,7 +101,7 @@ def recall_score( of the positive class in binary classification or weighted average of the recall of each class for the multiclass task. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def precision_score( @@ -132,7 +134,7 @@ def precision_score( Precision of the positive class in binary classification or weighted average of the precision of each class for the multiclass task. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def f1_score( @@ -167,4 +169,4 @@ def f1_score( average of the F1 scores of each class for the multiclass task. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/metrics/_ranking.py b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py index 693996070f0..ac919edbe33 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_ranking.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_ranking.py @@ -16,6 +16,8 @@ # Michal Karbownik # License: BSD 3 clause +from bigframes import constants + def auc(x, y) -> float: """Compute Area Under the Curve (AUC) using the trapezoidal rule. @@ -35,7 +37,7 @@ def auc(x, y) -> float: Returns: float: Area Under the Curve. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def roc_auc_score(y_true, y_score) -> float: @@ -60,7 +62,7 @@ def roc_auc_score(y_true, y_score) -> float: Returns: float: Area Under the Curve score. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def roc_curve( @@ -95,4 +97,4 @@ def roc_curve( fpr and tpr. `thresholds[0]` represents no instances being predicted and is arbitrarily set to `max(y_score) + 1`. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/metrics/_regression.py b/third_party/bigframes_vendored/sklearn/metrics/_regression.py index b90c4158878..9740c540e90 100644 --- a/third_party/bigframes_vendored/sklearn/metrics/_regression.py +++ b/third_party/bigframes_vendored/sklearn/metrics/_regression.py @@ -24,6 +24,8 @@ # Ohad Michel # License: BSD 3 clause +from bigframes import constants + def r2_score(y_true, y_pred, force_finite=True) -> float: """:math:`R^2` (coefficient of determination) regression score function. @@ -49,4 +51,4 @@ def r2_score(y_true, y_pred, force_finite=True) -> float: Returns: float: The :math:`R^2` score. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/pipeline.py b/third_party/bigframes_vendored/sklearn/pipeline.py index f8bbae86df6..4b8eb25a973 100644 --- a/third_party/bigframes_vendored/sklearn/pipeline.py +++ b/third_party/bigframes_vendored/sklearn/pipeline.py @@ -11,6 +11,7 @@ from abc import ABCMeta +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -47,7 +48,7 @@ def fit( Returns: Pipeline: Pipeline with fitted steps. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def score(self, X, y): @@ -67,7 +68,7 @@ def score(self, X, y): DataFrame: A DataFrame representing the result of calling `score` on the final estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def predict(self, X): @@ -81,4 +82,4 @@ def predict(self, X): bigframes.dataframe.DataFrame: A Dataframe representing predicted result. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index c57d1f22304..d0130434678 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -7,6 +7,7 @@ # Eric Chang # License: BSD 3 clause +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -59,7 +60,7 @@ def fit(self, X): Returns: StandardScaler: Fitted scaler. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def transform(self, X): """Perform standardization by centering and scaling. @@ -71,4 +72,4 @@ def transform(self, X): Returns: bigframes.dataframe.DataFrame: Transformed result. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py index a6c32d91c1d..b1cf17e539d 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -2,6 +2,7 @@ # Joris Van den Bossche # License: BSD 3 clause +from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -61,7 +62,7 @@ def fit(self, X): Returns: OneHotEncoder: Fitted encoder. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def transform(self, X): """Transform X using one-hot encoding. @@ -73,4 +74,4 @@ def transform(self, X): Returns: bigframes.dataframe.DataFrame: The result is categorized as index: number, value: number. Where index is the position of the dict that seeing the category, and value is 0 or 1.""" - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/xgboost/sklearn.py b/third_party/bigframes_vendored/xgboost/sklearn.py index fcb5d2ec593..620c87fa3db 100644 --- a/third_party/bigframes_vendored/xgboost/sklearn.py +++ b/third_party/bigframes_vendored/xgboost/sklearn.py @@ -2,6 +2,8 @@ from typing import Any +from bigframes import constants + from ..sklearn.base import BaseEstimator as XGBModelBase from ..sklearn.base import ClassifierMixin as XGBClassifierBase from ..sklearn.base import RegressorMixin as XGBRegressorBase @@ -18,7 +20,7 @@ def predict(self, X): Returns: DataFrame of shape (n_samples,): Returns predicted values. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def fit(self, X, y): """Fit gradient boosting model. @@ -42,7 +44,7 @@ def fit(self, X, y): Returns: XGBModel: Fitted Estimator. """ - raise NotImplementedError("abstract method") + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class XGBClassifierMixIn: From 7011a127722514fc562ff5a7ee259a5904f8c9ab Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 29 Aug 2023 11:40:36 -0500 Subject: [PATCH 08/80] chore: add templated GitHub config files (#12) Pulled manually from https://github.com/googleapis/synthtool/tree/master/synthtool/gcp/templates/python_library/.github Change-Id: I7f36912dffb427af2cd388abb8109670ac162701 --- .github/CODEOWNERS | 11 ++++ .github/CONTRIBUTING.md | 28 ++++++++++ .github/ISSUE_TEMPLATE/bug_report.md | 43 ++++++++++++++ .github/ISSUE_TEMPLATE/feature_request.md | 18 ++++++ .github/ISSUE_TEMPLATE/support_request.md | 7 +++ .github/PULL_REQUEST_TEMPLATE.md | 7 +++ .github/auto-approve.yml | 3 + .github/auto-label.yaml | 15 +++++ .github/header-checker-lint.yml | 15 +++++ .github/release-please.yml | 2 + .github/release-trigger.yml | 1 + .github/snippet-bot.yml | 0 .github/sync-repo-settings.yaml | 31 +++++++++++ .github/workflows/docs.yml | 38 +++++++++++++ .github/workflows/lint.yml | 25 +++++++++ .github/workflows/unittest.yml | 57 +++++++++++++++++++ SECURITY.md | 7 +++ owlbot.py | 68 +++++++++++++++++++++++ 18 files changed, 376 insertions(+) create mode 100644 .github/CODEOWNERS create mode 100644 .github/CONTRIBUTING.md create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md create mode 100644 .github/ISSUE_TEMPLATE/feature_request.md create mode 100644 .github/ISSUE_TEMPLATE/support_request.md create mode 100644 .github/PULL_REQUEST_TEMPLATE.md create mode 100644 .github/auto-approve.yml create mode 100644 .github/auto-label.yaml create mode 100644 .github/header-checker-lint.yml create mode 100644 .github/release-please.yml create mode 100644 .github/release-trigger.yml create mode 100644 .github/snippet-bot.yml create mode 100644 .github/sync-repo-settings.yaml create mode 100644 .github/workflows/docs.yml create mode 100644 .github/workflows/lint.yml create mode 100644 .github/workflows/unittest.yml create mode 100644 SECURITY.md create mode 100644 owlbot.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000000..cd904459a91 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,11 @@ +# 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/yoshi-python @googleapis/api-bigquery-dataframe are the default owners for changes in this repo +* @googleapis/yoshi-python @googleapis/api-bigquery-dataframe + +# @googleapis/python-samples-reviewers @googleapis/api-bigquery-dataframe are the default owners for samples changes +/samples/ @googleapis/python-samples-reviewers @googleapis/api-bigquery-dataframe 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..7b0900728e4 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,43 @@ +--- +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` + +#### 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-approve.yml b/.github/auto-approve.yml new file mode 100644 index 00000000000..311ebbb853a --- /dev/null +++ b/.github/auto-approve.yml @@ -0,0 +1,3 @@ +# https://github.com/googleapis/repo-automation-bots/tree/main/packages/auto-approve +processes: + - "OwlBotTemplateChanges" diff --git a/.github/auto-label.yaml b/.github/auto-label.yaml new file mode 100644 index 00000000000..b2016d119b4 --- /dev/null +++ b/.github/auto-label.yaml @@ -0,0 +1,15 @@ +# 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. +requestsize: + enabled: true diff --git a/.github/header-checker-lint.yml b/.github/header-checker-lint.yml new file mode 100644 index 00000000000..3058bec338c --- /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" + ] +} diff --git a/.github/release-please.yml b/.github/release-please.yml new file mode 100644 index 00000000000..466597e5b19 --- /dev/null +++ b/.github/release-please.yml @@ -0,0 +1,2 @@ +releaseType: python +handleGHRelease: true diff --git a/.github/release-trigger.yml b/.github/release-trigger.yml new file mode 100644 index 00000000000..d4ca94189e1 --- /dev/null +++ b/.github/release-trigger.yml @@ -0,0 +1 @@ +enabled: true diff --git a/.github/snippet-bot.yml b/.github/snippet-bot.yml new file mode 100644 index 00000000000..e69de29bb2d diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml new file mode 100644 index 00000000000..5004e872b4c --- /dev/null +++ b/.github/sync-repo-settings.yaml @@ -0,0 +1,31 @@ +# https://github.com/googleapis/repo-automation-bots/tree/main/packages/sync-repo-settings +# Rules for main branch protection +branchProtectionRules: +# Identifies the protection rule pattern. Name of the branch to be protected. +# Defaults to `main` +- pattern: main + requiresCodeOwnerReviews: true + requiresStrictStatusChecks: true + requiredStatusCheckContexts: + - 'cla/google' + - 'OwlBot Post Processor' + - 'docs' + - 'lint' + - 'unit (3.9)' + - 'unit (3.10)' + - 'unit (3.11)' + - 'cover' + # TODO(tswast): add Kokoro once we've enabled it +permissionRules: + - team: actools-python + permission: admin + - team: actools + permission: admin + - team: api-bigquery-dataframe + permission: push + - team: yoshi-python + permission: push + - team: python-samples-owners + permission: push + - team: python-samples-reviewers + permission: push diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000000..e97d89e484c --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,38 @@ +on: + pull_request: + branches: + - main +name: docs +jobs: + docs: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: "3.9" + - 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@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: "3.9" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run docfx + run: | + nox -s docfx diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000000..d2aee5b7d8e --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,25 @@ +on: + pull_request: + branches: + - main +name: lint +jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - 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/unittest.yml b/.github/workflows/unittest.yml new file mode 100644 index 00000000000..bb268fe6c66 --- /dev/null +++ b/.github/workflows/unittest.yml @@ -0,0 +1,57 @@ +on: + pull_request: + branches: + - main +name: unittest +jobs: + unit: + runs-on: ubuntu-latest + strategy: + matrix: + python: ['3.9', '3.10', '3.11'] + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + 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@v3 + with: + name: coverage-artifacts + path: .coverage-${{ matrix.python }} + + cover: + runs-on: ubuntu-latest + needs: + - unit + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - 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@v3 + with: + name: coverage-artifacts + path: .coverage-results/ + - name: Report coverage results + run: | + coverage combine .coverage-results/.coverage* + coverage report --show-missing --fail-under=40 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/owlbot.py b/owlbot.py new file mode 100644 index 00000000000..4ba7d14eb5e --- /dev/null +++ b/owlbot.py @@ -0,0 +1,68 @@ +# Copyright 2021 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. + +"""This script is used to synthesize generated parts of this library.""" + +import pathlib + +from synthtool import gcp +import synthtool as s +from synthtool.languages import python + +REPO_ROOT = pathlib.Path(__file__).parent.absolute() + +common = gcp.CommonTemplates() + +# ---------------------------------------------------------------------------- +# Add templated files +# ---------------------------------------------------------------------------- + +templated_files = common.py_library( + unit_test_python_versions=["3.9", "3.10", "3.11"], + system_test_python_versions=["3.9", "3.11"], + cov_level=40, + intersphinx_dependencies={ + "pandas": "https://pandas.pydata.org/pandas-docs/stable/", + "pydata-google-auth": "https://pydata-google-auth.readthedocs.io/en/latest/", + }, +) +s.move( + templated_files, + excludes=[ + # Multi-processing note isn't relevant, as pandas_gbq is responsible for + # creating clients, not the end user. + "docs/multiprocessing.rst", + "noxfile.py", + "README.rst", + ], +) + +# ---------------------------------------------------------------------------- +# Fixup files +# ---------------------------------------------------------------------------- + + +# ---------------------------------------------------------------------------- +# Samples templates +# ---------------------------------------------------------------------------- + +python.py_samples(skip_readmes=True) + +# ---------------------------------------------------------------------------- +# Final cleanup +# ---------------------------------------------------------------------------- + +s.shell.run(["nox", "-s", "blacken"], hide_output=False) +for noxfile in REPO_ROOT.glob("samples/**/noxfile.py"): + s.shell.run(["nox", "-s", "blacken"], cwd=noxfile.parent, hide_output=False) From bf6ecb81afeb199b3dad07d1fd2057668352f939 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 29 Aug 2023 10:16:44 -0700 Subject: [PATCH 09/80] chore(deps): bump cryptography from 41.0.1 to 41.0.3 in /.kokoro (#3) Bumps [cryptography](https://github.com/pyca/cryptography) from 41.0.1 to 41.0.3. - [Changelog](https://github.com/pyca/cryptography/blob/main/CHANGELOG.rst) - [Commits](https://github.com/pyca/cryptography/compare/41.0.1...41.0.3) --- updated-dependencies: - dependency-name: cryptography dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .kokoro/requirements.txt | 45 ++++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 37a007667f1..fd0e0b9d489 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -123,26 +123,30 @@ commonmark==0.9.1 \ --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 # via recommonmark -cryptography==41.0.1 \ - --hash=sha256:059e348f9a3c1950937e1b5d7ba1f8e968508ab181e75fc32b879452f08356db \ - --hash=sha256:1a5472d40c8f8e91ff7a3d8ac6dfa363d8e3138b961529c996f3e2df0c7a411a \ - --hash=sha256:1a8e6c2de6fbbcc5e14fd27fb24414507cb3333198ea9ab1258d916f00bc3039 \ - --hash=sha256:1fee5aacc7367487b4e22484d3c7e547992ed726d14864ee33c0176ae43b0d7c \ - --hash=sha256:5d092fdfedaec4cbbffbf98cddc915ba145313a6fdaab83c6e67f4e6c218e6f3 \ - --hash=sha256:5f0ff6e18d13a3de56f609dd1fd11470918f770c6bd5d00d632076c727d35485 \ - --hash=sha256:7bfc55a5eae8b86a287747053140ba221afc65eb06207bedf6e019b8934b477c \ - --hash=sha256:7fa01527046ca5facdf973eef2535a27fec4cb651e4daec4d043ef63f6ecd4ca \ - --hash=sha256:8dde71c4169ec5ccc1087bb7521d54251c016f126f922ab2dfe6649170a3b8c5 \ - --hash=sha256:8f4ab7021127a9b4323537300a2acfb450124b2def3756f64dc3a3d2160ee4b5 \ - --hash=sha256:948224d76c4b6457349d47c0c98657557f429b4e93057cf5a2f71d603e2fc3a3 \ - --hash=sha256:9a6c7a3c87d595608a39980ebaa04d5a37f94024c9f24eb7d10262b92f739ddb \ - --hash=sha256:b46e37db3cc267b4dea1f56da7346c9727e1209aa98487179ee8ebed09d21e43 \ - --hash=sha256:b4ceb5324b998ce2003bc17d519080b4ec8d5b7b70794cbd2836101406a9be31 \ - --hash=sha256:cb33ccf15e89f7ed89b235cff9d49e2e62c6c981a6061c9c8bb47ed7951190bc \ - --hash=sha256:d198820aba55660b4d74f7b5fd1f17db3aa5eb3e6893b0a41b75e84e4f9e0e4b \ - --hash=sha256:d34579085401d3f49762d2f7d6634d6b6c2ae1242202e860f4d26b046e3a1006 \ - --hash=sha256:eb8163f5e549a22888c18b0d53d6bb62a20510060a22fd5a995ec8a05268df8a \ - --hash=sha256:f73bff05db2a3e5974a6fd248af2566134d8981fd7ab012e5dd4ddb1d9a70699 +cryptography==41.0.3 \ + --hash=sha256:0d09fb5356f975974dbcb595ad2d178305e5050656affb7890a1583f5e02a306 \ + --hash=sha256:23c2d778cf829f7d0ae180600b17e9fceea3c2ef8b31a99e3c694cbbf3a24b84 \ + --hash=sha256:3fb248989b6363906827284cd20cca63bb1a757e0a2864d4c1682a985e3dca47 \ + --hash=sha256:41d7aa7cdfded09b3d73a47f429c298e80796c8e825ddfadc84c8a7f12df212d \ + --hash=sha256:42cb413e01a5d36da9929baa9d70ca90d90b969269e5a12d39c1e0d475010116 \ + --hash=sha256:4c2f0d35703d61002a2bbdcf15548ebb701cfdd83cdc12471d2bae80878a4207 \ + --hash=sha256:4fd871184321100fb400d759ad0cddddf284c4b696568204d281c902fc7b0d81 \ + --hash=sha256:5259cb659aa43005eb55a0e4ff2c825ca111a0da1814202c64d28a985d33b087 \ + --hash=sha256:57a51b89f954f216a81c9d057bf1a24e2f36e764a1ca9a501a6964eb4a6800dd \ + --hash=sha256:652627a055cb52a84f8c448185922241dd5217443ca194d5739b44612c5e6507 \ + --hash=sha256:67e120e9a577c64fe1f611e53b30b3e69744e5910ff3b6e97e935aeb96005858 \ + --hash=sha256:6af1c6387c531cd364b72c28daa29232162010d952ceb7e5ca8e2827526aceae \ + --hash=sha256:6d192741113ef5e30d89dcb5b956ef4e1578f304708701b8b73d38e3e1461f34 \ + --hash=sha256:7efe8041897fe7a50863e51b77789b657a133c75c3b094e51b5e4b5cec7bf906 \ + --hash=sha256:84537453d57f55a50a5b6835622ee405816999a7113267739a1b4581f83535bd \ + --hash=sha256:8f09daa483aedea50d249ef98ed500569841d6498aa9c9f4b0531b9964658922 \ + --hash=sha256:95dd7f261bb76948b52a5330ba5202b91a26fbac13ad0e9fc8a3ac04752058c7 \ + --hash=sha256:a74fbcdb2a0d46fe00504f571a2a540532f4c188e6ccf26f1f178480117b33c4 \ + --hash=sha256:a983e441a00a9d57a4d7c91b3116a37ae602907a7618b882c8013b5762e80574 \ + --hash=sha256:ab8de0d091acbf778f74286f4989cf3d1528336af1b59f3e5d2ebca8b5fe49e1 \ + --hash=sha256:aeb57c421b34af8f9fe830e1955bf493a86a7996cc1338fe41b30047d16e962c \ + --hash=sha256:ce785cf81a7bdade534297ef9e490ddff800d956625020ab2ec2780a556c313e \ + --hash=sha256:d0d651aa754ef58d75cec6edfbd21259d93810b73f6ec246436a21b7841908de # via # gcp-releasetool # secretstorage @@ -416,6 +420,7 @@ protobuf==3.20.3 \ # gcp-docuploader # gcp-releasetool # google-api-core + # googleapis-common-protos pyasn1==0.5.0 \ --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \ --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde From cccac8c16c2daeb4f2defe4562cf099b1c738e07 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Wed, 30 Aug 2023 11:42:04 -0700 Subject: [PATCH 10/80] chore: set up the kokoro release pipeline (#14) --- .kokoro/release.sh | 29 ++++++++++++++++++++++ .kokoro/release/common.cfg | 49 +++++++++++++++++++++++++++++++++++++ .kokoro/release/release.cfg | 1 + 3 files changed, 79 insertions(+) create mode 100644 .kokoro/release.sh create mode 100644 .kokoro/release/common.cfg create mode 100644 .kokoro/release/release.cfg diff --git a/.kokoro/release.sh b/.kokoro/release.sh new file mode 100644 index 00000000000..697b04009db --- /dev/null +++ b/.kokoro/release.sh @@ -0,0 +1,29 @@ +#!/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 + +# Start the releasetool reporter +python3 -m pip install --require-hashes -r github/python-bigquery-dataframes/.kokoro/requirements.txt +python3 -m releasetool publish-reporter-script > /tmp/publisher-script; source /tmp/publisher-script + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=1 + +# Move into the package, build the distribution and upload. +TWINE_PASSWORD=$(cat "${KOKORO_KEYSTORE_DIR}/73713_google-cloud-pypi-token-keystore-1") +cd github/python-bigquery-dataframes +python3 setup.py sdist bdist_wheel +twine upload --username __token__ --password "${TWINE_PASSWORD}" dist/* \ No newline at end of file diff --git a/.kokoro/release/common.cfg b/.kokoro/release/common.cfg new file mode 100644 index 00000000000..7ffa79c7a1d --- /dev/null +++ b/.kokoro/release/common.cfg @@ -0,0 +1,49 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# 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" + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-multi" +} +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/release.sh" +} + +# Fetch PyPI password +before_action { + fetch_keystore { + keystore_resource { + keystore_config_id: 73713 + keyname: "google-cloud-pypi-token-keystore-1" + } + } +} + +# Tokens needed to report release status back to GitHub +env_vars: { + key: "SECRET_MANAGER_KEYS" + value: "releasetool-publish-reporter-app,releasetool-publish-reporter-googleapis-installation,releasetool-publish-reporter-pem" +} + +# Store the packages we uploaded to PyPI. That way, we have a record of exactly +# what we published, which we can use to generate SBOMs and attestations. +action { + define_artifacts { + regex: "github/python-bigquery-dataframes/**/*.tar.gz" + strip_prefix: "github/python-bigquery-dataframes" + } +} \ No newline at end of file diff --git a/.kokoro/release/release.cfg b/.kokoro/release/release.cfg new file mode 100644 index 00000000000..8f43917d92f --- /dev/null +++ b/.kokoro/release/release.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file From 8fab75576757230bca5c7df10994837ac406300f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Sat, 2 Sep 2023 09:14:42 -0500 Subject: [PATCH 11/80] chore: sync changes from internal repo (#15) feat: support `DataFrame.isin` with list and dict inputs test: move flaky `reset_session` test to unit tests chore: don't run redundant tests in nightly build feat: support `DataFrame`-`DataFrame` binary operations feat: support `Series.map` feat: support `Index.is_monotonic` docs: update remote function notebook with read_gbq_function usage feat: use default session and connection in `ml.llm` and `ml.imported` chore: disable broken stack tests feat: support `pow()` and power operator in `DataFrame` and `Series` feat: support for `np.add`, `np.subtract`, `np.multiply`, `np.divide`, `np.power` perf: use `row_number()` filter for `head` / `tail` feat: support `bigframes.pandas.merge()` fix: make `X_train` argument names consistent across methods chore: refactor ml core feat: add `Series.dropna` and `DataFrame.fillna` chore: fix gcs notebooks upload in 'nightly' build chore: fix Kokoro build files to support GitHub directories chore: fix unit test to not require authentication --- .kokoro/build.sh | 8 +- .kokoro/continuous/nightly.cfg | 5 - .kokoro/release-nightly.sh | 29 +- OWNERS | 1 + bigframes/_config/bigquery_options.py | 22 +- bigframes/clients.py | 163 ++++ bigframes/core/__init__.py | 14 - bigframes/core/blocks.py | 71 ++ bigframes/core/indexes/index.py | 30 + bigframes/core/joins/__init__.py | 2 + bigframes/core/joins/merge.py | 67 ++ bigframes/dataframe.py | 180 ++++- bigframes/ml/cluster.py | 2 +- bigframes/ml/core.py | 58 +- bigframes/ml/decomposition.py | 2 +- bigframes/ml/imported.py | 16 +- bigframes/ml/llm.py | 88 +- bigframes/ml/sql.py | 16 +- bigframes/operations/__init__.py | 123 ++- bigframes/pandas/__init__.py | 34 +- bigframes/remote_function.py | 115 +-- bigframes/series.py | 118 ++- bigframes/session.py | 38 +- notebooks/experimental/longer_ml_demo.ipynb | 34 +- .../getting_started/ml_fundamentals.ipynb | 58 +- .../sklearn_linear_regression.ipynb | 8 +- .../remote_functions/remote_function.ipynb | 749 +++++++++++++----- tests/system/conftest.py | 18 +- tests/system/large/ml/test_core.py | 16 +- tests/system/large/ml/test_ensemble.py | 64 +- tests/system/large/ml/test_forecasting.py | 6 +- tests/system/large/ml/test_linear_model.py | 32 +- tests/system/small/ml/conftest.py | 40 +- tests/system/small/ml/test_ensemble.py | 72 +- tests/system/small/ml/test_imported.py | 10 + tests/system/small/ml/test_linear_model.py | 36 +- tests/system/small/ml/test_llm.py | 28 + tests/system/small/test_dataframe.py | 133 +++- tests/system/small/test_index.py | 14 + tests/system/small/test_multiindex.py | 68 +- tests/system/small/test_numpy.py | 66 ++ tests/system/small/test_pandas.py | 104 +++ tests/system/small/test_pandas_options.py | 43 - tests/system/small/test_remote_function.py | 35 +- tests/system/small/test_series.py | 89 +++ tests/unit/_config/test_bigquery_options.py | 4 +- tests/unit/conftest.py | 223 ------ .../core/{test_utils.py => test_bf_utils.py} | 0 tests/unit/resources.py | 73 ++ tests/unit/test_core.py | 62 +- tests/unit/test_pandas.py | 39 + tests/unit/test_session.py | 15 +- .../bigframes_vendored/pandas/core/frame.py | 98 +++ .../pandas/core/reshape/merge.py | 78 ++ .../bigframes_vendored/pandas/core/series.py | 102 ++- 55 files changed, 2617 insertions(+), 1002 deletions(-) create mode 100644 bigframes/clients.py create mode 100644 bigframes/core/joins/merge.py delete mode 100644 tests/unit/conftest.py rename tests/unit/core/{test_utils.py => test_bf_utils.py} (100%) create mode 100644 tests/unit/resources.py create mode 100644 third_party/bigframes_vendored/pandas/core/reshape/merge.py diff --git a/.kokoro/build.sh b/.kokoro/build.sh index f80cf6eab97..402ac0eb8c9 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -15,8 +15,14 @@ set -eo pipefail +if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then + PROJECT_SCM="github" +else + PROJECT_SCM="git" +fi + if [[ -z "${PROJECT_ROOT:-}" ]]; then - PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes" + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes" fi cd "${PROJECT_ROOT}" diff --git a/.kokoro/continuous/nightly.cfg b/.kokoro/continuous/nightly.cfg index ac34c4b0c60..63c3f51d053 100644 --- a/.kokoro/continuous/nightly.cfg +++ b/.kokoro/continuous/nightly.cfg @@ -1,8 +1,3 @@ # Format: //devtools/kokoro/config/proto/build.proto -env_vars: { - key: "NOX_SESSION" - value: "unit system cover lint lint_setup_py mypy format docs e2e notebook" -} - build_file: "bigframes/.kokoro/release-nightly.sh" diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh index e3b6b4d4498..582808a15ce 100755 --- a/.kokoro/release-nightly.sh +++ b/.kokoro/release-nightly.sh @@ -34,8 +34,14 @@ while [ $# -gt 0 ] ; do shift 1; done +if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then + PROJECT_SCM="github" +else + PROJECT_SCM="git" +fi + if [ -z "${PROJECT_ROOT:-}" ]; then - PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes" + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes" fi # Move into the package, build the distribution and upload to shared bucket. @@ -57,16 +63,6 @@ export PYTHONUNBUFFERED=1 # Install dependencies, as the following steps depend on it python3.10 -m pip install -e .[all] -# If NOX_SESSION is set, it only runs the specified session, -# otherwise run all the sessions. -if ! [ ${DRY_RUN} ]; then - if [ -n "${NOX_SESSION:-}" ]; then - python3.10 -m nox -s ${NOX_SESSION:-} - else - python3.10 -m nox - fi -fi - # Generate third party notices and include it in the licenses in setup.cfg # TODO(shobs): Don't include it in the package once vertex colab can pick it # from elsewhere @@ -138,15 +134,8 @@ if ! [ ${DRY_RUN} ]; then gsutil cp -v dist/* ${gcs_path} gsutil cp -v LICENSE ${gcs_path} gsutil cp -v ${THIRD_PARTY_NOTICES_FILE} ${gcs_path} - gsutil -m cp -v "notebooks/00 - Summary.ipynb" \ - "notebooks/01 - Getting Started.ipynb" \ - "notebooks/02 - DataFrame.ipynb" \ - "notebooks/03 - Using ML - ML fundamentals.ipynb" \ - "notebooks/04 - Using ML - SKLearn linear regression.ipynb" \ - "notebooks/05 - Using ML - Easy linear regression.ipynb" \ - "notebooks/06 - Using ML - Large Language Models.ipynb" \ - "notebooks/50 - Remote Function.ipynb" \ - ${gcs_path}notebooks/ + gsutil -m cp -r -v "notebooks/" ${gcs_path}notebooks/ + done # publish API coverage information to BigQuery diff --git a/OWNERS b/OWNERS index f2a0b4383d1..672da38afa5 100644 --- a/OWNERS +++ b/OWNERS @@ -5,6 +5,7 @@ henryjsolberg@google.com hormati@google.com huanc@google.com jiaxun@google.com +kemppeterson@google.com shobs@google.com swast@google.com tbergeron@google.com diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index a103abe1903..ea1864ed5f7 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -35,13 +35,13 @@ def __init__( credentials: Optional[google.auth.credentials.Credentials] = None, project: Optional[str] = None, location: Optional[str] = None, - remote_udf_connection: Optional[str] = None, + bq_connection: Optional[str] = None, use_regional_endpoints: bool = False, ): self._credentials = credentials self._project = project self._location = location - self._remote_udf_connection = remote_udf_connection + self._bq_connection = bq_connection self._use_regional_endpoints = use_regional_endpoints self._session_started = False @@ -82,23 +82,21 @@ def project(self, value: Optional[str]): self._project = value @property - def remote_udf_connection(self) -> Optional[str]: - """Name of the BigQuery connection to use for remote functions. + def bq_connection(self) -> Optional[str]: + """Name of the BigQuery connection to use. You should either have the connection already created in the location you have chosen, or you should have the Project IAM Admin role to enable the service to create the connection for you if you need it. """ - return self._remote_udf_connection + return self._bq_connection - @remote_udf_connection.setter - def remote_udf_connection(self, value: Optional[str]): - if self._session_started and self._remote_udf_connection != value: - raise ValueError( - SESSION_STARTED_MESSAGE.format(attribute="remote_udf_connection") - ) - self._remote_udf_connection = value + @bq_connection.setter + def bq_connection(self, value: Optional[str]): + if self._session_started and self._bq_connection != value: + raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="bq_connection")) + self._bq_connection = value @property def use_regional_endpoints(self) -> bool: diff --git a/bigframes/clients.py b/bigframes/clients.py new file mode 100644 index 00000000000..5c019e0fc89 --- /dev/null +++ b/bigframes/clients.py @@ -0,0 +1,163 @@ +# 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. + +"""BigQuery DataFrame clients to interact with other cloud resources""" + +from __future__ import annotations + +import logging +import time +from typing import Optional + +import google.api_core.exceptions +from google.cloud import bigquery_connection_v1, resourcemanager_v3 +from google.iam.v1 import iam_policy_pb2, policy_pb2 + +logging.basicConfig( + level=logging.INFO, format="[%(levelname)s][%(asctime)s][%(name)s] %(message)s" +) +logger = logging.getLogger(__name__) + + +class BqConnectionManager: + """Manager to handle operations with BQ connections.""" + + # Wait time (in seconds) for an IAM binding to take effect after creation + _IAM_WAIT_SECONDS = 120 + + def __init__( + self, + bq_connection_client: bigquery_connection_v1.ConnectionServiceClient, + cloud_resource_manager_client: resourcemanager_v3.ProjectsClient, + ): + self._bq_connection_client = bq_connection_client + self._cloud_resource_manager_client = cloud_resource_manager_client + + def create_bq_connection( + self, project_id: str, location: str, connection_id: str, iam_role: str + ): + """Create the BQ connection if not exist. In addition, try to add the IAM role to the connection to ensure required permissions. + + Args: + project_id: + ID of the project. + location: + Location of the connection. + connection_id: + ID of the connection. + iam_role: + str of the IAM role that the service account of the created connection needs to aquire. E.g. 'run.invoker', 'aiplatform.user' + """ + # TODO(shobs): The below command to enable BigQuery Connection API needs + # to be automated. Disabling for now since most target users would not + # have the privilege to enable API in a project. + # log("Making sure BigQuery Connection API is enabled") + # if os.system("gcloud services enable bigqueryconnection.googleapis.com"): + # raise ValueError("Failed to enable BigQuery Connection API") + # If the intended connection does not exist then create it + service_account_id = self._get_service_account_if_connection_exists( + project_id, location, connection_id + ) + if service_account_id: + logger.info( + f"Connector {project_id}.{location}.{connection_id} already exists" + ) + else: + connection_name, service_account_id = self._create_bq_connection( + project_id, location, connection_id + ) + logger.info( + f"Created BQ connection {connection_name} with service account id: {service_account_id}" + ) + # Ensure IAM role on the BQ connection + # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function + self._ensure_iam_binding(project_id, service_account_id, iam_role) + + # Introduce retries to accommodate transient errors like etag mismatch, + # which can be caused by concurrent operation on the same resource, and + # manifests with message like: + # google.api_core.exceptions.Aborted: 409 There were concurrent policy + # changes. Please retry the whole read-modify-write with exponential + # backoff. The request's ETag '\007\006\003,\264\304\337\272' did not match + # the current policy's ETag '\007\006\003,\3750&\363'. + @google.api_core.retry.Retry( + predicate=google.api_core.retry.if_exception_type( + google.api_core.exceptions.Aborted + ), + initial=10, + maximum=20, + multiplier=2, + timeout=60, + ) + def _ensure_iam_binding( + self, project_id: str, service_account_id: str, iam_role: str + ): + """Ensure necessary IAM role is configured on a service account.""" + project = f"projects/{project_id}" + service_account = f"serviceAccount:{service_account_id}" + role = f"roles/{iam_role}" + request = iam_policy_pb2.GetIamPolicyRequest(resource=project) + policy = self._cloud_resource_manager_client.get_iam_policy(request=request) + + # Check if the binding already exists, and if does, do nothing more + for binding in policy.bindings: + if binding.role == role: + if service_account in binding.members: + return + + # Create a new binding + new_binding = policy_pb2.Binding(role=role, members=[service_account]) + policy.bindings.append(new_binding) + request = iam_policy_pb2.SetIamPolicyRequest(resource=project, policy=policy) + self._cloud_resource_manager_client.set_iam_policy(request=request) + + # We would wait for the IAM policy change to take effect + # https://cloud.google.com/iam/docs/access-change-propagation + logger.info( + f"Waiting {self._IAM_WAIT_SECONDS} seconds for IAM to take effect.." + ) + time.sleep(self._IAM_WAIT_SECONDS) + + def _create_bq_connection(self, project_id: str, location: str, connection_id: str): + """Create the BigQuery Connection and returns corresponding service account id.""" + client = self._bq_connection_client + connection = bigquery_connection_v1.Connection( + cloud_resource=bigquery_connection_v1.CloudResourceProperties() + ) + request = bigquery_connection_v1.CreateConnectionRequest( + parent=client.common_location_path(project_id, location), + connection_id=connection_id, + connection=connection, + ) + connection = client.create_connection(request) + return connection.name, connection.cloud_resource.service_account_id + + def _get_service_account_if_connection_exists( + self, project_id: str, location: str, connection_id: str + ) -> Optional[str]: + """Check if the BigQuery Connection exists.""" + client = self._bq_connection_client + request = bigquery_connection_v1.GetConnectionRequest( + name=client.connection_path(project_id, location, connection_id) + ) + + service_account = None + try: + service_account = client.get_connection( + request=request + ).cloud_resource.service_account_id + except google.api_core.exceptions.NotFound: + pass + + return service_account diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index d6509e4c0a9..7086269af9b 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -1021,20 +1021,6 @@ def slice( if not step: step = 1 - # Special cases for head() and tail(), where we don't need to project - # offsets. LIMIT clause is much more efficient in BigQuery than a - # filter on row_number(). - if ( - (start is None or start == 0) - and step == 1 - and stop is not None - and stop > 0 - ): - return self.apply_limit(stop) - - if start is not None and start < 0 and step == 1 and stop is None: - return self.reversed().apply_limit(abs(start)).reversed() - expr_with_offsets = self.project_offsets() # start with True and reduce with start, stop, and step conditions diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index f23a4d0b5c9..482cfd0141e 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -62,6 +62,10 @@ _UNIFORM = "uniform" _SAMPLING_METHODS = (_HEAD, _UNIFORM) +# Monotonic Cache Names +_MONOTONIC_INCREASING = "monotonic_increasing" +_MONOTONIC_DECREASING = "monotonic_decreasing" + class BlockHolder(typing.Protocol): """Interface for mutable objects with state represented by a block value object.""" @@ -119,6 +123,8 @@ def __init__( } # TODO(kemppeterson) Add a cache for corr to parallel the single-column stats. + self._stats_cache[" ".join(self.index_columns)] = {} + @property def index(self) -> indexes.IndexValue: """Row identities for values in the Block.""" @@ -1408,6 +1414,71 @@ def _force_reproject(self) -> Block: index_labels=self.index.names, ) + def is_monotonic_increasing( + self, column_id: typing.Union[str, Sequence[str]] + ) -> bool: + return self._is_monotonic(column_id, increasing=True) + + def is_monotonic_decreasing( + self, column_id: typing.Union[str, Sequence[str]] + ) -> bool: + return self._is_monotonic(column_id, increasing=False) + + def _is_monotonic( + self, column_ids: typing.Union[str, Sequence[str]], increasing: bool + ) -> bool: + if isinstance(column_ids, str): + column_ids = (column_ids,) + + op_name = _MONOTONIC_INCREASING if increasing else _MONOTONIC_DECREASING + + column_name = " ".join(column_ids) + if op_name in self._stats_cache[column_name]: + return self._stats_cache[column_name][op_name] + + period = 1 + window = bigframes.core.WindowSpec( + preceding=period, + following=None, + ) + + # any NaN value means not monotonic + block, last_notna_id = self.apply_unary_op(column_ids[0], ops.notnull_op) + for column_id in column_ids[1:]: + block, notna_id = block.apply_unary_op(column_id, ops.notnull_op) + block, last_notna_id = block.apply_binary_op( + last_notna_id, notna_id, ops.and_op + ) + + # loop over all columns to check monotonicity + last_result_id = None + for column_id in column_ids[::-1]: + block, lag_result_id = block.apply_window_op( + column_id, agg_ops.ShiftOp(period), window + ) + block, strict_monotonic_id = block.apply_binary_op( + column_id, lag_result_id, ops.gt_op if increasing else ops.lt_op + ) + block, equal_id = block.apply_binary_op(column_id, lag_result_id, ops.eq_op) + if last_result_id is None: + block, last_result_id = block.apply_binary_op( + equal_id, strict_monotonic_id, ops.or_op + ) + continue + block, equal_monotonic_id = block.apply_binary_op( + equal_id, last_result_id, ops.and_op + ) + block, last_result_id = block.apply_binary_op( + equal_monotonic_id, strict_monotonic_id, ops.or_op + ) + + block, monotonic_result_id = block.apply_binary_op( + last_result_id, last_notna_id, ops.and_op # type: ignore + ) + result = block.get_stat(monotonic_result_id, agg_ops.all_op) + self._stats_cache[column_name].update({op_name: result}) + return result + def block_from_local(data, session=None, use_index=True) -> Block: # TODO(tbergeron): Handle duplicate column labels diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 4eb37e6d926..04b9a36b64a 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -67,6 +67,36 @@ def empty(self) -> bool: """Returns True if the Index is empty, otherwise returns False.""" return self.shape[0] == 0 + @property + def is_monotonic_increasing(self) -> bool: + """ + Return a boolean if the values are equal or increasing. + + Returns: + bool + """ + return typing.cast( + bool, + self._data._get_block().is_monotonic_increasing( + self._data._get_block().index_columns + ), + ) + + @property + def is_monotonic_decreasing(self) -> bool: + """ + Return a boolean if the values are equal or decreasing. + + Returns: + bool + """ + return typing.cast( + bool, + self._data._get_block().is_monotonic_decreasing( + self._data._get_block().index_columns + ), + ) + def __getitem__(self, key: int) -> typing.Any: if isinstance(key, int): result_pd_df, _ = self._data._get_block().slice(key, key + 1, 1).to_pandas() diff --git a/bigframes/core/joins/__init__.py b/bigframes/core/joins/__init__.py index 58a1c2cfd79..3f9447aef0d 100644 --- a/bigframes/core/joins/__init__.py +++ b/bigframes/core/joins/__init__.py @@ -14,10 +14,12 @@ """Helpers to join ArrayValue objects.""" +from bigframes.core.joins.merge import merge from bigframes.core.joins.row_identity import join_by_row_identity from bigframes.core.joins.single_column import join_by_column __all__ = ( "join_by_row_identity", "join_by_column", + "merge", ) diff --git a/bigframes/core/joins/merge.py b/bigframes/core/joins/merge.py new file mode 100644 index 00000000000..fac16b36078 --- /dev/null +++ b/bigframes/core/joins/merge.py @@ -0,0 +1,67 @@ +# 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. + +""" +Functions for Merging Data Structures in BigFrames. +""" + +from __future__ import annotations + +from typing import Literal, Optional + +from bigframes.dataframe import DataFrame +from bigframes.series import Series + + +def merge( + left: DataFrame, + right: DataFrame, + how: Literal[ + "inner", + "left", + "outer", + "right", + ] = "inner", + on: Optional[str] = None, + *, + left_on: Optional[str] = None, + right_on: Optional[str] = None, + sort: bool = False, + suffixes: tuple[str, str] = ("_x", "_y"), +) -> DataFrame: + left = _validate_operand(left) + right = _validate_operand(right) + + return left.merge( + right, + how=how, + on=on, + left_on=left_on, + right_on=right_on, + sort=sort, + suffixes=suffixes, + ) + + +def _validate_operand(obj: DataFrame | Series) -> DataFrame: + if isinstance(obj, DataFrame): + return obj + elif isinstance(obj, Series): + if obj.name is None: + raise ValueError("Cannot merge a Series without a name") + return obj.to_frame() + else: + raise TypeError( + f"Can only merge Series or DataFrame objects, a {type(obj)} was passed" + ) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index ef443db0797..31777f3face 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -529,7 +529,7 @@ def __setitem__(self, key: str, value: SingleItemValue): def _apply_binop( self, - other: float | int | bigframes.series.Series, + other: float | int | bigframes.series.Series | DataFrame, op, axis: str | int = "columns", ): @@ -537,6 +537,8 @@ def _apply_binop( return self._apply_scalar_binop(other, op) elif isinstance(other, bigframes.series.Series): return self._apply_series_binop(other, op, axis=axis) + elif isinstance(other, DataFrame): + return self._apply_dataframe_binop(other, op) raise NotImplementedError( f"binary operation is not implemented on the second operand of type {type(other).__name__}." f"{constants.FEEDBACK_LINK}" @@ -588,6 +590,47 @@ def _apply_series_binop( block = block.with_index_labels(self.index.names) return DataFrame(block) + def _apply_dataframe_binop( + self, + other: DataFrame, + op: ops.BinaryOp, + ) -> DataFrame: + # Join rows + joined_index, (get_column_left, get_column_right) = self._block.index.join( + other._block.index, how="outer" + ) + # join columns schema + columns, lcol_indexer, rcol_indexer = self.columns.join( + other.columns, how="outer", return_indexers=True + ) + + binop_result_ids = [] + block = joined_index._block + for left_index, right_index in zip(lcol_indexer, rcol_indexer): + if left_index >= 0 and right_index >= 0: # -1 indices indicate missing + left_col_id = self._block.value_columns[left_index] + right_col_id = other._block.value_columns[right_index] + block, result_col_id = block.apply_binary_op( + get_column_left(left_col_id), + get_column_right(right_col_id), + op, + ) + binop_result_ids.append(result_col_id) + elif left_index >= 0: + dtype = self.dtypes[left_index] + block, null_col_id = block.create_constant(None, dtype=dtype) + binop_result_ids.append(null_col_id) + elif right_index >= 0: + dtype = other.dtypes[right_index] + block, null_col_id = block.create_constant(None, dtype=dtype) + binop_result_ids.append(null_col_id) + else: + # Should not be possible + raise ValueError("No right or left index.") + + block = block.select_columns(binop_result_ids).with_column_labels(columns) + return DataFrame(block) + def eq(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: return self._apply_binop(other, ops.eq_op, axis=axis) @@ -619,7 +662,9 @@ def gt(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: __ge__ = ge def add( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: # TODO(swast): Support fill_value parameter. # TODO(swast): Support level parameter with MultiIndex. @@ -628,64 +673,92 @@ def add( __radd__ = __add__ = radd = add def sub( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.sub_op, axis=axis) __sub__ = subtract = sub def rsub( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.reverse(ops.sub_op), axis=axis) __rsub__ = rsub def mul( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.mul_op, axis=axis) __rmul__ = __mul__ = rmul = multiply = mul def truediv( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.div_op, axis=axis) div = divide = __truediv__ = truediv def rtruediv( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.reverse(ops.div_op), axis=axis) __rtruediv__ = rdiv = rtruediv def floordiv( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.floordiv_op, axis=axis) __floordiv__ = floordiv def rfloordiv( - self, other: float | int | bigframes.series.Series, axis: str | int = "columns" + self, + other: float | int | bigframes.series.Series | DataFrame, + axis: str | int = "columns", ) -> DataFrame: return self._apply_binop(other, ops.reverse(ops.floordiv_op), axis=axis) __rfloordiv__ = rfloordiv - def mod(self, other: int | bigframes.series.Series, axis: str | int = "columns") -> DataFrame: # type: ignore + def mod(self, other: int | bigframes.series.Series | DataFrame, axis: str | int = "columns") -> DataFrame: # type: ignore return self._apply_binop(other, ops.mod_op, axis=axis) - def rmod(self, other: int | bigframes.series.Series, axis: str | int = "columns") -> DataFrame: # type: ignore + def rmod(self, other: int | bigframes.series.Series | DataFrame, axis: str | int = "columns") -> DataFrame: # type: ignore return self._apply_binop(other, ops.reverse(ops.mod_op), axis=axis) __mod__ = mod __rmod__ = rmod + def pow( + self, other: int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.pow_op, axis=axis) + + def rpow( + self, other: int | bigframes.series.Series, axis: str | int = "columns" + ) -> DataFrame: + return self._apply_binop(other, ops.reverse(ops.pow_op), axis=axis) + + __pow__ = pow + + __rpow__ = rpow + def to_pandas( self, max_download_size: Optional[int] = None, @@ -1023,8 +1096,72 @@ def add_suffix(self, suffix: str, axis: int | str | None = None) -> DataFrame: axis = 1 if axis is None else axis return DataFrame(self._get_block().add_suffix(suffix, axis)) - def dropna(self) -> DataFrame: - return DataFrame(block_ops.dropna(self._block, how="any")) + def fillna(self, value=None) -> DataFrame: + return self._apply_binop(value, ops.fillna_op) + + def isin(self, values) -> DataFrame: + if utils.is_dict_like(values): + block = self._block + result_ids = [] + for col, label in zip(self._block.value_columns, self._block.column_labels): + if label in values.keys(): + value_for_key = values[label] + block, result_id = block.apply_unary_op( + col, ops.IsInOp(value_for_key, match_nulls=True), label + ) + result_ids.append(result_id) + else: + block, result_id = block.create_constant( + False, label=label, dtype=pandas.BooleanDtype() + ) + result_ids.append(result_id) + return DataFrame(block.select_columns(result_ids)).fillna(value=False) + elif utils.is_list_like(values): + return self._apply_unary_op(ops.IsInOp(values, match_nulls=True)).fillna( + value=False + ) + else: + raise TypeError( + "only list-like objects are allowed to be passed to " + f"isin(), you passed a [{type(values).__name__}]" + ) + + def dropna( + self, + *, + axis: int | str = 0, + inplace: bool = False, + how: str = "any", + ignore_index=False, + ) -> DataFrame: + if inplace: + raise NotImplementedError( + "'inplace'=True not supported. {constants.FEEDBACK_LINK}" + ) + if how not in ("any", "all"): + raise ValueError("'how' must be one of 'any', 'all'") + + axis_n = utils.get_axis_number(axis) + + if axis_n == 0: + result = block_ops.dropna(self._block, how=how) # type: ignore + if ignore_index: + result = result.reset_index() + return DataFrame(result) + else: + isnull_block = self._block.multi_apply_unary_op( + self._block.value_columns, ops.isnull_op + ) + if how == "any": + null_locations = DataFrame(isnull_block).any().to_pandas() + else: # 'all' + null_locations = DataFrame(isnull_block).all().to_pandas() + keep_columns = [ + col + for col, to_drop in zip(self._block.value_columns, null_locations) + if not to_drop + ] + return DataFrame(self._block.select_columns(keep_columns)) def any( self, @@ -1205,7 +1342,7 @@ def pivot( def stack(self): # TODO: support 'level' param by simply reordering levels such that selected level is last before passing to Block.stack. - # TODO: support 'dropna' param by executing dropna only conditionally + # TODO: match impl to pandas future_stack as described in pandas 2.1 release notes result_block = block_ops.dropna(self._block.stack(), how="all") if not isinstance(self.columns, pandas.MultiIndex): return bigframes.series.Series(result_block) @@ -1879,16 +2016,17 @@ def __array_ufunc__( """Used to support numpy ufuncs. See: https://numpy.org/doc/stable/reference/ufuncs.html """ - if ( - inputs[0] is not self - or method != "__call__" - or len(inputs) > 1 - or len(kwargs) > 0 - ): + if method != "__call__" or len(inputs) > 2 or len(kwargs) > 0: return NotImplemented - if ufunc in ops.NUMPY_TO_OP: + if len(inputs) == 1 and ufunc in ops.NUMPY_TO_OP: return self._apply_unary_op(ops.NUMPY_TO_OP[ufunc]) + if len(inputs) == 2 and ufunc in ops.NUMPY_TO_BINOP: + binop = ops.NUMPY_TO_BINOP[ufunc] + if inputs[0] is self: + return self._apply_binop(inputs[1], binop) + else: + return self._apply_binop(inputs[0], ops.reverse(binop)) return NotImplemented diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 9a20fd9ad74..2501d2b21f1 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -67,7 +67,7 @@ def fit( (X,) = utils.convert_to_dataframe(X) self._bqml_model = core.create_bqml_model( - train_X=X, + X_train=X, transforms=transforms, options=self._bqml_options, ) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 27727c9f81f..57f610c4c4b 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -22,7 +22,6 @@ from google.cloud import bigquery import bigframes -import bigframes.constants as constants from bigframes.ml import sql as ml_sql import bigframes.pandas as bpd @@ -53,10 +52,8 @@ def model(self) -> bigquery.Model: """Get the BQML model associated with this wrapper""" return self._model - @classmethod def _apply_sql( - cls, - session: bigframes.Session, + self, input_data: bpd.DataFrame, func: Callable[[str], str], ) -> bpd.DataFrame: @@ -78,13 +75,8 @@ def _apply_sql( include_index=True ) - if len(index_col_ids) != 1: - raise NotImplementedError( - f"Only exactly one index column is supported. {constants.FEEDBACK_LINK}" - ) - sql = func(source_sql) - df = session.read_gbq(sql, index_col=index_col_ids) + df = self._session.read_gbq(sql, index_col=index_col_ids) df.index.names = index_labels return df @@ -92,7 +84,6 @@ def _apply_sql( def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( - self._session, input_data, lambda source_sql: ml_sql.ml_predict( model_name=self.model_name, source_sql=source_sql @@ -102,7 +93,6 @@ def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame: def transform(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( - self._session, input_data, lambda source_sql: ml_sql.ml_transform( model_name=self.model_name, source_sql=source_sql @@ -116,7 +106,6 @@ def generate_text( ) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( - self._session, input_data, lambda source_sql: ml_sql.ml_generate_text( model_name=self.model_name, @@ -132,7 +121,6 @@ def generate_text_embedding( ) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( - self._session, input_data, lambda source_sql: ml_sql.ml_generate_text_embedding( model_name=self.model_name, @@ -209,16 +197,16 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: def create_bqml_model( - train_X: bpd.DataFrame, - train_y: Optional[bpd.DataFrame] = None, + X_train: bpd.DataFrame, + y_train: Optional[bpd.DataFrame] = None, transforms: Optional[Iterable[str]] = None, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> BqmlModel: """Create a session-temporary BQML model with the CREATE MODEL statement Args: - train_X: features columns for training - train_y: labels columns for training, if applicable + X_train: features columns for training + y_train: labels columns for training, if applicable transforms: an optional list of SQL expressions that implement preprocessing on top of the input data. Generates a BQML TRANSFORM clause options: a dict of options to configure the model. Generates a BQML OPTIONS @@ -227,19 +215,13 @@ def create_bqml_model( Returns: a BqmlModel, wrapping a trained model in BigQuery """ options = dict(options) - if train_y is None: - input_data = train_X + if y_train is None: + input_data = X_train else: - # TODO: handle case where train_y columns are renamed in the join - input_data = train_X.join(train_y, how="outer") - options.update({"INPUT_LABEL_COLS": train_y.columns.tolist()}) + input_data = X_train.join(y_train, how="outer") + options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) - # pickpocket session object from the dataframe - session = train_X._get_block().expr._session - - # TODO(garrettwu): add wrapper to select the feature columns - # for now, drop index to avoid including the index in feature columns - input_data = input_data.reset_index(drop=True) + session = X_train._get_block().expr._session source_sql = input_data.sql options_sql = ml_sql.options(**options) @@ -255,25 +237,25 @@ def create_bqml_model( def create_bqml_time_series_model( - train_X: bpd.DataFrame, - train_y: bpd.DataFrame, + X_train: bpd.DataFrame, + y_train: bpd.DataFrame, transforms: Optional[Iterable[str]] = None, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> BqmlModel: assert ( - train_X.columns.size == 1 + X_train.columns.size == 1 ), "Time series timestamp input must only contain 1 column." assert ( - train_y.columns.size == 1 + y_train.columns.size == 1 ), "Time stamp data input must only contain 1 column." options = dict(options) - input_data = train_X.join(train_y, how="outer") - options.update({"TIME_SERIES_TIMESTAMP_COL": train_X.columns.tolist()[0]}) - options.update({"TIME_SERIES_DATA_COL": train_y.columns.tolist()[0]}) - # pickpocket session object from the dataframe - session = train_X._get_block().expr._session + input_data = X_train.join(y_train, how="outer") + options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]}) + options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) + + session = X_train._get_block().expr._session source_sql = input_data.sql options_sql = ml_sql.options(**options) diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 16106d3a7b3..75b57f2e54f 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -61,7 +61,7 @@ def fit( (X,) = utils.convert_to_dataframe(X) self._bqml_model = core.create_bqml_model( - train_X=X, + X_train=X, transforms=transforms, options={ "model_type": "PCA", diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py index 89078f8267f..d4571eb3e5b 100644 --- a/bigframes/ml/imported.py +++ b/bigframes/ml/imported.py @@ -34,8 +34,12 @@ class TensorFlowModel(base.Predictor): model_path (str): GCS path that holds the model files.""" - def __init__(self, session: bigframes.Session, model_path: Optional[str] = None): - self.session = session + def __init__( + self, + session: Optional[bigframes.Session] = None, + model_path: Optional[str] = None, + ): + self.session = session or bpd.get_global_session() self.model_path = model_path self._bqml_model: Optional[core.BqmlModel] = None @@ -112,8 +116,12 @@ class ONNXModel(base.Predictor): model_path (str): Cloud Storage path that holds the model files.""" - def __init__(self, session: bigframes.Session, model_path: Optional[str] = None): - self.session = session + def __init__( + self, + session: Optional[bigframes.Session] = None, + model_path: Optional[str] = None, + ): + self.session = session or bpd.get_global_session() self.model_path = model_path self._bqml_model: Optional[core.BqmlModel] = None diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index def97b56ff0..345e3deb728 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -16,10 +16,10 @@ from __future__ import annotations -from typing import cast, Union +from typing import cast, Optional, Union import bigframes -import bigframes.constants as constants +from bigframes import clients, constants from bigframes.core import blocks from bigframes.ml import base, core, utils import bigframes.pandas as bpd @@ -35,17 +35,43 @@ class PaLM2TextGenerator(base.Predictor): """PaLM2 text generator LLM model. Args: - session (BigQuery Session): - BQ session to create the model - connection_name (str): - connection to connect with remote service. str of the format ..""" - - def __init__(self, session: bigframes.Session, connection_name: str): - self.session = session - self.connection_name = connection_name + session (bigframes.Session or None): + BQ session to create the model. If None, use the global default session. + connection_name (str or None): + connection to connect with remote service. str of the format ... + if None, use default connection in session context. + """ + + def __init__( + self, + session: Optional[bigframes.Session] = None, + connection_name: Optional[str] = None, + ): + self.session = session or bpd.get_global_session() + self.connection_name = connection_name or self.session._bq_connection + self._bq_connection_manager = clients.BqConnectionManager( + self.session.bqconnectionclient, self.session.resourcemanagerclient + ) self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): + # Parse and create connection if needed. + if not self.connection_name: + raise ValueError( + "Must provide connection_name, either in constructor or through session options." + ) + connection_name_parts = self.connection_name.split(".") + if len(connection_name_parts) != 3: + raise ValueError( + f"connection_name must be of the format .., got {self.connection_name}." + ) + self._bq_connection_manager.create_bq_connection( + project_id=connection_name_parts[0], + location=connection_name_parts[1], + connection_id=connection_name_parts[2], + iam_role="aiplatform.user", + ) + options = { "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, } @@ -140,17 +166,43 @@ class PaLM2TextEmbeddingGenerator(base.Predictor): """PaLM2 text embedding generator LLM model. Args: - session (BigQuery Session): - BQ session to create the model - connection_name (str): - connection to connect with remote service. str of the format ..""" - - def __init__(self, session: bigframes.Session, connection_name: str): - self.session = session - self.connection_name = connection_name + session (bigframes.Session or None): + BQ session to create the model. If None, use the global default session. + connection_name (str or None): + connection to connect with remote service. str of the format ... + if None, use default connection in session context. + """ + + def __init__( + self, + session: Optional[bigframes.Session] = None, + connection_name: Optional[str] = None, + ): + self.session = session or bpd.get_global_session() + self.connection_name = connection_name or self.session._bq_connection + self._bq_connection_manager = clients.BqConnectionManager( + self.session.bqconnectionclient, self.session.resourcemanagerclient + ) self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): + # Parse and create connection if needed. + if not self.connection_name: + raise ValueError( + "Must provide connection_name, either in constructor or through session options." + ) + connection_name_parts = self.connection_name.split(".") + if len(connection_name_parts) != 3: + raise ValueError( + f"connection_name must be of the format .., got {self.connection_name}." + ) + self._bq_connection_manager.create_bq_connection( + project_id=connection_name_parts[0], + location=connection_name_parts[1], + connection_id=connection_name_parts[2], + iam_role="aiplatform.user", + ) + options = { "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, } diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index bcd82435823..feb7ff7835b 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -34,8 +34,8 @@ def _encode_value(v: Union[str, int, float, Iterable[str]]) -> str: raise ValueError(f"Unexpected value type. {constants.FEEDBACK_LINK}") -def _build_param_Iterable(**kwargs: Union[str, int, float, Iterable[str]]) -> str: - """Encode a dict of values into a formatted Iterable of KVPs for SQL""" +def _build_parameters(**kwargs: Union[str, int, float, Iterable[str]]) -> str: + """Encode a dict of values into a formatted Iterable of key-value pairs for SQL""" indent_str = " " param_strs = [f"{k}={_encode_value(v)}" for k, v in kwargs.items()] return "\n" + indent_str + f",\n{indent_str}".join(param_strs) @@ -43,10 +43,10 @@ def _build_param_Iterable(**kwargs: Union[str, int, float, Iterable[str]]) -> st def options(**kwargs: Union[str, int, float, Iterable[str]]) -> str: """Encode the OPTIONS clause for BQML""" - return f"OPTIONS({_build_param_Iterable(**kwargs)})" + return f"OPTIONS({_build_parameters(**kwargs)})" -def _build_struct_param_Iterable(**kwargs: Union[int, float]) -> str: +def _build_structs(**kwargs: Union[int, float]) -> str: """Encode a dict of values into a formatted STRUCT items for SQL""" indent_str = " " param_strs = [f"{v} AS {k}" for k, v in kwargs.items()] @@ -55,10 +55,10 @@ def _build_struct_param_Iterable(**kwargs: Union[int, float]) -> str: def struct_options(**kwargs: Union[int, float]) -> str: """Encode a BQ STRUCT as options.""" - return f"STRUCT({_build_struct_param_Iterable(**kwargs)})" + return f"STRUCT({_build_structs(**kwargs)})" -def _build_expr_Iterable(*expr_sqls: str) -> str: +def _build_expressions(*expr_sqls: str) -> str: """Encode a Iterable of SQL expressions into a formatted Iterable for SQL""" indent_str = " " return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls) @@ -66,7 +66,7 @@ def _build_expr_Iterable(*expr_sqls: str) -> str: def transform(*expr_sqls: str) -> str: """Encode the TRANSFORM clause for BQML""" - return f"TRANSFORM({_build_expr_Iterable(*expr_sqls)})" + return f"TRANSFORM({_build_expressions(*expr_sqls)})" def connection(conn_name: str) -> str: @@ -137,7 +137,7 @@ def alter_model( return "\n".join(parts) -def ml_evaluate(model_name: str, source_sql: Union[str, None] = None) -> str: +def ml_evaluate(model_name: str, source_sql: Optional[str] = None) -> str: """Encode ML.EVALUATE for BQML""" if source_sql is None: return f"""SELECT * FROM ML.EVALUATE(MODEL `{model_name}`)""" diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 9305cf1ddab..fa43f725f64 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -38,6 +38,7 @@ # FLOAT64 has 11 exponent bits, so max values is about 2**(2**10) # ln(2**(2**10)) == (2**10)*ln(2) ~= 709.78, so EXP(x) for x>709.78 will overflow. _FLOAT64_EXP_BOUND = typing.cast(ibis_types.NumericValue, ibis_types.literal(709.78)) +_INT64_EXP_BOUND = typing.cast(ibis_types.NumericValue, ibis_types.literal(43.6)) BinaryOp = typing.Callable[[ibis_types.Value, ibis_types.Value], ibis_types.Value] TernaryOp = typing.Callable[ @@ -538,12 +539,27 @@ def __init__(self, values, match_nulls: bool = True): self._match_nulls = match_nulls def _as_ibis(self, x: ibis_types.Value): - if self._match_nulls and any(is_null(value) for value in self._values): - return x.isnull() | x.isin( - [val for val in self._values if not is_null(val)] - ) + contains_nulls = any(is_null(value) for value in self._values) + matchable_ibis_values = [] + for item in self._values: + if not is_null(item): + try: + # we want values that *could* be cast to the dtype, but we don't want + # to actually cast it, as that could be lossy (eg float -> int) + item_inferred_type = ibis.literal(item).type() + if ( + x.type() == item_inferred_type + or x.type().is_numeric() + and item_inferred_type.is_numeric() + ): + matchable_ibis_values.append(item) + except TypeError: + pass + + if self._match_nulls and contains_nulls: + return x.isnull() | x.isin(matchable_ibis_values) else: - return x.isin(self._values) + return x.isin(matchable_ibis_values) class BinopPartialRight(UnaryOp): @@ -746,6 +762,94 @@ def div_op( ) +@short_circuit_nulls(ibis_dtypes.float) +def pow_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + if x.type().is_integer() and y.type().is_integer(): + return _int_pow_op(x, y) + else: + return _float_pow_op(x, y) + + +def _int_pow_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + # Need to avoid any error cases - should produce NaN instead + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/mathematical_functions#pow + x_as_decimal = typing.cast( + ibis_types.NumericValue, + x.cast(ibis_dtypes.Decimal(precision=38, scale=9, nullable=True)), + ) + y_val = typing.cast(ibis_types.NumericValue, y) + + # BQ POW() function outputs FLOAT64, which can lose precision. + # Therefore, we do math in NUMERIC and cast back down after. + # Also, explicit bounds checks, pandas will silently overflow. + pow_result = x_as_decimal**y_val + overflow_cond = (pow_result > _ibis_num((2**63) - 1)) | ( + pow_result < _ibis_num(-(2**63)) + ) + + return ( + ibis.case() + .when((overflow_cond), ibis.null()) + .else_(pow_result.cast(ibis_dtypes.int64)) + .end() + ) + + +def _float_pow_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + # Most conditions here seek to prevent calling BQ POW with inputs that would generate errors. + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/mathematical_functions#pow + x_val = typing.cast(ibis_types.NumericValue, x) + y_val = typing.cast(ibis_types.NumericValue, y) + + overflow_cond = (x_val != _ZERO) & ((y_val * x_val.abs().ln()) > _FLOAT64_EXP_BOUND) + + # Float64 lose integer precision beyond 2**53, beyond this insufficient precision to get parity + exp_too_big = y_val.abs() > _ibis_num(2**53) + # Treat very large exponents as +=INF + norm_exp = exp_too_big.ifelse(_INF * y_val.sign(), y_val) + + pow_result = x_val**norm_exp + + # This cast is dangerous, need to only excuted where y_val has been bounds-checked + # Ibis needs try_cast binding to bq safe_cast + exponent_is_whole = y_val.cast(ibis_dtypes.int64) == y_val + odd_exponent = (x_val < _ZERO) & ( + y_val.cast(ibis_dtypes.int64) % _ibis_num(2) == _ibis_num(1) + ) + infinite_base = x_val.abs() == _INF + + return ( + ibis.case() + # Might be able to do something more clever with x_val==0 case + .when(y_val == _ZERO, _ibis_num(1)) + .when( + x_val == _ibis_num(1), _ibis_num(1) + ) # Need to ignore exponent, even if it is NA + .when( + (x_val == _ZERO) & (y_val < _ZERO), _INF + ) # This case would error POW function in BQ + .when(infinite_base, pow_result) + .when( + exp_too_big, pow_result + ) # Bigquery can actually handle the +-inf cases gracefully + .when((x_val < _ZERO) & (~exponent_is_whole), _NAN) + .when( + overflow_cond, _INF * odd_exponent.ifelse(_ibis_num(-1), _ibis_num(1)) + ) # finite overflows would cause bq to error + .else_(pow_result) + .end() + ) + + @short_circuit_nulls(ibis_dtypes.bool) def lt_op( x: ibis_types.Value, @@ -880,6 +984,15 @@ def partial_right(op: BinaryOp, scalar: typing.Any) -> UnaryOp: return BinopPartialRight(op, scalar) +NUMPY_TO_BINOP: typing.Final = { + np.add: add_op, + np.subtract: sub_op, + np.multiply: mul_op, + np.divide: div_op, + np.power: pow_op, +} + + # Ternary ops def where_op( original: ibis_types.Value, diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 280fce11126..e1a23e67a1c 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -52,6 +52,7 @@ import bigframes.series import bigframes.session import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat +import third_party.bigframes_vendored.pandas.core.reshape.merge as vendored_pandas_merge import third_party.bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile @@ -132,6 +133,37 @@ def cut( cut.__doc__ = vendored_pandas_tile.cut.__doc__ +def merge( + left: DataFrame, + right: DataFrame, + how: Literal[ + "inner", + "left", + "outer", + "right", + ] = "inner", + on: Optional[str] = None, + *, + left_on: Optional[str] = None, + right_on: Optional[str] = None, + sort: bool = False, + suffixes: tuple[str, str] = ("_x", "_y"), +) -> DataFrame: + return bigframes.core.joins.merge( + left, + right, + how=how, + on=on, + left_on=left_on, + right_on=right_on, + sort=sort, + suffixes=suffixes, + ) + + +merge.__doc__ = vendored_pandas_merge.merge.__doc__ + + def _set_default_session_location_if_possible(query): # Set the location as per the query if this is the first query the user is # running and: @@ -390,7 +422,6 @@ def read_gbq_function(function_name: str): read_gbq_function.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_function) - # pandas dtype attributes NA = pandas.NA BooleanDtype = pandas.BooleanDtype @@ -421,6 +452,7 @@ def read_gbq_function(function_name: str): __all___ = [ # Functions "concat", + "merge", "read_csv", "read_gbq", "read_gbq_function", diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 2a4b919dab0..6932e5b5805 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -26,7 +26,6 @@ import sys import tempfile import textwrap -import time from typing import List, NamedTuple, Optional, Sequence, TYPE_CHECKING if TYPE_CHECKING: @@ -49,6 +48,7 @@ import ibis.expr.operations as ops import ibis.expr.rules as rlz +from bigframes import clients import bigframes.constants as constants # TODO(shobs): Change the min log level to INFO after the development stabilizes @@ -167,35 +167,22 @@ def __init__( self._bq_location = bq_location self._bq_dataset = bq_dataset self._bq_client = bq_client - self._bq_connection_client = bq_connection_client self._bq_connection_id = bq_connection_id - self._cloud_resource_manager_client = cloud_resource_manager_client + self._bq_connection_manager = clients.BqConnectionManager( + bq_connection_client, cloud_resource_manager_client + ) def create_bq_remote_function( self, input_args, input_types, output_type, endpoint, bq_function_name ): """Create a BigQuery remote function given the artifacts of a user defined function and the http endpoint of a corresponding cloud function.""" - # TODO(shobs): The below command to enable BigQuery Connection API needs - # to be automated. Disabling for now since most target users would not - # have the privilege to enable API in a project. - # log("Making sure BigQuery Connection API is enabled") - # if os.system("gcloud services enable bigqueryconnection.googleapis.com"): - # raise ValueError("Failed to enable BigQuery Connection API") - - # If the intended connection does not exist then create it - service_account_id = self.get_service_account_if_connection_exists() - if service_account_id: - logger.info(f"Connector {self._bq_connection_id} already exists") - else: - connection_name, service_account_id = self.create_bq_connection() - logger.info( - f"Created BQ connection {connection_name} with service account id: {service_account_id}" - ) - - # Ensure IAM role on the BQ connection - # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function - self._ensure_iam_binding(service_account_id, "run.invoker") + self._bq_connection_manager.create_bq_connection( + self._gcp_project_id, + self._bq_location, + self._bq_connection_id, + "run.invoker", + ) # Create BQ function # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2 @@ -237,86 +224,6 @@ def get_cloud_function_endpoint(self, name): pass return None - # Introduce retries to accommodate transient errors like etag mismatch, - # which can be caused by concurrent operation on the same resource, and - # manifests with message like: - # google.api_core.exceptions.Aborted: 409 There were concurrent policy - # changes. Please retry the whole read-modify-write with exponential - # backoff. The request's ETag '\007\006\003,\264\304\337\272' did not match - # the current policy's ETag '\007\006\003,\3750&\363'. - @google.api_core.retry.Retry( - predicate=google.api_core.retry.if_exception_type( - google.api_core.exceptions.Aborted - ), - initial=10, - maximum=20, - multiplier=2, - timeout=60, - ) - def _ensure_iam_binding(self, service_account: str, role: str): - """Ensure necessary IAM role is configured on a service account.""" - project = f"projects/{self._gcp_project_id}" - service_account = f"serviceAccount:{service_account}" - role = f"roles/{role}" - request = google.iam.v1.iam_policy_pb2.GetIamPolicyRequest(resource=project) - policy = self._cloud_resource_manager_client.get_iam_policy(request=request) - - # Check if the binding already exists, and if does, do nothing more - for binding in policy.bindings: - if binding.role == role: - if service_account in binding.members: - return - - # Create a new binding - new_binding = google.iam.v1.policy_pb2.Binding( - role=role, members=[service_account] - ) - policy.bindings.append(new_binding) - request = google.iam.v1.iam_policy_pb2.SetIamPolicyRequest( - resource=project, policy=policy - ) - self._cloud_resource_manager_client.set_iam_policy(request=request) - - # We would wait for the IAM policy change to take effect - # https://cloud.google.com/iam/docs/access-change-propagation - logger.info( - f"Waiting {self._iam_wait_seconds} seconds for IAM to take effect.." - ) - time.sleep(self._iam_wait_seconds) - - def create_bq_connection(self): - """Create the BigQuery Connection and returns corresponding service account id.""" - client = self._bq_connection_client - connection = bigquery_connection_v1.Connection( - cloud_resource=bigquery_connection_v1.CloudResourceProperties() - ) - request = bigquery_connection_v1.CreateConnectionRequest( - parent=client.common_location_path(self._gcp_project_id, self._bq_location), - connection_id=self._bq_connection_id, - connection=connection, - ) - connection = client.create_connection(request) - return connection.name, connection.cloud_resource.service_account_id - - def get_service_account_if_connection_exists(self) -> Optional[str]: - """Check if the BigQuery Connection exists.""" - client = self._bq_connection_client - request = bigquery_connection_v1.GetConnectionRequest( - name=client.connection_path( - self._gcp_project_id, self._bq_location, self._bq_connection_id - ) - ) - - service_account = None - try: - service_account = client.get_connection( - request=request - ).cloud_resource.service_account_id - except google.api_core.exceptions.NotFound: - pass - - return service_account - def generate_udf_code(self, def_, dir): """Generate serialized bytecode using cloudpickle given a udf.""" udf_code_file_name = "udf.py" @@ -825,7 +732,7 @@ def remote_function( # A connection is required for BQ remote function # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function if not bigquery_connection and session: - bigquery_connection = session._remote_udf_connection # type: ignore + bigquery_connection = session._bq_connection # type: ignore if not bigquery_connection: raise ValueError( "BigQuery connection must be provided, either directly or via session. " diff --git a/bigframes/series.py b/bigframes/series.py index 49b0a5b1f08..8eadee37ed0 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -387,9 +387,24 @@ def rank( ) -> Series: return Series(block_ops.rank(self._block, method, na_option, ascending)) - def fillna(self, value=None) -> "Series" | None: + def fillna(self, value=None) -> Series: return self._apply_binary_op(value, ops.fillna_op) + def dropna( + self, + *, + axis: int = 0, + inplace: bool = False, + how: typing.Optional[str] = None, + ignore_index: bool = False, + ) -> Series: + if inplace: + raise NotImplementedError("'inplace'=True not supported") + result = block_ops.dropna(self._block, how="any") + if ignore_index: + result = result.reset_index() + return Series(result) + def head(self, n: int = 5) -> Series: return typing.cast(Series, self.iloc[0:n]) @@ -547,6 +562,18 @@ def floordiv(self, other: float | int | Series) -> Series: def rfloordiv(self, other: float | int | Series) -> Series: return self._apply_binary_op(other, ops.reverse(ops.floordiv_op)) + def __pow__(self, other: float | int | Series) -> Series: + return self.pow(other) + + def __rpow__(self, other: float | int | Series) -> Series: + return self.rpow(other) + + def pow(self, other: float | int | Series) -> Series: + return self._apply_binary_op(other, ops.pow_op) + + def rpow(self, other: float | int | Series) -> Series: + return self._apply_binary_op(other, ops.reverse(ops.pow_op)) + def __lt__(self, other: float | int | Series) -> Series: # type: ignore return self.lt(other) @@ -843,23 +870,15 @@ def argmin(self) -> scalars.Scalar: @property def is_monotonic_increasing(self) -> bool: - period = 1 - window = bigframes.core.WindowSpec( - preceding=period, - following=None, + return typing.cast( + bool, self._block.is_monotonic_increasing(self._value_column) ) - shifted_series = self._apply_window_op(agg_ops.ShiftOp(period), window) - return self.notna().__and__(self >= shifted_series).all() @property def is_monotonic_decreasing(self) -> bool: - period = 1 - window = bigframes.core.WindowSpec( - preceding=period, - following=None, + return typing.cast( + bool, self._block.is_monotonic_decreasing(self._value_column) ) - shifted_series = self._apply_window_op(agg_ops.ShiftOp(period), window) - return self.notna().__and__(self <= shifted_series).all() def __getitem__(self, indexer): # TODO: enforce stricter alignment, should fail if indexer is missing any keys. @@ -1105,9 +1124,12 @@ def mask(self, cond, other=None) -> Series: ) return self.where(~cond, other) - def to_frame(self) -> bigframes.dataframe.DataFrame: + def to_frame(self, name: blocks.Label = None) -> bigframes.dataframe.DataFrame: + provided_name = name if name else self.name # To be consistent with Pandas, it assigns 0 as the column name if missing. 0 is the first element of RangeIndex. - block = self._block.with_column_labels([self.name] if self.name else ["0"]) + block = self._block.with_column_labels( + [provided_name] if provided_name else ["0"] + ) return bigframes.dataframe.DataFrame(block) def to_csv(self, path_or_buf=None, **kwargs) -> typing.Optional[str]: @@ -1191,6 +1213,57 @@ def to_string( def to_xarray(self): return self.to_pandas().to_xarray() + def _throw_if_index_contains_duplicates( + self, error_message: typing.Optional[str] = None + ) -> None: + duplicates_block, _ = block_ops.indicate_duplicates( + self._get_block(), self._get_block().index_columns + ) + duplicates_block = duplicates_block.with_column_labels( + ["values", "is_duplicate"] + ) + duplicates_df = bigframes.dataframe.DataFrame(duplicates_block) + if duplicates_df["is_duplicate"].any(): + error_message = ( + error_message + if error_message + else "Index contains duplicate entries, but uniqueness is required." + ) + raise pandas.errors.InvalidIndexError(error_message) + + def map( + self, + arg: typing.Union[Mapping, Series], + na_action: Optional[str] = None, + *, + verify_integrity: bool = False, + ) -> Series: + if na_action: + raise NotImplementedError( + f"Non-None na_action argument is not yet supported for Series.map. {constants.FEEDBACK_LINK}" + ) + if isinstance(arg, Series): + if verify_integrity: + error_message = "When verify_integrity is True in Series.map, index of arg parameter must not have duplicate entries." + arg._throw_if_index_contains_duplicates(error_message=error_message) + map_df = bigframes.dataframe.DataFrame(arg._block) + map_df = map_df.rename(columns={arg.name: self.name}) + elif isinstance(arg, Mapping): + map_df = bigframes.dataframe.DataFrame( + {"keys": list(arg.keys()), self.name: list(arg.values())}, + session=self._get_block().expr._session, + ) + map_df = map_df.set_index("keys") + elif callable(arg): + return self.apply(arg) + else: + # Mirroring pandas, call the uncallable object + arg() # throws TypeError: object is not callable + + self_df = self.to_frame(name="series") + result_df = self_df.join(map_df, on="series") + return result_df[self.name] + def __array_ufunc__( self, ufunc: numpy.ufunc, method: str, *inputs, **kwargs ) -> Series: @@ -1198,16 +1271,17 @@ def __array_ufunc__( See: https://numpy.org/doc/stable/reference/ufuncs.html """ # Only __call__ supported with zero arguments - if ( - inputs[0] is not self - or method != "__call__" - or len(inputs) > 1 - or len(kwargs) > 0 - ): + if method != "__call__" or len(inputs) > 2 or len(kwargs) > 0: return NotImplemented - if ufunc in ops.NUMPY_TO_OP: + if len(inputs) == 1 and ufunc in ops.NUMPY_TO_OP: return self._apply_unary_op(ops.NUMPY_TO_OP[ufunc]) + if len(inputs) == 2 and ufunc in ops.NUMPY_TO_BINOP: + binop = ops.NUMPY_TO_BINOP[ufunc] + if inputs[0] is self: + return self._apply_binary_op(inputs[1], binop) + else: + return self._apply_binary_op(inputs[0], ops.reverse(binop)) return NotImplemented diff --git a/bigframes/session.py b/bigframes/session.py index ac2f8fa53ab..17444077726 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -270,9 +270,22 @@ class Session( third_party_pandas_readers.ReaderIOMixin, ): """Establishes a BigQuery connection to capture a group of job activities related to - DataFrames.""" + DataFrames. - def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): + Args: + context (bigframes._config.bigquery_options.BigQueryOptions): + Configuration adjusting how to connect to BigQuery and related + APIs. Note that some options are ignored if ``clients_provider`` is + set. + clients_provider (bigframes.session.ClientsProvider): + An object providing client library objects. + """ + + def __init__( + self, + context: Optional[bigquery_options.BigQueryOptions] = None, + clients_provider: Optional[ClientsProvider] = None, + ): if context is None: context = bigquery_options.BigQueryOptions() @@ -288,12 +301,15 @@ def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): # Instantiate a clients provider to help with cloud clients that will be # used in the future operations in the session - self._clients_provider = ClientsProvider( - project=context.project, - location=self._location, - use_regional_endpoints=context.use_regional_endpoints, - credentials=context.credentials, - ) + if clients_provider: + self._clients_provider = clients_provider + else: + self._clients_provider = ClientsProvider( + project=context.project, + location=self._location, + use_regional_endpoints=context.use_regional_endpoints, + credentials=context.credentials, + ) self._create_and_bind_bq_session() self.ibis_client = typing.cast( @@ -305,7 +321,7 @@ def __init__(self, context: Optional[bigquery_options.BigQueryOptions] = None): ), ) - self._remote_udf_connection = context.remote_udf_connection + self._bq_connection = context.bq_connection # Now that we're starting the session, don't allow the options to be # changed. @@ -381,10 +397,10 @@ def close(self): try: query_job = self.bqclient.query(abort_session_query) query_job.result() # blocks until finished - except google.api_core.exceptions.BadRequest as e: + except google.api_core.exceptions.BadRequest as exc: # Ignore the exception when the BQ session itself has expired # https://cloud.google.com/bigquery/docs/sessions-terminating#auto-terminate_a_session - if not e.message.startswith( + if not exc.message.startswith( f"Session {self._session_id} has expired and is no longer available." ): raise diff --git a/notebooks/experimental/longer_ml_demo.ipynb b/notebooks/experimental/longer_ml_demo.ipynb index c4d133421f6..793ff58ecdf 100644 --- a/notebooks/experimental/longer_ml_demo.ipynb +++ b/notebooks/experimental/longer_ml_demo.ipynb @@ -1373,9 +1373,9 @@ } ], "source": [ - "train_x = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", - "train_y = training_data[['body_mass_g']]\n", - "model.fit(train_x, train_y)\n", + "X_train = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "y_train = training_data[['body_mass_g']]\n", + "model.fit(X_train, y_train)\n", "model" ] }, @@ -1453,7 +1453,7 @@ } ], "source": [ - "model.score(train_x, train_y)" + "model.score(X_train, y_train)" ] }, { @@ -1490,12 +1490,12 @@ "# lets define a preprocessing step that adjust the linear measurements to use the cube\n", "'''\n", "def cubify(penguin_df):\n", - " penguin_df.culmen_length_mm = train_x.culmen_length_mm.pow(3)\n", - " penguin_df.culmen_depth_mm = train_x.culmen_depth_mm.pow(3)\n", - " penguin_df.flipper_length_mm = train_x.flipper_length_mm.pow(3)\n", + " penguin_df.culmen_length_mm = X_train.culmen_length_mm.pow(3)\n", + " penguin_df.culmen_depth_mm = X_train.culmen_depth_mm.pow(3)\n", + " penguin_df.flipper_length_mm = X_train.flipper_length_mm.pow(3)\n", "\n", - "cubify(train_x)\n", - "train_x\n", + "cubify(X_train)\n", + "X_train\n", "'''" ] }, @@ -1519,7 +1519,7 @@ "source": [ "# AS ABOVE, SKIP FOR NOW\n", "'''\n", - "model.fit(train_x, train_y)\n", + "model.fit(X_train, y_train)\n", "model.evaluate()\n", "'''" ] @@ -1756,15 +1756,15 @@ "training_data = training_data.dropna()\n", "\n", "# And we'll include species in our features\n", - "train_x = training_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", - "train_y = training_data[['body_mass_g']]\n", - "model.fit(train_x, train_y)\n", + "X_train = training_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "y_train = training_data[['body_mass_g']]\n", + "model.fit(X_train, y_train)\n", "\n", "# And we'll evaluate it on the Adelie penguins only\n", "adelie_data = training_data[training_data.species == \"Adelie Penguin (Pygoscelis adeliae)\"]\n", - "test_x = adelie_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", - "test_y = adelie_data[['body_mass_g']]\n", - "model.score(test_x, test_y)" + "X_test = adelie_data[['species', 'island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", + "y_test = adelie_data[['body_mass_g']]\n", + "model.score(X_test, y_test)" ] }, { @@ -1852,7 +1852,7 @@ " ('linreg', LinearRegression())\n", "])\n", "\n", - "pipe.fit(train_x, train_y)\n", + "pipe.fit(X_train, y_train)\n", "pipe.evaluate()" ] }, diff --git a/notebooks/getting_started/ml_fundamentals.ipynb b/notebooks/getting_started/ml_fundamentals.ipynb index ade50bcbc2e..e48aff1d57c 100644 --- a/notebooks/getting_started/ml_fundamentals.ipynb +++ b/notebooks/getting_started/ml_fundamentals.ipynb @@ -567,10 +567,10 @@ "name": "stdout", "output_type": "stream", "text": [ - "train_X shape: (267, 6)\n", - "test_X shape: (67, 6)\n", - "train_y shape: (267, 1)\n", - "test_y shape: (67, 1)\n" + "X_train shape: (267, 6)\n", + "X_test shape: (67, 6)\n", + "y_train shape: (267, 1)\n", + "y_test shape: (67, 1)\n" ] } ], @@ -584,14 +584,14 @@ "\n", "# This will split X and y into test and training sets, with 20% of the rows in the test set,\n", "# and the rest in the training set\n", - "train_X, test_X, train_y, test_y = train_test_split(\n", + "X_train, X_test, y_train, y_test = train_test_split(\n", " X, y, test_size=0.2)\n", "\n", "# Show the shape of the data after the split\n", - "print(f\"\"\"train_X shape: {train_X.shape}\n", - "test_X shape: {test_X.shape}\n", - "train_y shape: {train_y.shape}\n", - "test_y shape: {test_y.shape}\"\"\")" + "print(f\"\"\"X_train shape: {X_train.shape}\n", + "X_test shape: {X_test.shape}\n", + "y_train shape: {y_train.shape}\n", + "y_test shape: {y_test.shape}\"\"\")" ] }, { @@ -758,7 +758,7 @@ "source": [ "# If we look at the data, we can see that random rows were selected for\n", "# each side of the split\n", - "test_X.head(5)" + "X_test.head(5)" ] }, { @@ -880,8 +880,8 @@ } ], "source": [ - "# Note that this matches the rows in test_X\n", - "test_y.head(5)" + "# Note that this matches the rows in X_test\n", + "y_test.head(5)" ] }, { @@ -1242,11 +1242,11 @@ "numeric_columns = [\"culmen_length_mm\", \"culmen_depth_mm\", \"flipper_length_mm\"]\n", "\n", "scaler = StandardScaler()\n", - "scaler.fit(train_X[numeric_columns])\n", + "scaler.fit(X_train[numeric_columns])\n", "\n", "# Now, standardscaler should transform the numbers to have mean of zero\n", "# and standard deviation of one:\n", - "scaler.transform(train_X[numeric_columns])" + "scaler.transform(X_train[numeric_columns])" ] }, { @@ -1580,7 +1580,7 @@ ], "source": [ "# We can then repeat this transformation on new data\n", - "scaler.transform(test_X[numeric_columns])" + "scaler.transform(X_test[numeric_columns])" ] }, { @@ -2070,12 +2070,12 @@ " (\"encode\", OneHotEncoder(), [\"species\", \"sex\", \"island\"])])\n", "\n", "# Now we can fit all columns of the training data\n", - "preproc.fit(train_X)\n", + "preproc.fit(X_train)\n", "\n", - "processed_train_X = preproc.transform(train_X)\n", - "processed_test_X = preproc.transform(test_X)\n", + "processed_X_train = preproc.transform(X_train)\n", + "processed_X_test = preproc.transform(X_test)\n", "\n", - "processed_train_X" + "processed_X_train" ] }, { @@ -2347,12 +2347,12 @@ "linreg = LinearRegression()\n", "\n", "# Learn from the training data how to predict output y\n", - "linreg.fit(processed_train_X, train_y)\n", + "linreg.fit(processed_X_train, y_train)\n", "\n", "# Predict y for the test data\n", - "predicted_test_y = linreg.predict(processed_test_X)\n", + "predicted_y_test = linreg.predict(processed_X_test)\n", "\n", - "predicted_test_y" + "predicted_y_test" ] }, { @@ -2617,9 +2617,9 @@ "\n", "kmeans = KMeans(n_clusters=4)\n", "\n", - "kmeans.fit(processed_train_X)\n", + "kmeans.fit(processed_X_train)\n", "\n", - "kmeans.predict(processed_test_X)" + "kmeans.predict(processed_X_test)" ] }, { @@ -2924,10 +2924,10 @@ } ], "source": [ - "pipeline.fit(train_X, train_y)\n", + "pipeline.fit(X_train, y_train)\n", "\n", - "predicted_test_y = pipeline.predict(test_X)\n", - "predicted_test_y" + "predicted_y_test = pipeline.predict(X_test)\n", + "predicted_y_test" ] }, { @@ -3084,7 +3084,7 @@ ], "source": [ "# In the case of a pipeline, this will be equivalent to calling .score on the contained LinearRegression\n", - "pipeline.score(test_X, test_y)" + "pipeline.score(X_test, y_test)" ] }, { @@ -3156,7 +3156,7 @@ "source": [ "from bigframes.ml.metrics import r2_score\n", "\n", - "r2_score(test_y, predicted_test_y)" + "r2_score(y_test, predicted_y_test)" ] }, { diff --git a/notebooks/regression/sklearn_linear_regression.ipynb b/notebooks/regression/sklearn_linear_regression.ipynb index 8c0a21c0ff2..beb77ef092d 100644 --- a/notebooks/regression/sklearn_linear_regression.ipynb +++ b/notebooks/regression/sklearn_linear_regression.ipynb @@ -867,7 +867,7 @@ "feature_columns = training_data[['island', 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'sex']]\n", "label_columns = training_data[['body_mass_g']] \n", "\n", - "train_X, test_X, train_y, test_y = train_test_split(\n", + "X_train, X_test, y_train, y_test = train_test_split(\n", " feature_columns, label_columns, test_size=0.2)" ] }, @@ -940,7 +940,7 @@ "metadata": {}, "outputs": [], "source": [ - "pipeline.fit(train_X, train_y)" + "pipeline.fit(X_train, y_train)" ] }, { @@ -1040,9 +1040,9 @@ "source": [ "from bigframes.ml.metrics import r2_score\n", "\n", - "pred_y = pipeline.predict(test_X)\n", + "pred_y = pipeline.predict(X_test)\n", "\n", - "r2_score(test_y, pred_y)" + "r2_score(y_test, pred_y)" ] }, { diff --git a/notebooks/remote_functions/remote_function.ipynb b/notebooks/remote_functions/remote_function.ipynb index 76c87f8629c..06be0e7293e 100644 --- a/notebooks/remote_functions/remote_function.ipynb +++ b/notebooks/remote_functions/remote_function.ipynb @@ -2,27 +2,38 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 2, "id": "3613b1cd", "metadata": {}, "outputs": [], "source": [ + "# BigQuery table data on which notebook should be run\n", "TABLE='bigquery-public-data.stackoverflow.comments'\n", - "MAX_ROWS=1000000 # 1 Million" + "\n", + "# Change this up to test the scale, down to run the notebook faster\n", + "MAX_ROWS=100000" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 3, "id": "f1175247", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 30.5 s, sys: 4.25 s, total: 34.8 s\n", - "Wall time: 1min 49s\n" + "CPU times: user 25.4 s, sys: 2.5 s, total: 27.9 s\n", + "Wall time: 2min 31s\n" ] }, { @@ -130,7 +141,7 @@ "9 154 Sure, but what about a solution using O(1) mem... 8" ] }, - "execution_count": 2, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } @@ -149,7 +160,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 4, "id": "fd8a04a3", "metadata": {}, "outputs": [], @@ -180,7 +191,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 5, "id": "2b5e4568", "metadata": {}, "outputs": [ @@ -188,8 +199,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 5.48 s, sys: 27.6 ms, total: 5.51 s\n", - "Wall time: 5.49 s\n" + "CPU times: user 4.22 s, sys: 18.2 ms, total: 4.24 s\n", + "Wall time: 4.26 s\n" ] }, { @@ -308,7 +319,7 @@ "9 154 Sure, but what about a solution using O(1) mem... 8 19" ] }, - "execution_count": 4, + "execution_count": 5, "metadata": {}, "output_type": "execute_result" } @@ -322,18 +333,84 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 6, "id": "b81feaef", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n", + "/usr/local/google/home/shobs/code/bigframes1/venv/lib/python3.10/site-packages/google/auth/_default.py:78: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n", + " warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "2b1c9d671db14d2ca3be6a0b0c698430", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 6b0a39de-40a0-4dd4-be88-248bd8ebcd77 is RUNNING. \n", " 0\n", " \n", " \n", " 1\n", - " 35156124\n", - " Sorry I didn't include my timeout method befor...\n", + " 11013760\n", + " You *should* be concerned with the disk being ...\n", " 0\n", " \n", " \n", " 2\n", - " 35157401\n", - " As soon as I defined some sort of primary Key ...\n", + " 11013784\n", + " have you looked at `Integrate` or `NIntegrate`?\n", " 0\n", " \n", " \n", " 3\n", - " 35158649\n", - " @user3355243 I've edited it to give $values an...\n", + " 11015512\n", + " sorry, is a typo. The variable name is dist. (...\n", " 0\n", " \n", " \n", " 4\n", - " 35162039\n", - " I pasted my exes @Matt.\n", + " 11016238\n", + " Pfff, I'm having trouble with that formula too...\n", " 0\n", " \n", " \n", " 5\n", - " 35162396\n", - " @Gene - I do have separate fields, but I also ...\n", + " 11016276\n", + " Thanks thinksteep! Does this mean that by usin...\n", " 0\n", " \n", " \n", " 6\n", - " 35162907\n", - " could you please provide any kind of sketch ho...\n", + " 11016551\n", + " Jason, thanks for the reply. I've been workin...\n", " 0\n", " \n", " \n", " 7\n", - " 35166498\n", - " We use PhoneGap. Using Angular.\n", + " 11017973\n", + " I assume an `off` of 0.5 would put be exactly ...\n", " 0\n", " \n", " \n", " 8\n", - " 35170138\n", - " `decltype((int(Foo::*)(int))(&Foo::foo))` That...\n", + " 11018225\n", + " Thank you very much. I do worry too much abou...\n", " 0\n", " \n", " \n", " 9\n", - " 35172348\n", - " Yes, I found that one already. And I understan...\n", + " 11018370\n", + " @IanClelland, I edited my question a bit. The ...\n", " 0\n", " \n", " \n", "\n", + "

10 rows × 3 columns

\n", "[10 rows x 3 columns in total]" ], "text/plain": [ " id text score\n", - "0 35153602 Are you trying to access a nested array? 0\n", - "1 35156124 Sorry I didn't include my timeout method befor... 0\n", - "2 35157401 As soon as I defined some sort of primary Key ... 0\n", - "3 35158649 @user3355243 I've edited it to give $values an... 0\n", - "4 35162039 I pasted my exes @Matt. 0\n", - "5 35162396 @Gene - I do have separate fields, but I also ... 0\n", - "6 35162907 could you please provide any kind of sketch ho... 0\n", - "7 35166498 We use PhoneGap. Using Angular. 0\n", - "8 35170138 `decltype((int(Foo::*)(int))(&Foo::foo))` That... 0\n", - "9 35172348 Yes, I found that one already. And I understan... 0\n", + "0 11012908 you're welcome! according to the docs it shoul... 0\n", + "1 11013760 You *should* be concerned with the disk being ... 0\n", + "2 11013784 have you looked at `Integrate` or `NIntegrate`? 0\n", + "3 11015512 sorry, is a typo. The variable name is dist. (... 0\n", + "4 11016238 Pfff, I'm having trouble with that formula too... 0\n", + "5 11016276 Thanks thinksteep! Does this mean that by usin... 0\n", + "6 11016551 Jason, thanks for the reply. I've been workin... 0\n", + "7 11017973 I assume an `off` of 0.5 would put be exactly ... 0\n", + "8 11018225 Thank you very much. I do worry too much abou... 0\n", + "9 11018370 @IanClelland, I edited my question a bit. The ... 0\n", "\n", "[10 rows x 3 columns]" ] }, - "execution_count": 5, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -461,7 +539,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 7, "id": "55ed241e", "metadata": {}, "outputs": [ @@ -474,59 +552,73 @@ "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True)\n", " Decorator to turn a user defined function into a BigQuery remote function.\n", " \n", + " .. note::\n", + " Please make sure following is setup before using this API:\n", + " \n", + " 1. Have the below APIs enabled for your project:\n", + " \n", + " * BigQuery Connection API\n", + " * Cloud Functions API\n", + " * Cloud Run API\n", + " * Cloud Build API\n", + " * Artifact Registry API\n", + " * Cloud Resource Manager API\n", + " \n", + " This can be done from the cloud console (change `PROJECT_ID` to yours):\n", + " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", + " \n", + " Or from the gcloud CLI:\n", + " \n", + " `$ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com`\n", + " \n", + " 2. Have following IAM roles enabled for you:\n", + " \n", + " * BigQuery Data Editor (roles/bigquery.dataEditor)\n", + " * BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", + " * Cloud Functions Developer (roles/cloudfunctions.developer)\n", + " * Service Account User (roles/iam.serviceAccountUser)\n", + " * Storage Object Viewer (roles/storage.objectViewer)\n", + " * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", + " \n", + " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", + " \n", + " 1. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", + " 2. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", + " \n", + " Alternatively, the IAM could also be setup via the gcloud CLI:\n", + " \n", + " `$ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"`.\n", + " \n", " Args:\n", - " input_types : list(type)\n", + " input_types (list(type)):\n", " List of input data types in the user defined function.\n", - " output_type : type\n", + " output_type (type):\n", " Data type of the output in the user defined function.\n", - " dataset : str, Optional\n", - " Dataset to use to create a BigQuery function. It should be in\n", + " dataset (str, Optional):\n", + " Dataset in which to create a BigQuery remote function. It should be in\n", " `.` or `` format. If this\n", - " param is not provided then session dataset id would be used.\n", - " bigquery_connection : str, Optional\n", - " Name of the BigQuery connection. If it is pre created in the same\n", - " location as the `bigquery_client.location` then it would be used,\n", - " otherwise it would be created dynamically assuming the user has\n", - " necessary priviliges. If this param is not provided then the\n", - " bigquery connection from the session would be used.\n", - " reuse : bool, Optional\n", + " parameter is not provided then session dataset id is used.\n", + " bigquery_connection (str, Optional):\n", + " Name of the BigQuery connection. You should either have the\n", + " connection already created in the `location` you have chosen, or\n", + " you should have the Project IAM Admin role to enable the service\n", + " to create the connection for you if you need it.If this parameter is\n", + " not provided then the BigQuery connection from the session is used.\n", + " reuse (bool, Optional):\n", " Reuse the remote function if already exists.\n", " `True` by default, which will result in reusing an existing remote\n", " function (if any) that was previously created for the same udf.\n", " Setting it to false would force creating a unique remote function.\n", " If the required remote function does not exist then it would be\n", " created irrespective of this param.\n", + " Returns:\n", + " callable: A remote function object pointing to the cloud assets created\n", + " in the background to support the remote execution. The cloud assets can be\n", + " located through the following properties set in the object:\n", " \n", - " Notes:\n", - " Please make sure following is setup before using this API:\n", - " \n", - " 1. Have the below APIs enabled for your project:\n", - " a. BigQuery Connection API\n", - " b. Cloud Functions API\n", - " c. Cloud Run API\n", - " d. Cloud Build API\n", - " e. Artifact Registry API\n", - " f. Cloud Resource Manager API\n", + " `bigframes_cloud_function` - The google cloud function deployed for the user defined code.\n", " \n", - " This can be done from the cloud console (change PROJECT_ID to yours):\n", - " https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID\n", - " Or from the gcloud CLI:\n", - " $ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com\n", - " \n", - " 2. Have following IAM roles enabled for you:\n", - " a. BigQuery Data Editor (roles/bigquery.dataEditor)\n", - " b. BigQuery Connection Admin (roles/bigquery.connectionAdmin)\n", - " c. Cloud Functions Developer (roles/cloudfunctions.developer)\n", - " d. Service Account User (roles/iam.serviceAccountUser)\n", - " e. Storage Object Viewer (roles/storage.objectViewer)\n", - " f. Project IAM Admin (roles/resourcemanager.projectIamAdmin)\n", - " (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)\n", - " \n", - " 3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:\n", - " a. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection\n", - " b. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function\n", - " Alternatively, the IAM could also be setup via the gcloud CLI:\n", - " $ gcloud projects add-iam-policy-binding PROJECT_ID --member=\"serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID\" --role=\"roles/run.invoker\"\n", + " `bigframes_remote_function` - The bigquery remote function capable of calling into `bigframes_cloud_function`.\n", "\n" ] } @@ -539,7 +631,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 8, "id": "c9a8d03d", "metadata": {}, "outputs": [], @@ -555,7 +647,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 9, "id": "fbc27f81", "metadata": {}, "outputs": [ @@ -563,69 +655,17 @@ "name": "stderr", "output_type": "stream", "text": [ - "[INFO][2023-06-29 01:03:30,557][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-ac72c931423f68dddb4f84f6754e2b28 --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmposspvw3v --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", - "Preparing function...\n", - ".done.\n", - "Deploying function...\n", - "[Build]......................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................done\n", - "[Service]....................................................................................................done\n", - "Done.\n", - "You can view your function in the Cloud Console here: https://console.cloud.google.com/functions/details/us-central1/bigframes-ac72c931423f68dddb4f84f6754e2b28?project=bigframes-dev\n", - "\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "buildConfig:\n", - " build: projects/1084210331973/locations/us-central1/builds/57859622-4287-47bc-ab86-3cd3f4c47548\n", - " entryPoint: udf_http\n", - " runtime: python310\n", - " source:\n", - " storageSource:\n", - " bucket: gcf-v2-sources-1084210331973-us-central1\n", - " object: bigframes-ac72c931423f68dddb4f84f6754e2b28/function-source.zip\n", - " sourceProvenance:\n", - " resolvedStorageSource:\n", - " bucket: gcf-v2-sources-1084210331973-us-central1\n", - " generation: '1688000614202977'\n", - " object: bigframes-ac72c931423f68dddb4f84f6754e2b28/function-source.zip\n", - "environment: GEN_2\n", - "labels:\n", - " deployment-tool: cli-gcloud\n", - "name: projects/bigframes-dev/locations/us-central1/functions/bigframes-ac72c931423f68dddb4f84f6754e2b28\n", - "serviceConfig:\n", - " allTrafficOnLatestRevision: true\n", - " availableCpu: '0.1666'\n", - " availableMemory: 256M\n", - " ingressSettings: ALLOW_ALL\n", - " maxInstanceCount: 100\n", - " maxInstanceRequestConcurrency: 1\n", - " revision: bigframes-ac72c931423f68dddb4f84f6754e2b28-00001-faf\n", - " service: projects/bigframes-dev/locations/us-central1/services/bigframes-ac72c931423f68dddb4f84f6754e2b28\n", - " serviceAccountEmail: 1084210331973-compute@developer.gserviceaccount.com\n", - " timeoutSeconds: 60\n", - " uri: https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app\n", - "state: ACTIVE\n", - "updateTime: '2023-06-29T01:04:34.306070995Z'\n", - "url: https://us-central1-bigframes-dev.cloudfunctions.net/bigframes-ac72c931423f68dddb4f84f6754e2b28\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "[INFO][2023-06-29 01:04:37,239][bigframes.remote_function] Successfully created cloud function bigframes-ac72c931423f68dddb4f84f6754e2b28 with uri (https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app)\n", - "[INFO][2023-06-29 01:04:42,402][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", - "[INFO][2023-06-29 01:04:42,404][bigframes.remote_function] Creating BQ remote function: \n", - " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_ac72c931423f68dddb4f84f6754e2b28(n INT64)\n", - " RETURNS INT64\n", - " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", - " OPTIONS (\n", - " endpoint = \"https://bigframes-ac72c931423f68dddb4f84f6754e2b28-7krlje3eoq-uc.a.run.app\"\n", - " )\n", - "[INFO][2023-06-29 01:04:43,456][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_ac72c931423f68dddb4f84f6754e2b28\n" + "[INFO][2023-08-18 21:23:29,687][bigframes.remote_function] Creating new cloud function: gcloud functions deploy bigframes-b0feb1fbaf8188b64d7e70118d93c5d4 --gen2 --runtime=python310 --project=bigframes-dev --region=us-central1 --source=/tmp/tmpl2ewfnue --entry-point=udf_http --trigger-http --no-allow-unauthenticated\n", + "[INFO][2023-08-18 21:24:43,689][bigframes.remote_function] Successfully created cloud function bigframes-b0feb1fbaf8188b64d7e70118d93c5d4 with uri (https://bigframes-b0feb1fbaf8188b64d7e70118d93c5d4-7krlje3eoq-uc.a.run.app)\n", + "[INFO][2023-08-18 21:24:57,348][bigframes.remote_function] Connector bigframes-rf-conn already exists\n", + "[INFO][2023-08-18 21:24:57,351][bigframes.remote_function] Creating BQ remote function: \n", + " CREATE OR REPLACE FUNCTION `bigframes-dev.bigframes_temp_us`.bigframes_b0feb1fbaf8188b64d7e70118d93c5d4(n INT64)\n", + " RETURNS INT64\n", + " REMOTE WITH CONNECTION `bigframes-dev.us.bigframes-rf-conn`\n", + " OPTIONS (\n", + " endpoint = \"https://bigframes-b0feb1fbaf8188b64d7e70118d93c5d4-7krlje3eoq-uc.a.run.app\"\n", + " )\n", + "[INFO][2023-08-18 21:24:58,300][bigframes.remote_function] Created remote function bigframes-dev.bigframes_temp_us.bigframes_b0feb1fbaf8188b64d7e70118d93c5d4\n" ] }, { @@ -633,7 +673,7 @@ "output_type": "stream", "text": [ "\n", - "Wall time: 76.2659 s\n" + "Wall time: 89.0601 s\n" ] } ], @@ -672,7 +712,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 10, "id": "c1c9355f", "metadata": {}, "outputs": [ @@ -680,10 +720,52 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 22.2 ms, sys: 0 ns, total: 22.2 ms\n", - "Wall time: 22.4 ms\n" + "CPU times: user 16.8 ms, sys: 61 µs, total: 16.8 ms\n", + "Wall time: 17 ms\n" ] }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "2f840ad27c514ed19c759a004b32de33", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 0f421233-9d02-4746-bb39-86a3b0880aba is RUNNING. \n", " \n", " 0\n", - " 35153602\n", - " Are you trying to access a nested array?\n", + " 11012908\n", + " you're welcome! according to the docs it shoul...\n", " 0\n", " -1\n", " \n", " \n", " 1\n", - " 35156124\n", - " Sorry I didn't include my timeout method befor...\n", + " 11013760\n", + " You *should* be concerned with the disk being ...\n", " 0\n", " -1\n", " \n", " \n", " 2\n", - " 35157401\n", - " As soon as I defined some sort of primary Key ...\n", + " 11013784\n", + " have you looked at `Integrate` or `NIntegrate`?\n", " 0\n", " -1\n", " \n", " \n", " 3\n", - " 35158649\n", - " @user3355243 I've edited it to give $values an...\n", + " 11015512\n", + " sorry, is a typo. The variable name is dist. (...\n", " 0\n", " -1\n", " \n", " \n", " 4\n", - " 35162039\n", - " I pasted my exes @Matt.\n", + " 11016238\n", + " Pfff, I'm having trouble with that formula too...\n", " 0\n", " -1\n", " \n", " \n", " 5\n", - " 35162396\n", - " @Gene - I do have separate fields, but I also ...\n", + " 11016276\n", + " Thanks thinksteep! Does this mean that by usin...\n", " 0\n", " -1\n", " \n", " \n", " 6\n", - " 35162907\n", - " could you please provide any kind of sketch ho...\n", + " 11016551\n", + " Jason, thanks for the reply. I've been workin...\n", " 0\n", " -1\n", " \n", " \n", " 7\n", - " 35166498\n", - " We use PhoneGap. Using Angular.\n", + " 11017973\n", + " I assume an `off` of 0.5 would put be exactly ...\n", " 0\n", " -1\n", " \n", " \n", " 8\n", - " 35170138\n", - " `decltype((int(Foo::*)(int))(&Foo::foo))` That...\n", + " 11018225\n", + " Thank you very much. I do worry too much abou...\n", " 0\n", " -1\n", " \n", " \n", " 9\n", - " 35172348\n", - " Yes, I found that one already. And I understan...\n", + " 11018370\n", + " @IanClelland, I edited my question a bit. The ...\n", " 0\n", " -1\n", " \n", " \n", "\n", + "

10 rows × 4 columns

\n", "[10 rows x 4 columns in total]" ], "text/plain": [ " id text score n_prime\n", - "0 35153602 Are you trying to access a nested array? 0 -1\n", - "1 35156124 Sorry I didn't include my timeout method befor... 0 -1\n", - "2 35157401 As soon as I defined some sort of primary Key ... 0 -1\n", - "3 35158649 @user3355243 I've edited it to give $values an... 0 -1\n", - "4 35162039 I pasted my exes @Matt. 0 -1\n", - "5 35162396 @Gene - I do have separate fields, but I also ... 0 -1\n", - "6 35162907 could you please provide any kind of sketch ho... 0 -1\n", - "7 35166498 We use PhoneGap. Using Angular. 0 -1\n", - "8 35170138 `decltype((int(Foo::*)(int))(&Foo::foo))` That... 0 -1\n", - "9 35172348 Yes, I found that one already. And I understan... 0 -1\n", + "0 11012908 you're welcome! according to the docs it shoul... 0 -1\n", + "1 11013760 You *should* be concerned with the disk being ... 0 -1\n", + "2 11013784 have you looked at `Integrate` or `NIntegrate`? 0 -1\n", + "3 11015512 sorry, is a typo. The variable name is dist. (... 0 -1\n", + "4 11016238 Pfff, I'm having trouble with that formula too... 0 -1\n", + "5 11016276 Thanks thinksteep! Does this mean that by usin... 0 -1\n", + "6 11016551 Jason, thanks for the reply. I've been workin... 0 -1\n", + "7 11017973 I assume an `off` of 0.5 would put be exactly ... 0 -1\n", + "8 11018225 Thank you very much. I do worry too much abou... 0 -1\n", + "9 11018370 @IanClelland, I edited my question a bit. The ... 0 -1\n", "\n", "[10 rows x 4 columns]" ] }, - "execution_count": 9, + "execution_count": 10, "metadata": {}, "output_type": "execute_result" } @@ -810,9 +893,299 @@ "source": [ "%%time\n", "\n", + "# Let's apply the function to the dataframe\n", "df = df.assign(n_prime=df['score'].apply(nth_prime))\n", "df.head(10)" ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "2701cb81", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "bigframes-dev.bigframes_temp_us.bigframes_b0feb1fbaf8188b64d7e70118d93c5d4\n", + "projects/bigframes-dev/locations/us-central1/functions/bigframes-b0feb1fbaf8188b64d7e70118d93c5d4\n" + ] + } + ], + "source": [ + "# We can see the path to the BQ remote function and the google cloud function\n", + "# that was created under the hood\n", + "print(nth_prime.bigframes_remote_function)\n", + "print(nth_prime.bigframes_cloud_function)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "920fa18e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Help on function read_gbq_function in module bigframes.pandas:\n", + "\n", + "read_gbq_function(function_name: 'str')\n", + " Loads a BigQuery function from BigQuery.\n", + " \n", + " Then it can be applied to a DataFrame or Series.\n", + " \n", + " Args:\n", + " function_name (str):\n", + " the function's name in BigQuery in the format\n", + " `project_id.dataset_id.function_name`, or\n", + " `dataset_id.function_name` to load from the default project, or\n", + " `function_name` to load from the default project and the dataset\n", + " associated with the current session.\n", + " \n", + " Returns:\n", + " callable: A function object pointing to the BigQuery function read\n", + " from BigQuery.\n", + " \n", + " The object is similar to the one created by the `remote_function`\n", + " decorator, including the `bigframes_remote_function` property, but\n", + " not including the `bigframes_cloud_function` property.\n", + "\n" + ] + } + ], + "source": [ + "# Let's try to simulate a scenario in which user shares this remote funciton to\n", + "# their colleague who simply wants to reuse it. BigFrames provides an API to do\n", + "# so via `read_gbq_function`. Usage details are available via `help` command.\n", + "help(pd.read_gbq_function)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "a6c9da0a", + "metadata": {}, + "outputs": [], + "source": [ + "EXISTING_REMOTE_FUNCTION=nth_prime.bigframes_remote_function\n", + "\n", + "# Let's read the existing remote function in bigframes\n", + "nth_prime_existing = pd.read_gbq_function(EXISTING_REMOTE_FUNCTION)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "d7e7de7f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 10.9 ms, sys: 0 ns, total: 10.9 ms\n", + "Wall time: 11.4 ms\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "73d1a73593cb4115821ab128c221a48d", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job bec5f7d1-3df1-4292-8c68-c396bce7dc5d is RUNNING. \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtextscoren_primen_prime_again
011012908you're welcome! according to the docs it shoul...0-1-1
111013760You *should* be concerned with the disk being ...0-1-1
211013784have you looked at `Integrate` or `NIntegrate`?0-1-1
311015512sorry, is a typo. The variable name is dist. (...0-1-1
411016238Pfff, I'm having trouble with that formula too...0-1-1
511016276Thanks thinksteep! Does this mean that by usin...0-1-1
611016551Jason, thanks for the reply. I've been workin...0-1-1
711017973I assume an `off` of 0.5 would put be exactly ...0-1-1
811018225Thank you very much. I do worry too much abou...0-1-1
911018370@IanClelland, I edited my question a bit. The ...0-1-1
\n", + "

10 rows × 5 columns

\n", + "[10 rows x 5 columns in total]" + ], + "text/plain": [ + " id text score \\\n", + "0 11012908 you're welcome! according to the docs it shoul... 0 \n", + "1 11013760 You *should* be concerned with the disk being ... 0 \n", + "2 11013784 have you looked at `Integrate` or `NIntegrate`? 0 \n", + "3 11015512 sorry, is a typo. The variable name is dist. (... 0 \n", + "4 11016238 Pfff, I'm having trouble with that formula too... 0 \n", + "5 11016276 Thanks thinksteep! Does this mean that by usin... 0 \n", + "6 11016551 Jason, thanks for the reply. I've been workin... 0 \n", + "7 11017973 I assume an `off` of 0.5 would put be exactly ... 0 \n", + "8 11018225 Thank you very much. I do worry too much abou... 0 \n", + "9 11018370 @IanClelland, I edited my question a bit. The ... 0 \n", + "\n", + " n_prime n_prime_again \n", + "0 -1 -1 \n", + "1 -1 -1 \n", + "2 -1 -1 \n", + "3 -1 -1 \n", + "4 -1 -1 \n", + "5 -1 -1 \n", + "6 -1 -1 \n", + "7 -1 -1 \n", + "8 -1 -1 \n", + "9 -1 -1 \n", + "\n", + "[10 rows x 5 columns]" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "\n", + "# Let's apply the existing function to the dataframe\n", + "df = df.assign(n_prime_again=df['score'].apply(nth_prime_existing))\n", + "df.head(10)" + ] } ], "metadata": { @@ -831,7 +1204,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.9" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/tests/system/conftest.py b/tests/system/conftest.py index bc94e32e12d..3153bd1559e 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -849,7 +849,7 @@ def floats_pd(): { "float64_col": [ float("-inf"), - float("-inf"), + float("inf"), float("nan"), float(-234239487.4), float(-1.0), @@ -863,17 +863,33 @@ def floats_pd(): float(math.e), float(math.pi), float(234239487.4), + float(1.23124 * (2**70)), pd.NA, ] }, dtype=pd.Float64Dtype(), ) + # Index helps debug failed cases df.index = df.float64_col # Upload fails if index name same as column name df.index.name = None return df.float64_col +@pytest.fixture() +def floats_product_pd(floats_pd): + df = pd.merge(floats_pd, floats_pd, how="cross") + # Index helps debug failed cases + df = df.set_index([df.float64_col_x, df.float64_col_y]) + df.index.names = ["left", "right"] + return df + + @pytest.fixture() def floats_bf(session, floats_pd): return session.read_pandas(floats_pd.to_frame()).float64_col + + +@pytest.fixture() +def floats_product_bf(session, floats_product_pd): + return session.read_pandas(floats_product_pd) diff --git a/tests/system/large/ml/test_core.py b/tests/system/large/ml/test_core.py index ab33e5d7185..133af2dae4a 100644 --- a/tests/system/large/ml/test_core.py +++ b/tests/system/large/ml/test_core.py @@ -20,7 +20,7 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_df): df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -30,10 +30,10 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d "sex", ] ] - train_y = df[["body_mass_g"]] + y_train = df[["body_mass_g"]] model = bigframes.ml.core.create_bqml_model( - train_X, train_y, options={"model_type": "linear_reg"} + X_train, y_train, options={"model_type": "linear_reg"} ) # no data - report evaluation from the automatic data split @@ -85,22 +85,22 @@ def test_bqml_manual_preprocessing_e2e( session, dataset_id, penguins_df_default_index, new_penguins_df ): df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", ] ] - train_y = df[["body_mass_g"]] + y_train = df[["body_mass_g"]] transforms = [ bigframes.ml.sql.ml_standard_scaler(column, column) - for column in train_X.columns.tolist() + for column in X_train.columns.tolist() ] - transforms.extend(train_y.columns.tolist()) + transforms.extend(y_train.columns.tolist()) options = {"model_type": "linear_reg"} model = bigframes.ml.core.create_bqml_model( - train_X, train_y, transforms=transforms, options=options + X_train, y_train, transforms=transforms, options=options ) # no data - report evaluation from the automatic data split diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py index 88c5ccd2f0d..9b2872d6737 100644 --- a/tests/system/large/ml/test_ensemble.py +++ b/tests/system/large/ml/test_ensemble.py @@ -25,7 +25,7 @@ def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.XGBRegressor() df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -35,11 +35,11 @@ def test_xgbregressor_default_params(penguins_df_default_index, dataset_id): "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [97.368139], @@ -86,7 +86,7 @@ def test_xgbregressor_dart_booster_multiple_params( ) df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -96,11 +96,11 @@ def test_xgbregressor_dart_booster_multiple_params( "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -144,7 +144,7 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.XGBClassifier() df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -153,11 +153,11 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): "flipper_length_mm", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -201,7 +201,7 @@ def test_xgbclassifier_dart_booster_multiple_params( ) df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -210,11 +210,11 @@ def test_xgbclassifier_dart_booster_multiple_params( "flipper_length_mm", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -258,7 +258,7 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset model = bigframes.ml.ensemble.RandomForestRegressor() df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -268,11 +268,11 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -311,7 +311,7 @@ def test_randomforestregressor_multiple_params(penguins_df_default_index, datase ) df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -321,11 +321,11 @@ def test_randomforestregressor_multiple_params(penguins_df_default_index, datase "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "mean_absolute_error", @@ -366,7 +366,7 @@ def test_randomforestclassifier_default_params(penguins_df_default_index, datase model = bigframes.ml.ensemble.RandomForestClassifier() df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -375,11 +375,11 @@ def test_randomforestclassifier_default_params(penguins_df_default_index, datase "flipper_length_mm", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -418,7 +418,7 @@ def test_randomforestclassifier_multiple_params(penguins_df_default_index, datas ) df = penguins_df_default_index.dropna().sample(n=70) - train_X = df[ + X_train = df[ [ "species", "island", @@ -427,11 +427,11 @@ def test_randomforestclassifier_multiple_params(penguins_df_default_index, datas "flipper_length_mm", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", diff --git a/tests/system/large/ml/test_forecasting.py b/tests/system/large/ml/test_forecasting.py index d1e2d12296d..33b835e8522 100644 --- a/tests/system/large/ml/test_forecasting.py +++ b/tests/system/large/ml/test_forecasting.py @@ -21,9 +21,9 @@ def test_arima_plus_model_fit_score( time_series_df_default_index, dataset_id, new_time_series_df ): model = forecasting.ARIMAPlus() - train_X = time_series_df_default_index[["parsed_date"]] - train_y = time_series_df_default_index[["total_visits"]] - model.fit(train_X, train_y) + X_train = time_series_df_default_index[["parsed_date"]] + y_train = time_series_df_default_index[["total_visits"]] + model.fit(X_train, y_train) result = model.score( new_time_series_df[["parsed_date"]], new_time_series_df[["total_visits"]] diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py index 332b460fe51..3b90568450b 100644 --- a/tests/system/large/ml/test_linear_model.py +++ b/tests/system/large/ml/test_linear_model.py @@ -21,7 +21,7 @@ def test_linear_regression_configure_fit_score(penguins_df_default_index, datase model = bigframes.ml.linear_model.LinearRegression(fit_intercept=False) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -31,11 +31,11 @@ def test_linear_regression_configure_fit_score(penguins_df_default_index, datase "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [225.735767], @@ -66,7 +66,7 @@ def test_linear_regression_manual_split_configure_fit_score( model = bigframes.ml.linear_model.LinearRegression(fit_intercept=True) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -76,11 +76,11 @@ def test_linear_regression_manual_split_configure_fit_score( "sex", ] ] - train_y = df[["body_mass_g"]] - model.fit(train_X, train_y) + y_train = df[["body_mass_g"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { "mean_absolute_error": [225.735767], @@ -108,7 +108,7 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( ): model = bigframes.ml.linear_model.LogisticRegression() df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -117,11 +117,11 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( "flipper_length_mm", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { "precision": [0.58085], @@ -155,7 +155,7 @@ def test_logistic_regression_manual_split_configure_fit_score( model = bigframes.ml.linear_model.LogisticRegression(fit_intercept=True) df = penguins_df_default_index.dropna() - train_X = df[ + X_train = df[ [ "species", "island", @@ -165,11 +165,11 @@ def test_logistic_regression_manual_split_configure_fit_score( "body_mass_g", ] ] - train_y = df[["sex"]] - model.fit(train_X, train_y) + y_train = df[["sex"]] + model.fit(X_train, y_train) # Check score to ensure the model was fitted - result = model.score(train_X, train_y).to_pandas() + result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { "precision": [0.616753], diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index ebe768b685e..9ca5a2fd0ef 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -32,8 +32,8 @@ @pytest.fixture(scope="session") -def ml_connection() -> str: - return "bigframes-dev.us.bigframes-ml" +def bq_connection() -> str: + return "bigframes-dev.us.bigframes-rf-conn" @pytest.fixture(scope="session") @@ -198,33 +198,33 @@ def llm_text_df(session, llm_text_pandas_df): @pytest.fixture(scope="session") -def bqml_palm2_text_generator_model(session, ml_connection) -> core.BqmlModel: +def bqml_palm2_text_generator_model(session, bq_connection) -> core.BqmlModel: options = { "remote_service_type": "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1", } return core.create_bqml_remote_model( - session=session, connection_name=ml_connection, options=options + session=session, connection_name=bq_connection, options=options ) @pytest.fixture(scope="session") -def palm2_text_generator_model(session, ml_connection) -> llm.PaLM2TextGenerator: - return llm.PaLM2TextGenerator(session=session, connection_name=ml_connection) +def palm2_text_generator_model(session, bq_connection) -> llm.PaLM2TextGenerator: + return llm.PaLM2TextGenerator(session=session, connection_name=bq_connection) @pytest.fixture(scope="function") def ephemera_palm2_text_generator_model( - session, ml_connection + session, bq_connection ) -> llm.PaLM2TextGenerator: - return llm.PaLM2TextGenerator(session=session, connection_name=ml_connection) + return llm.PaLM2TextGenerator(session=session, connection_name=bq_connection) @pytest.fixture(scope="session") def palm2_embedding_generator_model( - session, ml_connection + session, bq_connection ) -> llm.PaLM2TextEmbeddingGenerator: return llm.PaLM2TextEmbeddingGenerator( - session=session, connection_name=ml_connection + session=session, connection_name=bq_connection ) @@ -247,10 +247,22 @@ def time_series_arima_plus_model( @pytest.fixture(scope="session") -def imported_tensorflow_model(session) -> imported.TensorFlowModel: +def imported_tensorflow_model_path() -> str: + return "gs://cloud-training-demos/txtclass/export/exporter/1549825580/*" + + +@pytest.fixture(scope="session") +def imported_onnx_model_path() -> str: + return "gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx" + + +@pytest.fixture(scope="session") +def imported_tensorflow_model( + session, imported_tensorflow_model_path +) -> imported.TensorFlowModel: return imported.TensorFlowModel( session=session, - model_path="gs://cloud-training-demos/txtclass/export/exporter/1549825580/*", + model_path=imported_tensorflow_model_path, ) @@ -263,8 +275,8 @@ def ephemera_imported_tensorflow_model(session) -> imported.TensorFlowModel: @pytest.fixture(scope="session") -def imported_onnx_model(session) -> imported.ONNXModel: +def imported_onnx_model(session, imported_onnx_model_path) -> imported.ONNXModel: return imported.ONNXModel( session=session, - model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx", + model_path=imported_onnx_model_path, ) diff --git a/tests/system/small/ml/test_ensemble.py b/tests/system/small/ml/test_ensemble.py index fde3cc431e1..bba083d98d9 100644 --- a/tests/system/small/ml/test_ensemble.py +++ b/tests/system/small/ml/test_ensemble.py @@ -25,7 +25,7 @@ def test_xgbregressor_model_score( penguins_xgbregressor_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -35,8 +35,8 @@ def test_xgbregressor_model_score( "body_mass_g", ] ] - test_y = df[["sex"]] - result = penguins_xgbregressor_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = penguins_xgbregressor_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [108.77582], @@ -62,7 +62,7 @@ def test_xgbregressor_model_score_series( penguins_xgbregressor_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -72,8 +72,8 @@ def test_xgbregressor_model_score_series( "body_mass_g", ] ] - test_y = df["sex"] - result = penguins_xgbregressor_model.score(test_X, test_y).to_pandas() + y_test = df["sex"] + result = penguins_xgbregressor_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [108.77582], @@ -120,7 +120,7 @@ def test_to_gbq_saved_xgbregressor_model_scores( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -130,8 +130,8 @@ def test_to_gbq_saved_xgbregressor_model_scores( "body_mass_g", ] ] - test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [109.016973], @@ -165,7 +165,7 @@ def test_xgbclassifier_model_score( penguins_xgbclassifier_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -175,8 +175,8 @@ def test_xgbclassifier_model_score( "body_mass_g", ] ] - test_y = df[["sex"]] - result = penguins_xgbclassifier_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = penguins_xgbclassifier_model.score(X_test, y_test).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -193,7 +193,7 @@ def test_xgbclassifier_model_score_series( penguins_xgbclassifier_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -203,8 +203,8 @@ def test_xgbclassifier_model_score_series( "body_mass_g", ] ] - test_y = df["sex"] - result = penguins_xgbclassifier_model.score(test_X, test_y).to_pandas() + y_test = df["sex"] + result = penguins_xgbclassifier_model.score(X_test, y_test).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -242,7 +242,7 @@ def test_to_gbq_saved_xgbclassifier_model_scores( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -252,8 +252,8 @@ def test_to_gbq_saved_xgbclassifier_model_scores( "body_mass_g", ] ] - test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "precision": [1.0], @@ -289,7 +289,7 @@ def test_randomforestregressor_model_score( penguins_randomforest_regressor_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -299,8 +299,8 @@ def test_randomforestregressor_model_score( "body_mass_g", ] ] - test_y = df[["sex"]] - result = penguins_randomforest_regressor_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = penguins_randomforest_regressor_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [317.031042], @@ -326,7 +326,7 @@ def test_randomforestregressor_model_score_series( penguins_randomforest_regressor_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -336,8 +336,8 @@ def test_randomforestregressor_model_score_series( "body_mass_g", ] ] - test_y = df["sex"] - result = penguins_randomforest_regressor_model.score(test_X, test_y).to_pandas() + y_test = df["sex"] + result = penguins_randomforest_regressor_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [317.031042], @@ -385,7 +385,7 @@ def test_to_gbq_saved_randomforestregressor_model_scores( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -395,8 +395,8 @@ def test_to_gbq_saved_randomforestregressor_model_scores( "body_mass_g", ] ] - test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [319.239235], @@ -434,7 +434,7 @@ def test_randomforestclassifier_model_score( penguins_randomforest_classifier_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -444,8 +444,8 @@ def test_randomforestclassifier_model_score( "body_mass_g", ] ] - test_y = df[["sex"]] - result = penguins_randomforest_classifier_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = penguins_randomforest_classifier_model.score(X_test, y_test).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -462,7 +462,7 @@ def test_randomforestclassifier_model_score_series( penguins_randomforest_classifier_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -472,8 +472,8 @@ def test_randomforestclassifier_model_score_series( "body_mass_g", ] ] - test_y = df["sex"] - result = penguins_randomforest_classifier_model.score(test_X, test_y).to_pandas() + y_test = df["sex"] + result = penguins_randomforest_classifier_model.score(X_test, y_test).to_pandas() TestCase().assertSequenceEqual(result.shape, (1, 6)) for col_name in [ "precision", @@ -512,7 +512,7 @@ def test_to_gbq_saved_randomforestclassifier_model_scores( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -522,8 +522,8 @@ def test_to_gbq_saved_randomforestclassifier_model_scores( "body_mass_g", ] ] - test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "precision": [0.636746], diff --git a/tests/system/small/ml/test_imported.py b/tests/system/small/ml/test_imported.py index 6274ab12450..d3055670668 100644 --- a/tests/system/small/ml/test_imported.py +++ b/tests/system/small/ml/test_imported.py @@ -25,6 +25,11 @@ def test_tensorflow_create_model(imported_tensorflow_model): assert imported_tensorflow_model is not None +def test_tensorflow_create_model_default_session(imported_tensorflow_model_path): + model = imported.TensorFlowModel(model_path=imported_tensorflow_model_path) + assert model is not None + + def test_tensorflow_model_predict(imported_tensorflow_model, llm_text_df): df = llm_text_df.rename(columns={"prompt": "input"}) result = imported_tensorflow_model.predict(df).to_pandas() @@ -61,6 +66,11 @@ def test_onnx_create_model(imported_onnx_model): assert imported_onnx_model is not None +def test_onnx_create_model_default_session(imported_onnx_model_path): + model = imported.TensorFlowModel(model_path=imported_onnx_model_path) + assert model is not None + + def test_onnx_model_predict(imported_onnx_model, onnx_iris_df): result = imported_onnx_model.predict(onnx_iris_df).to_pandas() value1 = np.array([0.9999993443489075, 0.0, 0.0]) diff --git a/tests/system/small/ml/test_linear_model.py b/tests/system/small/ml/test_linear_model.py index bbb7e2820c3..3a8232ed9ec 100644 --- a/tests/system/small/ml/test_linear_model.py +++ b/tests/system/small/ml/test_linear_model.py @@ -19,7 +19,7 @@ def test_linear_reg_model_score(penguins_linear_model, penguins_df_default_index): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -29,8 +29,8 @@ def test_linear_reg_model_score(penguins_linear_model, penguins_df_default_index "sex", ] ] - test_y = df[["body_mass_g"]] - result = penguins_linear_model.score(test_X, test_y).to_pandas() + y_test = df[["body_mass_g"]] + result = penguins_linear_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [225.817334], @@ -56,7 +56,7 @@ def test_linear_reg_model_score_series( penguins_linear_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -66,8 +66,8 @@ def test_linear_reg_model_score_series( "sex", ] ] - test_y = df["body_mass_g"] - result = penguins_linear_model.score(test_X, test_y).to_pandas() + y_test = df["body_mass_g"] + result = penguins_linear_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [225.817334], @@ -111,7 +111,7 @@ def test_to_gbq_saved_linear_reg_model_scores( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -121,8 +121,8 @@ def test_to_gbq_saved_linear_reg_model_scores( "sex", ] ] - test_y = df[["body_mass_g"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["body_mass_g"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "mean_absolute_error": [227.01223], @@ -152,7 +152,7 @@ def test_to_gbq_replace(penguins_linear_model, dataset_id): def test_logistic_model_score(penguins_logistic_model, penguins_df_default_index): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -162,8 +162,8 @@ def test_logistic_model_score(penguins_logistic_model, penguins_df_default_index "body_mass_g", ] ] - test_y = df[["sex"]] - result = penguins_logistic_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = penguins_logistic_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "precision": [0.616753], @@ -189,7 +189,7 @@ def test_logistic_model_score_series( penguins_logistic_model, penguins_df_default_index ): df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -199,8 +199,8 @@ def test_logistic_model_score_series( "body_mass_g", ] ] - test_y = df["sex"] - result = penguins_logistic_model.score(test_X, test_y).to_pandas() + y_test = df["sex"] + result = penguins_logistic_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "precision": [0.616753], @@ -244,7 +244,7 @@ def test_logsitic_model_to_gbq_saved_score( f"{dataset_id}.test_penguins_model", replace=True ) df = penguins_df_default_index.dropna() - test_X = df[ + X_test = df[ [ "species", "island", @@ -254,8 +254,8 @@ def test_logsitic_model_to_gbq_saved_score( "body_mass_g", ] ] - test_y = df[["sex"]] - result = saved_model.score(test_X, test_y).to_pandas() + y_test = df[["sex"]] + result = saved_model.score(X_test, y_test).to_pandas() expected = pandas.DataFrame( { "precision": [0.616753], diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 181678ebcb2..7486277487d 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -17,10 +17,25 @@ import numpy as np import pytest +from bigframes.ml import llm + def test_create_text_generator_model(palm2_text_generator_model): # Model creation doesn't return error assert palm2_text_generator_model is not None + assert palm2_text_generator_model._bqml_model is not None + + +def test_create_text_generator_model_defaults(bq_connection): + import bigframes.pandas as bpd + + bpd.reset_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextGenerator() + assert model is not None + assert model._bqml_model is not None # Marked as flaky only because BQML LLM is in preview, the service only has limited capacity, not stable enough. @@ -74,6 +89,19 @@ def test_text_generator_predict_with_params_success( def test_create_embedding_generator_model(palm2_embedding_generator_model): # Model creation doesn't return error assert palm2_embedding_generator_model is not None + assert palm2_embedding_generator_model._bqml_model is not None + + +def test_create_text_embedding_generator_model_defaults(bq_connection): + import bigframes.pandas as bpd + + bpd.reset_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextEmbeddingGenerator() + assert model is not None + assert model._bqml_model is not None @pytest.mark.flaky(retries=2, delay=120) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 01305adb204..85c3cce1d77 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -551,13 +551,69 @@ def test_assign_callable_lambda(scalars_dfs): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) -def test_dropna(scalars_dfs): +@pytest.mark.parametrize( + ("axis", "how", "ignore_index"), + [ + (0, "any", False), + (0, "any", True), + (1, "any", False), + (1, "all", False), + ], +) +def test_df_dropna(scalars_dfs, axis, how, ignore_index): + if pd.__version__.startswith("1."): + pytest.skip("ignore_index parameter not supported in pandas 1.x.") scalars_df, scalars_pandas_df = scalars_dfs - df = scalars_df.dropna() + df = scalars_df.dropna(axis=axis, how=how, ignore_index=ignore_index) bf_result = df.to_pandas() - pd_result = scalars_pandas_df.dropna() + pd_result = scalars_pandas_df.dropna(axis=axis, how=how, ignore_index=ignore_index) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_fillna(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + df = scalars_df[["int64_col", "float64_col"]].fillna(3) + bf_result = df.to_pandas() + pd_result = scalars_pandas_df[["int64_col", "float64_col"]].fillna(3) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_isin_list(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + values = ["Hello, World!", 55555, 2.51, pd.NA, True] + bf_result = ( + scalars_df[["int64_col", "float64_col", "string_col", "bool_col"]] + .isin(values) + .to_pandas() + ) + pd_result = scalars_pandas_df[ + ["int64_col", "float64_col", "string_col", "bool_col"] + ].isin(values) + + pandas.testing.assert_frame_equal(bf_result, pd_result.astype("boolean")) + + +def test_df_isin_dict(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + values = { + "string_col": ["Hello, World!", 55555, 2.51, pd.NA, True], + "int64_col": [5555, 2.51], + "bool_col": [pd.NA], + } + bf_result = ( + scalars_df[["int64_col", "float64_col", "string_col", "bool_col"]] + .isin(values) + .to_pandas() + ) + pd_result = scalars_pandas_df[ + ["int64_col", "float64_col", "string_col", "bool_col"] + ].isin(values) + + pandas.testing.assert_frame_equal(bf_result, pd_result.astype("boolean")) @pytest.mark.parametrize( @@ -1084,50 +1140,43 @@ def test_series_binop_axis_index( @pytest.mark.parametrize( - ("op"), + ("left_labels", "right_labels"), [ - (lambda x, y: x.add(y, axis="index")), - (lambda x, y: x.radd(y, axis="index")), - (lambda x, y: x.sub(y, axis="index")), - (lambda x, y: x.rsub(y, axis="index")), - (lambda x, y: x.mul(y, axis="index")), - (lambda x, y: x.rmul(y, axis="index")), - (lambda x, y: x.truediv(y, axis="index")), - (lambda x, y: x.rtruediv(y, axis="index")), - (lambda x, y: x.floordiv(y, axis="index")), - (lambda x, y: x.floordiv(y, axis="index")), - (lambda x, y: x.gt(y, axis="index")), - (lambda x, y: x.ge(y, axis="index")), - (lambda x, y: x.lt(y, axis="index")), - (lambda x, y: x.le(y, axis="index")), + (["a", "a", "b"], ["c", "c", "d"]), + (["a", "b", "c"], ["c", "a", "b"]), + (["a", "c", "c"], ["c", "a", "c"]), ], ids=[ - "add", - "radd", - "sub", - "rsub", - "mul", - "rmul", - "truediv", - "rtruediv", - "floordiv", - "rfloordiv", - "gt", - "ge", - "lt", - "le", + "no_overlap", + "one_one_match", + "multi_match", ], ) -def test_dataframe_binop_axis_index_throws_not_implemented( - scalars_dfs, - op, +def test_binop_df_df_binary_op( + scalars_df_index, + scalars_df_2_index, + scalars_pandas_df_index, + left_labels, + right_labels, ): - scalars_df, scalars_pandas_df = scalars_dfs - df_columns = ["int64_col", "float64_col"] - other_df_columns = ["int64_too"] - - with pytest.raises(NotImplementedError): - op(scalars_df[df_columns], scalars_df[other_df_columns]).to_pandas() + if pd.__version__.startswith("1."): + pytest.skip("pd.NA vs NaN not handled well in pandas 1.x.") + columns = ["int64_too", "int64_col", "float64_col"] + + bf_df_a = scalars_df_index[columns] + bf_df_a.columns = left_labels + bf_df_b = scalars_df_2_index[columns] + bf_df_b.columns = right_labels + bf_result = (bf_df_a - bf_df_b).to_pandas() + + pd_df_a = scalars_pandas_df_index[columns] + pd_df_a.columns = left_labels + pd_df_b = scalars_pandas_df_index[columns] + pd_df_b.columns = right_labels + pd_result = pd_df_a - pd_df_b + + # Some dtype inconsistency for all-NULL columns + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) # Differnt table will only work for explicit index, since default index orders are arbitrary. diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py index ac1f8c7220d..558dd12e697 100644 --- a/tests/system/small/test_index.py +++ b/tests/system/small/test_index.py @@ -50,3 +50,17 @@ def test_index_getitem_int(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.index[-2] pd_result = scalars_pandas_df_index.index[-2] assert bf_result == pd_result + + +def test_is_monotonic_increasing(scalars_df_index, scalars_pandas_df_index): + assert ( + scalars_df_index.index.is_monotonic_increasing + == scalars_pandas_df_index.index.is_monotonic_increasing + ) + + +def test_is_monotonic_decreasing(scalars_df_index, scalars_pandas_df_index): + assert ( + scalars_df_index.index.is_monotonic_increasing + == scalars_pandas_df_index.index.is_monotonic_increasing + ) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 1baf3e66508..25d1e2ad498 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -634,12 +634,17 @@ def test_column_multi_index_stack(scalars_df_index, scalars_pandas_df_index): pd_df.columns = multi_columns bf_result = bf_df.stack().to_pandas() + # Shifting sort behavior in stack pd_result = pd_df.stack() # Pandas produces NaN, where bq dataframes produces pd.NA - pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + # Column ordering seems to depend on pandas version + pandas.testing.assert_frame_equal( + bf_result.sort_index(axis=1), pd_result.sort_index(axis=1), check_dtype=False + ) +@pytest.mark.skip(reason="Pandas fails in newer versions.") def test_column_multi_index_w_na_stack(scalars_df_index, scalars_pandas_df_index): columns = ["int64_too", "int64_col", "rowindex_2"] level1 = pandas.Index(["b", pandas.NA, pandas.NA]) @@ -656,3 +661,64 @@ def test_column_multi_index_w_na_stack(scalars_df_index, scalars_pandas_df_index # Pandas produces NaN, where bq dataframes produces pd.NA pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("index_names",), + [ + (["rowindex_2", "int64_too"],), + (["int64_too", "rowindex_2"],), + ], +) +def test_is_monotonic_increasing( + scalars_df_index, scalars_pandas_df_index, index_names +): + bf_result = scalars_df_index.set_index(index_names).index + pd_result = scalars_pandas_df_index.set_index(index_names).index + + assert bf_result.is_monotonic_increasing == pd_result.is_monotonic_increasing + + +@pytest.mark.parametrize( + ("indexes",), + [ + ({"A": [1, 2, 3], "B": [1, 2, 3], "C": [1, 2, 3]},), + ({"A": [1, 2, 3], "B": [1, 2, 3], "C": [1, None, 3]},), + ({"A": [1, 2, 2], "B": [1, 2, 1], "C": [1, 2, 3]},), + ({"A": [1, 2, 2], "B": [1, 2, 3], "C": [1, 2, 1]},), + ({"A": [1, 2, 1], "B": [1, 2, 3], "C": [1, 2, 1]},), + ({"A": [3, 2, 1], "B": [3, 2, 1], "C": [2, 2, 1]},), + ], +) +def test_is_monotonic_increasing_extra(indexes): + bf_result = bpd.DataFrame(indexes) + bf_result = bf_result.set_index(["A", "B", "C"]) + pd_result = pandas.DataFrame(indexes) + pd_result = pd_result.set_index(["A", "B", "C"]) + + assert ( + bf_result.index.is_monotonic_increasing + == pd_result.index.is_monotonic_increasing + ) + + +@pytest.mark.parametrize( + ("indexes",), + [ + ({"A": [3, 2, 1], "B": [3, 2, 1], "C": [3, 2, 1]},), + ({"A": [3, 2, 1], "B": [3, 2, 1], "C": [3, None, 1]},), + ({"A": [2, 2, 1], "B": [1, 2, 1], "C": [3, 2, 1]},), + ({"A": [2, 2, 1], "B": [3, 2, 1], "C": [1, 2, 1]},), + ({"A": [1, 2, 1], "B": [3, 2, 1], "C": [1, 2, 1]},), + ], +) +def test_is_monotonic_decreasing_extra(indexes): + bf_result = bpd.DataFrame(indexes) + bf_result = bf_result.set_index(["A", "B", "C"]) + pd_result = pandas.DataFrame(indexes) + pd_result = pd_result.set_index(["A", "B", "C"]) + + assert ( + bf_result.index.is_monotonic_decreasing + == pd_result.index.is_monotonic_decreasing + ) diff --git a/tests/system/small/test_numpy.py b/tests/system/small/test_numpy.py index fff689cabac..5c2a93ec394 100644 --- a/tests/system/small/test_numpy.py +++ b/tests/system/small/test_numpy.py @@ -67,3 +67,69 @@ def test_df_ufuncs(scalars_dfs, opname): pd_result = getattr(np, opname)(scalars_pandas_df[["float64_col", "int64_col"]]) pd.testing.assert_frame_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("opname",), + [ + ("add",), + ("subtract",), + ("multiply",), + ("divide",), + ("power",), + ], +) +def test_series_binary_ufuncs(floats_product_pd, floats_product_bf, opname): + bf_result = getattr(np, opname)( + floats_product_bf.float64_col_x, floats_product_bf.float64_col_y + ).to_pandas() + pd_result = getattr(np, opname)( + floats_product_pd.float64_col_x, floats_product_pd.float64_col_y + ) + pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("opname",), + [ + ("add",), + ("subtract",), + ("multiply",), + ("divide",), + ("power",), + ], +) +def test_df_binary_ufuncs(scalars_dfs, opname): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = getattr(np, opname)( + scalars_df[["float64_col", "int64_col"]], 5.1 + ).to_pandas() + pd_result = getattr(np, opname)( + scalars_pandas_df[["float64_col", "int64_col"]], 5.1 + ) + + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_series_binary_ufuncs_reverse(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + # Could be any non-symmetric binary op + bf_result = np.subtract(5.1, scalars_df["int64_col"]).to_pandas() + pd_result = np.subtract(5.1, scalars_pandas_df["int64_col"]) + + pd.testing.assert_series_equal(bf_result, pd_result) + + +def test_df_binary_ufuncs_reverse(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + # Could be any non-symmetric binary op + bf_result = np.subtract(5.1, scalars_df[["float64_col", "int64_col"]]).to_pandas() + pd_result = np.subtract( + 5.1, + scalars_pandas_df[["float64_col", "int64_col"]], + ) + + pd.testing.assert_frame_equal(bf_result, pd_result) diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index 98bafc6392c..e451d5c3a2c 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -16,6 +16,7 @@ import pytest import bigframes.pandas as bpd +from tests.system.utils import assert_pandas_df_equal_ignore_ordering def test_concat_dataframe(scalars_dfs): @@ -105,3 +106,106 @@ def test_concat_axis_1(scalars_dfs, how): pd_result = pd.concat([pd_part1, pd_part2, pd_part3], join=how, axis=1) pd.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + on = "rowindex_2" + left_columns = ["int64_col", "float64_col", "rowindex_2"] + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_columns] + # Offset the rows somewhat so that outer join can have an effect. + right = scalars_df[right_columns].assign(rowindex_2=scalars_df["rowindex_2"] + 2) + + df = bpd.merge(left, right, merge_how, on, sort=True) + bf_result = df.to_pandas() + + pd_result = pd.merge( + scalars_pandas_df[left_columns], + scalars_pandas_df[right_columns].assign( + rowindex_2=scalars_pandas_df["rowindex_2"] + 2 + ), + merge_how, + on, + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge_left_on_right_on(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + left_columns = ["int64_col", "float64_col", "int64_too"] + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_columns] + right = scalars_df[right_columns] + + df = bpd.merge( + left, right, merge_how, left_on="int64_too", right_on="rowindex_2", sort=True + ) + bf_result = df.to_pandas() + + pd_result = pd.merge( + scalars_pandas_df[left_columns], + scalars_pandas_df[right_columns], + merge_how, + left_on="int64_too", + right_on="rowindex_2", + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("merge_how",), + [ + ("inner",), + ("outer",), + ("left",), + ("right",), + ], +) +def test_merge_series(scalars_dfs, merge_how): + scalars_df, scalars_pandas_df = scalars_dfs + left_column = "int64_too" + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_column] + right = scalars_df[right_columns] + + df = bpd.merge( + left, right, merge_how, left_on="int64_too", right_on="rowindex_2", sort=True + ) + bf_result = df.to_pandas() + + pd_result = pd.merge( + scalars_pandas_df[left_column], + scalars_pandas_df[right_columns], + merge_how, + left_on="int64_too", + right_on="rowindex_2", + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index 9a3d55aed23..6510c4fa271 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -254,49 +254,6 @@ def test_read_gbq_must_comply_with_set_location_non_US( assert df is not None -def test_reset_session_after_bq_session_ended(): - # Use a simple test query to verify that default session works to interact - # with BQ - test_query = "SELECT 1" - - # Confirm that there is a session id in the default session - session = bpd.get_global_session() - assert session._session_id - - # Confirm that session works as usual - df = bpd.read_gbq(test_query) - assert df is not None - - # Abort the session to simulate the auto-expiration - # https://cloud.google.com/bigquery/docs/sessions-terminating#auto-terminate_a_session - abort_session_query = "CALL BQ.ABORT_SESSION()" - query_job = session.bqclient.query(abort_session_query) - query_job.result() # blocks until finished - - # Confirm that session is unusable to run any jobs - with pytest.raises( - google.api_core.exceptions.BadRequest, - match=f"Session {session._session_id} has expired and is no longer available.", - ): - query_job = session.bqclient.query(test_query) - query_job.result() # blocks until finished - - # Confirm that as a result bigframes.pandas interface is unusable - with pytest.raises( - google.api_core.exceptions.BadRequest, - match=f"Session {session._session_id} has expired and is no longer available.", - ): - bpd.read_gbq(test_query) - - # Now try to reset session and verify that it works - bpd.reset_session() - assert bigframes.core.global_session._global_session is None - - # Now verify that use is able to start over - df = bpd.read_gbq(test_query) - assert df is not None - - def test_reset_session_after_credentials_need_reauthentication(monkeypatch): # Use a simple test query to verify that default session works to interact # with BQ diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index c60d270fcaa..77fb81d2c93 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -65,9 +65,7 @@ def bq_cf_connection_location_project_mismatched() -> str: @pytest.fixture(scope="module") def session_with_bq_connection(bq_cf_connection) -> bigframes.Session: - return bigframes.Session( - bigframes.BigQueryOptions(remote_udf_connection=bq_cf_connection) - ) + return bigframes.Session(bigframes.BigQueryOptions(bq_connection=bq_cf_connection)) @pytest.fixture(scope="module") @@ -75,7 +73,7 @@ def session_with_bq_connection_location_specified( bq_cf_connection_location, ) -> bigframes.Session: return bigframes.Session( - bigframes.BigQueryOptions(remote_udf_connection=bq_cf_connection_location) + bigframes.BigQueryOptions(bq_connection=bq_cf_connection_location) ) @@ -84,9 +82,7 @@ def session_with_bq_connection_location_mistached( bq_cf_connection_location_mistached, ) -> bigframes.Session: return bigframes.Session( - bigframes.BigQueryOptions( - remote_udf_connection=bq_cf_connection_location_mistached - ) + bigframes.BigQueryOptions(bq_connection=bq_cf_connection_location_mistached) ) @@ -95,9 +91,7 @@ def session_with_bq_connection_location_project_specified( bq_cf_connection_location_project, ) -> bigframes.Session: return bigframes.Session( - bigframes.BigQueryOptions( - remote_udf_connection=bq_cf_connection_location_project - ) + bigframes.BigQueryOptions(bq_connection=bq_cf_connection_location_project) ) @@ -432,7 +426,7 @@ def test_remote_function_via_session_context_connection_setter( # Creating a session scoped only to this test as we would be setting a # property in it context = bigframes.BigQueryOptions() - context.remote_udf_connection = bq_cf_connection + context.bq_connection = bq_cf_connection session = bigframes.connect(context) # Without an explicit bigquery connection, the one present in Session, @@ -523,6 +517,25 @@ def add_one(x): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.flaky(retries=2, delay=120) +def test_series_map(session_with_bq_connection, scalars_dfs): + def add_one(x): + return x + 1 + + remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = scalars_df.int64_too.map(remote_add_one).to_pandas() + pd_result = scalars_pandas_df.int64_too.map(add_one) + pd_result = pd_result.astype("Int64") # pandas type differences + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + @pytest.mark.flaky(retries=2, delay=120) def test_read_gbq_function_detects_invalid_function(bigquery_client, dataset_id): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 88ad2245c9b..07dc892ddc3 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -186,6 +186,23 @@ def test_fillna(scalars_dfs): ) +@pytest.mark.parametrize( + ("ignore_index",), + ( + (True,), + (False,), + ), +) +def test_series_dropna(scalars_dfs, ignore_index): + if pd.__version__.startswith("1."): + pytest.skip("ignore_index parameter not supported in pandas 1.x.") + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = scalars_df[col_name].dropna(ignore_index=ignore_index).to_pandas() + pd_result = scalars_pandas_df[col_name].dropna(ignore_index=ignore_index) + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + def test_series_agg_single_string(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = scalars_df["int64_col"].agg("sum") @@ -365,6 +382,24 @@ def test_series_int_int_operators_scalar( assert_series_equal_ignoring_order(pd_result, bf_result) +def test_series_pow_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = (scalars_df["int64_col"] ** 2).to_pandas() + pd_result = scalars_pandas_df["int64_col"] ** 2 + + assert_series_equal_ignoring_order(pd_result, bf_result) + + +def test_series_pow_scalar_reverse(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = (0.8 ** scalars_df["int64_col"]).to_pandas() + pd_result = 0.8 ** scalars_pandas_df["int64_col"] + + assert_series_equal_ignoring_order(pd_result, bf_result) + + @pytest.mark.parametrize( ("operator"), [ @@ -2463,3 +2498,57 @@ def test_is_monotonic_decreasing(series_input): assert ( scalars_df.is_monotonic_decreasing == scalars_pandas_df.is_monotonic_decreasing ) + + +def test_map_dict_input(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + local_map = dict() + # construct a local map, incomplete to cover behavior + for s in scalars_pandas_df.string_col[:-3]: + if isinstance(s, str): + local_map[s] = ord(s[0]) + + pd_result = scalars_pandas_df.string_col.map(local_map) + pd_result = pd_result.astype("Int64") # pandas type differences + bf_result = scalars_df.string_col.map(local_map) + + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_map_series_input(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + new_index = scalars_pandas_df.int64_too.drop_duplicates() + pd_map_series = scalars_pandas_df.string_col.iloc[0 : len(new_index)] + pd_map_series.index = new_index + bf_map_series = series.Series( + pd_map_series, session=scalars_df._get_block().expr._session + ) + + pd_result = scalars_pandas_df.int64_too.map(pd_map_series) + bf_result = scalars_df.int64_too.map(bf_map_series) + + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_map_series_input_duplicates_error(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + new_index = scalars_pandas_df.int64_too + pd_map_series = scalars_pandas_df.string_col.iloc[0 : len(new_index)] + pd_map_series.index = new_index + bf_map_series = series.Series( + pd_map_series, session=scalars_df._get_block().expr._session + ) + + with pytest.raises(pd.errors.InvalidIndexError): + scalars_pandas_df.int64_too.map(pd_map_series) + with pytest.raises(pd.errors.InvalidIndexError): + scalars_df.int64_too.map(bf_map_series, verify_integrity=True) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index 43b5663bf75..aeee0583195 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -26,7 +26,7 @@ ("credentials", object(), object()), ("location", "us-east1", "us-central1"), ("project", "my-project", "my-other-project"), - ("remote_udf_connection", "path/to/connection/1", "path/to/connection/2"), + ("bq_connection", "path/to/connection/1", "path/to/connection/2"), ], ) def test_setter_raises_if_session_started(attribute, original_value, new_value): @@ -56,7 +56,7 @@ def test_setter_raises_if_session_started(attribute, original_value, new_value): "credentials", "location", "project", - "remote_udf_connection", + "bq_connection", ] ], ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py deleted file mode 100644 index dcf2d918a5e..00000000000 --- a/tests/unit/conftest.py +++ /dev/null @@ -1,223 +0,0 @@ -# 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 math -from typing import Callable, Optional, Tuple, Union -from unittest import mock - -import google.api_core.exceptions -import google.auth -import google.cloud.bigquery as bigquery -import google.cloud.bigquery.table -import google.oauth2.credentials # type: ignore -import ibis.expr.types as ibis_types -import pandas -import pytest - -import bigframes -import bigframes.core -import bigframes.dataframe - -SCALARS_TABLE_ID = "project.dataset.scalars_table" - - -@pytest.fixture -def scalars_pandas_df_default_index() -> pandas.DataFrame: - # Note: as of 2023-02-07, using nullable dtypes with the ibis pandas - # backend requires running ibis at HEAD. See: - # https://github.com/ibis-project/ibis/pull/5345 - return pandas.DataFrame( - { - "rowindex": pandas.Series( - [ - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - ], - dtype="Int64", - ), - "bool_col": pandas.Series( - [ - True, - None, - False, - True, - None, - False, - True, - None, - False, - True, - ], - dtype="boolean", - ), - "int64_col": pandas.Series( - [ - 1, - 2, - 3, - None, - 0, - -1, - -2, - 2**63 - 1, - -(2**63), - None, - ], - dtype="Int64", - ), - "float64_col": pandas.Series( - [ - None, - 1, - math.pi, - math.e * 1e10, - 0, - float("nan"), - float("inf"), - float("-inf"), - -2.23e-308, - 1.8e308, - ], - dtype="Float64", - ), - "string_col": pandas.Series( - [ - "abc", - "XYZ", - "aBcDeFgHiJkLmNoPqRsTuVwXyZ", - "1_2-3+4=5~6*7/8&9%10#11@12$" "", - None, - "こんにちは", - "你好", - "வணக்கம்", - "שלום", - ], - dtype="string[pyarrow]", - ), - } - ) - - -# We parameterize the fixtures at this point with the real pandas -# dataframes and deferred bigframes dataframes as we have the following -# chain of dependencies: -# -> index/default_index parameterization -# -> pandas dataframe -# -> bqclient mock -# -> session -# -> bigframes dataframe -@pytest.fixture -def scalars_testdata_setup( - scalars_pandas_df_default_index, -) -> Tuple[ - pandas.DataFrame, Callable[[bigframes.Session], bigframes.dataframe.DataFrame] -]: - return ( - scalars_pandas_df_default_index.set_index("rowindex"), - lambda session: session.read_gbq(SCALARS_TABLE_ID, index_col=["rowindex"]), - ) - - -@pytest.fixture(autouse=True) -def mock_bigquery_client(monkeypatch, scalars_testdata_setup) -> bigquery.Client: - scalars_pandas_df, _ = scalars_testdata_setup - mock_client = mock.create_autospec(bigquery.Client) - # Constructor returns the mock itself, so this mock can be treated as the - # constructor or the instance. - mock_client.return_value = mock_client - mock_client.project = "default-project" - most_recent_table = None - - def mock_bigquery_client_get_table( - table_ref: Union[google.cloud.bigquery.table.TableReference, str] - ): - global most_recent_table - - if isinstance(table_ref, google.cloud.bigquery.table.TableReference): - table_name = table_ref.__str__() - else: - table_name = table_ref - - schema = [ - {"mode": "NULLABLE", "name": "rowindex", "type": "INTEGER"}, - { - "mode": "NULLABLE", - "name": "bigframes_ordering_id", - "type": "INTEGER", - }, - ] - - if table_name == SCALARS_TABLE_ID: - schema += [ - {"mode": "NULLABLE", "name": "bool_col", "type": "BOOL"}, - {"mode": "NULLABLE", "name": "int64_col", "type": "INTEGER"}, - {"mode": "NULLABLE", "name": "float64_col", "type": "FLOAT"}, - {"mode": "NULLABLE", "name": "string_col", "type": "STRING"}, - ] - else: - raise google.api_core.exceptions.NotFound("Not Found Table") - - most_recent_table = bigquery.Table(table_name, schema) # type: ignore - return most_recent_table # type: ignore - - def mock_query( - sql: str, - job_config: Optional[bigquery.QueryJobConfig] = None, - location: str = "US", - ) -> bigquery.QueryJob: - global most_recent_table - - def mock_result(max_results=None): - mock_rows = mock.create_autospec(google.cloud.bigquery.table.RowIterator) - mock_rows.total_rows = len(scalars_pandas_df.index) - mock_rows.schema = [ - bigquery.SchemaField(name=name, field_type="INT64") - for name in scalars_pandas_df.columns - ] - # Use scalars_pandas_df instead of ibis_expr.execute() to preserve dtypes. - mock_rows.to_dataframe.return_value = scalars_pandas_df.head(n=max_results) - return mock_rows - - mock_job = mock.create_autospec(bigquery.QueryJob) - mock_job.result = mock_result - return mock_job - - mock_client.get_table = mock_bigquery_client_get_table - mock_client.query.side_effect = mock_query - monkeypatch.setattr(bigquery, "Client", mock_client) - mock_client.reset_mock() - return mock_client - - -@pytest.fixture -def session() -> bigframes.Session: - return bigframes.Session( - context=bigframes.BigQueryOptions( - credentials=mock.create_autospec(google.oauth2.credentials.Credentials), - project="unit-test-project", - ) - ) - - -@pytest.fixture -def scalars_ibis_table(session) -> ibis_types.Table: - return session.ibis_client.table(SCALARS_TABLE_ID) diff --git a/tests/unit/core/test_utils.py b/tests/unit/core/test_bf_utils.py similarity index 100% rename from tests/unit/core/test_utils.py rename to tests/unit/core/test_bf_utils.py diff --git a/tests/unit/resources.py b/tests/unit/resources.py new file mode 100644 index 00000000000..c8ed6e86ed6 --- /dev/null +++ b/tests/unit/resources.py @@ -0,0 +1,73 @@ +# 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. + +from typing import Dict, List, Optional +import unittest.mock as mock + +import google.auth.credentials +import google.cloud.bigquery +import ibis +import pandas + +import bigframes +import bigframes.core as core + +"""Utilities for creating test resources.""" + + +def create_bigquery_session( + bqclient: Optional[google.cloud.bigquery.Client] = None, session_id: str = "abcxyz" +) -> bigframes.Session: + credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + if bqclient is None: + bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) + bqclient.project = "test-project" + + clients_provider = mock.create_autospec(bigframes.session.ClientsProvider) + type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient) + clients_provider._credentials = credentials + + bqoptions = bigframes.BigQueryOptions( + credentials=credentials, location="test-region" + ) + session = bigframes.Session(context=bqoptions, clients_provider=clients_provider) + session._session_id = session_id + return session + + +def create_pandas_session(tables: Dict[str, pandas.DataFrame]) -> bigframes.Session: + # TODO(tswast): Refactor to make helper available for all tests. Consider + # providing a proper "local Session" for use by downstream developers. + session = mock.create_autospec(bigframes.Session, instance=True) + ibis_client = ibis.pandas.connect(tables) + type(session).ibis_client = mock.PropertyMock(return_value=ibis_client) + return session + + +def create_arrayvalue( + df: pandas.DataFrame, total_ordering_columns: List[str] +) -> bigframes.core.ArrayValue: + session = create_pandas_session({"test_table": df}) + ibis_table = session.ibis_client.table("test_table") + columns = tuple(ibis_table[key] for key in ibis_table.columns) + ordering = core.ExpressionOrdering( + [core.OrderingColumnReference(column) for column in total_ordering_columns], + total_ordering_columns=frozenset(total_ordering_columns), + ) + return core.ArrayValue( + session=session, table=ibis_table, columns=columns, ordering=ordering + ) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 123dae79395..e01638e22e1 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -13,39 +13,55 @@ # limitations under the License. import ibis -from ibis.expr.types import Table +import pandas -from bigframes import core +import bigframes.core as core -ORDERING = core.ExpressionOrdering( - [ - core.OrderingColumnReference("int64_col"), - core.OrderingColumnReference("string_col"), - ], - total_ordering_columns=frozenset(["int64_col", "string_col"]), -) +from . import resources -def test_constructor_from_ibis_table_adds_all_columns( - session, scalars_ibis_table: Table -): - columns = tuple(scalars_ibis_table[key] for key in scalars_ibis_table.columns) +def test_arrayvalue_constructor_from_ibis_table_adds_all_columns(): + session = resources.create_pandas_session( + { + "test_table": pandas.DataFrame( + { + "col1": [1, 2, 3], + "not_included": [True, False, True], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ) + } + ) + ibis_table = session.ibis_client.table("test_table") + columns = (ibis_table["col1"], ibis_table["col2"], ibis_table["col3"]) + ordering = core.ExpressionOrdering( + [core.OrderingColumnReference("col1")], + total_ordering_columns=frozenset(["col1"]), + ) actual = core.ArrayValue( - session=session, table=scalars_ibis_table, columns=columns, ordering=ORDERING + session=session, table=ibis_table, columns=columns, ordering=ordering ) - assert actual._table is scalars_ibis_table - assert len(actual._columns) == len(scalars_ibis_table.columns) + assert actual.table is ibis_table + assert len(actual.columns) == 3 -def test_to_ibis_expr_with_projection(session, scalars_ibis_table: Table): - columns = tuple(scalars_ibis_table[key] for key in scalars_ibis_table.columns) - expr = core.ArrayValue( - session=session, table=scalars_ibis_table, columns=columns, ordering=ORDERING - ).projection( +def test_arrayvalue_to_ibis_expr_with_projection(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.projection( [ - scalars_ibis_table["int64_col"], + (value.table["col1"] + ibis.literal(-1)).name("int64_col"), ibis.literal(123456789).name("literals"), - scalars_ibis_table["string_col"], + value.table["col2"].name("string_col"), ] ) actual = expr.to_ibis_expr() diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index a178a454388..2325fc96a0a 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -17,12 +17,17 @@ import sys import unittest.mock as mock +import google.api_core.exceptions +import google.cloud.bigquery import pandas as pd import pytest +import bigframes.core.global_session import bigframes.pandas as bpd import bigframes.session +from . import resources + leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE) @@ -109,3 +114,37 @@ def test_pandas_attribute(): assert bpd.Int64Dtype is pd.Int64Dtype assert bpd.StringDtype is pd.StringDtype assert bpd.ArrowDtype is pd.ArrowDtype + + +def test_reset_session_after_bq_session_ended(monkeypatch): + bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) + bqclient.project = "test-project" + session = resources.create_bigquery_session( + bqclient=bqclient, session_id="JUST_A_TEST" + ) + + # Simulate that the session has already expired. + # Note: this needs to be done after the Session is constructed, as the + # initializer sends a query to start the BigQuery Session. + query_job = mock.create_autospec(google.cloud.bigquery.QueryJob, instance=True) + query_job.result.side_effect = google.api_core.exceptions.BadRequest( + "Session JUST_A_TEST has expired and is no longer available." + ) + bqclient.query.return_value = query_job + + # Simulate that the session has already started. + monkeypatch.setattr(bigframes.core.global_session, "_global_session", session) + bpd.options.bigquery._session_started = True + + # Confirm that as a result bigframes.pandas interface is unusable + with pytest.raises( + google.api_core.exceptions.BadRequest, + match="Session JUST_A_TEST has expired and is no longer available.", + ): + bpd.read_gbq("SELECT 1") + + # Even though the query to stop the session raises an exception, we should + # still be able to reset it without raising an error to the user. + bpd.reset_session() + assert "CALL BQ.ABORT_SESSION('JUST_A_TEST')" in bqclient.query.call_args.args[0] + assert bigframes.core.global_session._global_session is None diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index ab573c4c11f..e39a316e5b6 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -20,9 +20,13 @@ import bigframes +from . import resources + @pytest.mark.parametrize("missing_parts_table_id", [(""), ("table")]) -def test_read_gbq_missing_parts(session, missing_parts_table_id): +def test_read_gbq_missing_parts(missing_parts_table_id): + session = resources.create_bigquery_session() + with pytest.raises(ValueError): session.read_gbq(missing_parts_table_id) @@ -31,7 +35,14 @@ def test_read_gbq_missing_parts(session, missing_parts_table_id): "not_found_table_id", [("unknown.dataset.table"), ("project.unknown.table"), ("project.dataset.unknown")], ) -def test_read_gdb_not_found_tables(session, not_found_table_id): +def test_read_gdb_not_found_tables(not_found_table_id): + bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) + bqclient.project = "test-project" + bqclient.get_table.side_effect = google.api_core.exceptions.NotFound( + "table not found" + ) + session = resources.create_bigquery_session(bqclient=bqclient) + with pytest.raises(google.api_core.exceptions.NotFound): session.read_gbq(not_found_table_id) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 5a812dae7e7..8c81b23b6ca 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -380,11 +380,44 @@ def dropna( ) -> DataFrame: """Remove missing values. + Args: + axis ({0 or 'index', 1 or 'columns'}, default 'columns'): + Determine if rows or columns which contain missing values are + removed. + + * 0, or 'index' : Drop rows which contain missing values. + * 1, or 'columns' : Drop columns which contain missing value. + how ({'any', 'all'}, default 'any'): + Determine if row or column is removed from DataFrame, when we have + at least one NA or all NA. + + * 'any' : If any NA values are present, drop that row or column. + * 'all' : If all values are NA, drop that row or column. + ignore_index (bool, default ``False``): + If ``True``, the resulting axis will be labeled 0, 1, …, n - 1. + + Returns: bigframes.dataframe.DataFrame: DataFrame with NA entries dropped from it. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def isin(self, values): + """ + Whether each element in the DataFrame is contained in values. + + Args: + values (iterable, or dict): + The result will only be true at a location if all the + labels match. If `values` is a dict, the keys must be + the column names, which must match. + + Returns: + DataFrame: DataFrame of booleans showing whether each element + in the DataFrame is contained in values. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Sorting @@ -808,6 +841,54 @@ def rmod(self, other, axis: str | int = "columns") -> DataFrame: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def pow(self, other, axis: str | int = "columns") -> DataFrame: + """Get Exponential power of dataframe and other, element-wise (binary operator `pow`). + + Equivalent to ``dataframe ** other``, but with support to substitute a fill_value + for missing data in one of the inputs. With reverse version, `rpow`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + .. note:: + Mismatched indices will be unioned together. + + Args: + other (float, int, or Series): + Any single or multiple element data structure, or list-like object. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. + + Returns: + DataFrame: DataFrame result of the arithmetic operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def rpow(self, other, axis: str | int = "columns") -> DataFrame: + """Get Exponential power of dataframe and other, element-wise (binary operator `rpow`). + + Equivalent to ``other ** dataframe``, but with support to substitute a fill_value + for missing data in one of the inputs. With reverse version, `pow`. + + Among flexible wrappers (`add`, `sub`, `mul`, `div`, `mod`, `pow`) to + arithmetic operators: `+`, `-`, `*`, `/`, `//`, `%`, `**`. + + .. note:: + Mismatched indices will be unioned together. + + Args: + other (float, int, or Series): + Any single or multiple element data structure, or list-like object. + axis ({0 or 'index', 1 or 'columns'}): + Whether to compare by the index (0 or 'index') or columns. + (1 or 'columns'). For Series input, axis to match Series index on. + + Returns: + DataFrame: DataFrame result of the arithmetic operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Data reshaping @@ -1338,3 +1419,20 @@ def value_counts( Series: Series containing counts of unique rows in the DataFrame """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def fillna(self, value): + """ + Fill NA/NaN values using the specified method. + + Args: + value (scalar, Series): + Value to use to fill holes (e.g. 0), alternately a + Series of values specifying which value to use for + each index (for a Series) or column (for a DataFrame). Values not + in the Series will not be filled. This value cannot + be a list. + + Returns: + DataFrame: Object with missing values filled + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/reshape/merge.py b/third_party/bigframes_vendored/pandas/core/reshape/merge.py new file mode 100644 index 00000000000..ee02d698da3 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/reshape/merge.py @@ -0,0 +1,78 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/merge.py +""" +SQL-style merge routines +""" +from __future__ import annotations + + +def merge( + left, + right, + how="inner", + on=None, + *, + left_on=None, + right_on=None, + sort=False, + suffixes=("_x", "_y"), +): + + """ + Merge DataFrame objects with a database-style join. + + The join is done on columns or indexes. If joining columns on + columns, the DataFrame indexes *will be ignored*. Otherwise if joining indexes + on indexes or indexes on a column or columns, the index will be passed on. + When performing a cross merge, no column specifications to merge on are + allowed. + + .. note:: + A named Series object is treated as a DataFrame with a single named column. + + .. warning:: + If both key columns contain rows where the key is a null value, those + rows will be matched against each other. This is different from usual SQL + join behaviour and can lead to unexpected results. + + Args: + left: + The primary object to be merged. + right: + Object to merge with. + how: + ``{'left', 'right', 'outer', 'inner'}, default 'inner'`` + Type of merge to be performed. + ``left``: use only keys from left frame, similar to a SQL left outer join; + preserve key order. + ``right``: use only keys from right frame, similar to a SQL right outer join; + preserve key order. + ``outer``: use union of keys from both frames, similar to a SQL full outer + join; sort keys lexicographically. + ``inner``: use intersection of keys from both frames, similar to a SQL inner + join; preserve the order of the left keys. + + on: + Column join on. It must be found in both DataFrames. Either on or left_on + right_on + must be passed in. + left_on: + Column join on in the left DataFrame. Either on or left_on + right_on + must be passed in. + right_on: + Column join on in the right DataFrame. Either on or left_on + right_on + must be passed in. + sort: + Default False. Sort the join keys lexicographically in the + result DataFrame. If False, the order of the join keys depends + on the join type (how keyword). + suffixes: + Default ``("_x", "_y")``. A length-2 sequence where each + element is optionally a string indicating the suffix to add to + overlapping column names in `left` and `right` respectively. + Pass a value of `None` instead of a string to indicate that the + column name from `left` or `right` should be left as-is, with + no suffix. At least one of the values must not be None. + + Returns: + bigframes.dataframe.DataFrame: A DataFrame of the two merged objects. + """ + raise NotImplementedError("abstract method") diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 8d505c1eadb..76fb46a7003 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -207,10 +207,16 @@ def to_dict(self, into: type[dict] = dict) -> Mapping: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def to_frame(self) -> DataFrame: + def to_frame(self, name=None) -> DataFrame: """ Convert Series to DataFrame. + The column in the new dataframe will be named name (the keyword parameter) + if the name parameter is provided and not None. + + Args: + name (Hashable, default None) + Returns: bigframes.dataframe.DataFrame: DataFrame representation of Series. """ @@ -830,6 +836,23 @@ def fillna( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def dropna(self, *, axis=0, inplace: bool = False, how=None) -> Series: + """ + Return a new Series with missing values removed. + + Args: + axis (0 or 'index'): + Unused. Parameter needed for compatibility with DataFrame. + inplace (bool, default False): + Unsupported, do not set. + how (str, optional): + Not in use. Kept for compatibility. + + Returns: + Series: Series with NA entries dropped from it. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def between( self, left, @@ -1185,9 +1208,39 @@ def mod(self, other) -> Series: raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def rmod(self, other) -> Series: - """Get modulo of Series and other, element-wise (binary operator `rmod`). + """Return modulo of Series and other, element-wise (binary operator mod). + + Equivalent to ``series % other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other (Series, or scalar value): + + Returns: + bigframes.series.Series: The result of the operation. + + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - Equivalent to ``other % series``, but with support to substitute a fill_value for + def pow(self, other) -> Series: + """Return Exponential power of series and other, element-wise (binary operator `pow`). + + Equivalent to ``series ** other``, but with support to substitute a fill_value for + missing data in either one of the inputs. + + Args: + other (Series, or scalar value): + + Returns: + bigframes.series.Series: The result of the operation. + + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def rpow(self, other) -> Series: + """Return Exponential power of series and other, element-wise (binary operator `rpow`). + + Equivalent to ``other ** series``, but with support to substitute a fill_value for missing data in either one of the inputs. Args: @@ -1640,3 +1693,46 @@ def is_monotonic_decreasing(self) -> bool: bool """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def map( + self, + arg, + na_action=None, + *, + verify_integrity=False, + ) -> Series: + """ + Map values of Series according to an input mapping or function. + + Used for substituting each value in a Series with another value, + that may be derived from a remote function, ``dict``, or a :class:`Series`. + + If arg is a remote function, the overhead for remote functions + applies. If mapping with a dict, fully deferred computation is possible. + If mapping with a Series, fully deferred computation is only possible if + verify_integrity=False. + + .. note:: + Bigframes does not yet support ``dict`` subclasses that define + ``__missing__`` (i.e. provide a method for default values). These + are treated the same as ``dict``. + + Args: + arg (function, Mapping, Series): + remote function, collections.abc.Mapping subclass or Series + Mapping correspondence. + na_action: (str, default None) + Only None is currently supported, indicating that arg may + map values to scalars. values won't be ignored. + Passing 'ignore' will raise NotImplementedException. + verify_integrity: (bool, default False) + Only applies when arg is a Series. If True, throw if the Series + index contains duplicate entries (this matches pandas behavior). + If False, skip the expensive computation, and any duplicate + index entries will produce duplicate rows in the result for each + index entry. + + Returns: + Series: Same index as caller. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 106c75c79e6ba33f613075053c84ac7fe7d22c2c Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Tue, 5 Sep 2023 06:20:31 -0500 Subject: [PATCH 12/80] chore(main): release 0.3.0 (#13) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 67 ++++++++++++++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ab2b05d571..7770534cad4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,73 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.3.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.2.0...v0.3.0) (2023-09-02) + + +### Features + +* Add `bigframes.get_global_session()` and `bigframes.reset_session()` aliases ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Add `bigframes.pandas.read_pickle` function ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Add `components_`, `explained_variance_`, and `explained_variance_ratio_` properties to `bigframes.ml.decomposition.PCA` ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Add `fit_transform` to `bigquery.ml` transformers ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Add `Series.dropna` and `DataFrame.fillna` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Add `Series.str` methods `isalpha`, `isdigit`, `isdecimal`, `isalnum`, `isspace`, `islower`, `isupper`, `zfill`, `center` ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Support `bigframes.pandas.merge()` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support `DataFrame.isin` with list and dict inputs ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support `DataFrame.pivot` ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Support `DataFrame.stack` ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `DataFrame`-`DataFrame` binary operations ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support `df[my_column] = [a python list]` ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `Index.is_monotonic` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support `np.arcsin`, `np.arccos`, `np.arctan`, `np.sinh`, `np.cosh`, `np.tanh`, `np.arcsinh`, `np.arccosh`, `np.arctanh`, `np.exp` with Series argument ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `np.sin`, `np.cos`, `np.tan`, `np.log`, `np.log10`, `np.sqrt`, `np.abs` with Series argument ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `pow()` and power operator in `DataFrame` and `Series` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support `read_json` with `engine=bigquery` for newline-delimited JSON files ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `Series.corr` ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Support `Series.map` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support for `np.add`, `np.subtract`, `np.multiply`, `np.divide`, `np.power` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Support MultiIndex for DataFrame columns ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Use `pandas.Index` for column labels ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Use default session and connection in `ml.llm` and `ml.imported` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) + + +### Bug Fixes + +* Add error message to `set_index` ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Align column names with pandas in `DataFrame.agg` results ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Allow (but still not recommended) `ORDER BY` in `read_gbq` input when an `index_col` is defined ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Check for IAM role on the BigQuery connection when initializing a `remote_function` ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Check that types are specified in `read_gbq_function` ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Don't use query cache for Session construction ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Include survey link in abstract `NotImplementedError` exception messages ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Label temp table creation jobs with `source=bigquery-dataframes-temp` label ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Make `X_train` argument names consistent across methods ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) +* Raise AttributeError for unimplemented pandas methods ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Raise exception for invalid function in `read_gbq_function` ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Support spaces in column names in `DataFrame` initializater ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) + + +### Performance Improvements + +* Add local cache for `__repr_*__` methods ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Lazily instantiate client library objects ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Use `row_number()` filter for `head` / `tail` ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) + + +### Documentation + +* Add ML section under Overview ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Add release status to table of contents ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Add samples and best practices to `read_gbq` docs ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Correct the return types of Dataframe and Series ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Create subfolders for notebooks ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Fix link to GitHub ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Highlight bigframes is open-source ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Sample ML Drug Name Generation notebook ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Set `options.bigquery.project` in sample code ([89b9503](https://github.com/googleapis/python-bigquery-dataframes/commit/89b95033d6b449bfc21249057d7c024d096c80d0)) +* Transform remote function user guide into sample code ([a32b747](https://github.com/googleapis/python-bigquery-dataframes/commit/a32b74751785c8e8aec40ce01df639dd7c4fbb77)) +* Update remote function notebook with read_gbq_function usage ([8fab755](https://github.com/googleapis/python-bigquery-dataframes/commit/8fab75576757230bca5c7df10994837ac406300f)) + ## 0.2.0 (2023-08-17) ### Features diff --git a/bigframes/version.py b/bigframes/version.py index d386742d59f..4cc4639705d 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.2.0" +__version__ = "0.3.0" From 99722e7642fee629e1fc0558d7c2d6c71915f1cf Mon Sep 17 00:00:00 2001 From: Jeff Ching Date: Tue, 5 Sep 2023 12:29:40 -0700 Subject: [PATCH 13/80] build: enable multi-scm for release job (#16) --- .github/release-trigger.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/release-trigger.yml b/.github/release-trigger.yml index d4ca94189e1..b0a6cadefca 100644 --- a/.github/release-trigger.yml +++ b/.github/release-trigger.yml @@ -1 +1,2 @@ enabled: true +multiScmName: bigframes From 3dd25b379ed832ea062e188f483d2789830de67b Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 5 Sep 2023 15:17:39 -0500 Subject: [PATCH 14/80] fix(release): use correct directory name for release build config (#17) multi_scm name is "bigframes" --- .kokoro/release/common.cfg | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.kokoro/release/common.cfg b/.kokoro/release/common.cfg index 7ffa79c7a1d..021198fb269 100644 --- a/.kokoro/release/common.cfg +++ b/.kokoro/release/common.cfg @@ -11,7 +11,7 @@ action { 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" +build_file: "bigframes/.kokoro/trampoline_v2.sh" # Configure the docker image for kokoro-trampoline. env_vars: { @@ -20,7 +20,7 @@ env_vars: { } env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: "github/python-bigquery-dataframes/.kokoro/release.sh" + value: "github/bigframes/.kokoro/release.sh" } # Fetch PyPI password @@ -43,7 +43,7 @@ env_vars: { # what we published, which we can use to generate SBOMs and attestations. action { define_artifacts { - regex: "github/python-bigquery-dataframes/**/*.tar.gz" - strip_prefix: "github/python-bigquery-dataframes" + regex: "github/bigframes/**/*.tar.gz" + strip_prefix: "github/bigframes" } -} \ No newline at end of file +} From 9e581edbc7352b8a9567c1a30864b54e102e1f7c Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Tue, 5 Sep 2023 15:21:57 -0500 Subject: [PATCH 15/80] chore(main): release 0.3.1 (#18) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 7 +++++++ bigframes/version.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7770534cad4..1bcfeacfd0e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.3.1](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.0...v0.3.1) (2023-09-05) + + +### Bug Fixes + +* **release:** Use correct directory name for release build config ([#17](https://github.com/googleapis/python-bigquery-dataframes/issues/17)) ([3dd25b3](https://github.com/googleapis/python-bigquery-dataframes/commit/3dd25b379ed832ea062e188f483d2789830de67b)) + ## [0.3.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.2.0...v0.3.0) (2023-09-02) diff --git a/bigframes/version.py b/bigframes/version.py index 4cc4639705d..90bf336b9d6 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.3.0" +__version__ = "0.3.1" From 995161068b118a639903878acfde3202087c25f8 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 6 Sep 2023 08:56:43 -0500 Subject: [PATCH 16/80] fix: make release.sh script for PyPI upload executable (#20) Change-Id: I45939017295e5c6c4942087afecb5659b423d4e9 --- .kokoro/release.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) mode change 100644 => 100755 .kokoro/release.sh diff --git a/.kokoro/release.sh b/.kokoro/release.sh old mode 100644 new mode 100755 index 697b04009db..2bce143e1c9 --- a/.kokoro/release.sh +++ b/.kokoro/release.sh @@ -16,7 +16,7 @@ set -eo pipefail # Start the releasetool reporter -python3 -m pip install --require-hashes -r github/python-bigquery-dataframes/.kokoro/requirements.txt +python3 -m pip install --require-hashes -r github/bigframes/.kokoro/requirements.txt python3 -m releasetool publish-reporter-script > /tmp/publisher-script; source /tmp/publisher-script # Disable buffering, so that the logs stream through. @@ -24,6 +24,6 @@ export PYTHONUNBUFFERED=1 # Move into the package, build the distribution and upload. TWINE_PASSWORD=$(cat "${KOKORO_KEYSTORE_DIR}/73713_google-cloud-pypi-token-keystore-1") -cd github/python-bigquery-dataframes +cd github/bigframes python3 setup.py sdist bdist_wheel -twine upload --username __token__ --password "${TWINE_PASSWORD}" dist/* \ No newline at end of file +twine upload --username __token__ --password "${TWINE_PASSWORD}" dist/* From e76b5e7a0738c8d408e336ea21ad0c86e5b4511a Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 6 Sep 2023 10:04:03 -0500 Subject: [PATCH 17/80] chore(main): release 0.3.2 (#21) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 7 +++++++ bigframes/version.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bcfeacfd0e..aa736877f4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.3.2](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.1...v0.3.2) (2023-09-06) + + +### Bug Fixes + +* Make release.sh script for PyPI upload executable ([#20](https://github.com/googleapis/python-bigquery-dataframes/issues/20)) ([9951610](https://github.com/googleapis/python-bigquery-dataframes/commit/995161068b118a639903878acfde3202087c25f8)) + ## [0.3.1](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.0...v0.3.1) (2023-09-05) diff --git a/bigframes/version.py b/bigframes/version.py index 90bf336b9d6..ef4f01f6294 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.3.1" +__version__ = "0.3.2" From ffad464a99219a2139b1fd456703041a0e185750 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 7 Sep 2023 13:07:44 -0500 Subject: [PATCH 18/80] chore: enable owlbot (#23) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: enable owlbot Change-Id: Idb2a03ab01d414f136a846d18f90329a9a50e4aa * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * revert some owlbot changes Change-Id: Iabe654bf693d82f5f1afdef7e9663577ded96025 * only exclude common.cfg Change-Id: I1199737edca971b648ec2cfddb82b4f378b3f97d * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * use correct path in build.sh Change-Id: I407dc0500f52c52ccd90e108ba821fd518a7fe1d * decrease coverage and use repo name in 'nightly' build Change-Id: I7d450364dd234c640d905c9ac2c2a96e7df08d02 * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * restore pre-commit checks Change-Id: Iaf059d5b197d9cc9e87e3d17adf6c7c48ca41cb8 --------- Co-authored-by: Owl Bot --- .coveragerc | 39 ++ .flake8 | 4 +- .github/.OwlBot.lock.yaml | 17 + .github/.OwlBot.yaml | 18 + .github/CODEOWNERS | 1 + .github/header-checker-lint.yml | 14 +- .github/sync-repo-settings.yaml | 2 +- .github/workflows/lint.yml | 2 +- .github/workflows/unittest.yml | 4 +- .gitignore | 7 - .kokoro/build.sh | 6 +- .kokoro/continuous/continuous.cfg | 2 +- .kokoro/continuous/prerelease-deps.cfg | 7 + .kokoro/docker/docs/Dockerfile | 2 +- .kokoro/docs/common.cfg | 14 +- .kokoro/docs/docs-presubmit.cfg | 9 +- .kokoro/docs/docs.cfg | 2 +- .kokoro/populate-secrets.sh | 43 ++ .kokoro/presubmit/common.cfg | 2 +- .kokoro/presubmit/prerelease-deps.cfg | 7 + .kokoro/presubmit/presubmit.cfg | 2 +- .kokoro/publish-docs.sh | 19 + .kokoro/release-nightly.sh | 49 +-- .kokoro/release.sh | 4 +- .kokoro/release/common.cfg | 8 +- .kokoro/requirements.in | 3 - .kokoro/requirements.txt | 384 +++++++----------- .kokoro/samples/lint/common.cfg | 34 ++ .kokoro/samples/lint/continuous.cfg | 6 + .kokoro/samples/lint/periodic.cfg | 6 + .kokoro/samples/lint/presubmit.cfg | 6 + .kokoro/samples/python3.10/common.cfg | 40 ++ .kokoro/samples/python3.10/continuous.cfg | 6 + .kokoro/samples/python3.10/periodic-head.cfg | 11 + .kokoro/samples/python3.10/periodic.cfg | 6 + .kokoro/samples/python3.10/presubmit.cfg | 6 + .kokoro/samples/python3.11/common.cfg | 40 ++ .kokoro/samples/python3.11/continuous.cfg | 6 + .kokoro/samples/python3.11/periodic-head.cfg | 11 + .kokoro/samples/python3.11/periodic.cfg | 6 + .kokoro/samples/python3.11/presubmit.cfg | 6 + .kokoro/samples/python3.7/common.cfg | 40 ++ .kokoro/samples/python3.7/continuous.cfg | 6 + .kokoro/samples/python3.7/periodic-head.cfg | 11 + .kokoro/samples/python3.7/periodic.cfg | 6 + .kokoro/samples/python3.7/presubmit.cfg | 6 + .kokoro/samples/python3.8/common.cfg | 40 ++ .kokoro/samples/python3.8/continuous.cfg | 6 + .kokoro/samples/python3.8/periodic-head.cfg | 11 + .kokoro/samples/python3.8/periodic.cfg | 6 + .kokoro/samples/python3.8/presubmit.cfg | 6 + .kokoro/samples/python3.9/common.cfg | 40 ++ .kokoro/samples/python3.9/continuous.cfg | 6 + .kokoro/samples/python3.9/periodic-head.cfg | 11 + .kokoro/samples/python3.9/periodic.cfg | 6 + .kokoro/samples/python3.9/presubmit.cfg | 6 + .kokoro/test-samples-against-head.sh | 26 ++ .kokoro/test-samples-impl.sh | 102 +++++ .kokoro/test-samples.sh | 44 ++ .kokoro/trampoline.sh | 28 ++ .kokoro/trampoline_v2.sh | 0 .pre-commit-config.yaml | 8 +- .trampolinerc | 3 + CODE_OF_CONDUCT.md | 3 +- CONTRIBUTING.rst | 277 +++++++++++++ MANIFEST.in | 26 ++ docs/_static/custom.css | 5 + docs/_templates/layout.html | 5 + docs/conf.py | 41 +- noxfile.py | 3 +- owlbot.py | 54 ++- renovate.json | 12 + scripts/decrypt-secrets.sh | 46 +++ scripts/readme-gen/readme_gen.py | 69 ++++ scripts/readme-gen/templates/README.tmpl.rst | 87 ++++ scripts/readme-gen/templates/auth.tmpl.rst | 9 + .../templates/auth_api_key.tmpl.rst | 14 + .../templates/install_deps.tmpl.rst | 29 ++ .../templates/install_portaudio.tmpl.rst | 35 ++ scripts/upload_to_google_drive.py | 5 +- setup.cfg | 16 +- testing/.gitignore | 3 + 82 files changed, 1652 insertions(+), 366 deletions(-) create mode 100644 .coveragerc create mode 100644 .github/.OwlBot.lock.yaml create mode 100644 .github/.OwlBot.yaml create mode 100644 .kokoro/continuous/prerelease-deps.cfg create mode 100755 .kokoro/populate-secrets.sh create mode 100644 .kokoro/presubmit/prerelease-deps.cfg create mode 100644 .kokoro/samples/lint/common.cfg create mode 100644 .kokoro/samples/lint/continuous.cfg create mode 100644 .kokoro/samples/lint/periodic.cfg create mode 100644 .kokoro/samples/lint/presubmit.cfg create mode 100644 .kokoro/samples/python3.10/common.cfg create mode 100644 .kokoro/samples/python3.10/continuous.cfg create mode 100644 .kokoro/samples/python3.10/periodic-head.cfg create mode 100644 .kokoro/samples/python3.10/periodic.cfg create mode 100644 .kokoro/samples/python3.10/presubmit.cfg create mode 100644 .kokoro/samples/python3.11/common.cfg create mode 100644 .kokoro/samples/python3.11/continuous.cfg create mode 100644 .kokoro/samples/python3.11/periodic-head.cfg create mode 100644 .kokoro/samples/python3.11/periodic.cfg create mode 100644 .kokoro/samples/python3.11/presubmit.cfg create mode 100644 .kokoro/samples/python3.7/common.cfg create mode 100644 .kokoro/samples/python3.7/continuous.cfg create mode 100644 .kokoro/samples/python3.7/periodic-head.cfg create mode 100644 .kokoro/samples/python3.7/periodic.cfg create mode 100644 .kokoro/samples/python3.7/presubmit.cfg create mode 100644 .kokoro/samples/python3.8/common.cfg create mode 100644 .kokoro/samples/python3.8/continuous.cfg create mode 100644 .kokoro/samples/python3.8/periodic-head.cfg create mode 100644 .kokoro/samples/python3.8/periodic.cfg create mode 100644 .kokoro/samples/python3.8/presubmit.cfg create mode 100644 .kokoro/samples/python3.9/common.cfg create mode 100644 .kokoro/samples/python3.9/continuous.cfg create mode 100644 .kokoro/samples/python3.9/periodic-head.cfg create mode 100644 .kokoro/samples/python3.9/periodic.cfg create mode 100644 .kokoro/samples/python3.9/presubmit.cfg create mode 100755 .kokoro/test-samples-against-head.sh create mode 100755 .kokoro/test-samples-impl.sh create mode 100755 .kokoro/test-samples.sh create mode 100755 .kokoro/trampoline.sh mode change 100644 => 100755 .kokoro/trampoline_v2.sh create mode 100644 CONTRIBUTING.rst create mode 100644 MANIFEST.in create mode 100644 renovate.json create mode 100755 scripts/decrypt-secrets.sh create mode 100644 scripts/readme-gen/readme_gen.py create mode 100644 scripts/readme-gen/templates/README.tmpl.rst create mode 100644 scripts/readme-gen/templates/auth.tmpl.rst create mode 100644 scripts/readme-gen/templates/auth_api_key.tmpl.rst create mode 100644 scripts/readme-gen/templates/install_deps.tmpl.rst create mode 100644 scripts/readme-gen/templates/install_portaudio.tmpl.rst create mode 100644 testing/.gitignore diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 00000000000..dffe22a6f46 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +# +# 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. + +# Generated by synthtool. DO NOT EDIT! +[run] +branch = True +omit = + google/__init__.py + google/cloud/__init__.py + +[report] +fail_under = 35 +show_missing = True +exclude_lines = + # Re-enable the standard pragma + pragma: NO COVER + # Ignore debug-only repr + def __repr__ + # Ignore abstract methods + raise NotImplementedError +omit = + */gapic/*.py + */proto/*.py + */core/*.py + */site-packages/*.py + google/cloud/__init__.py diff --git a/.flake8 b/.flake8 index 7826eceab6a..87f6e408c47 100644 --- a/.flake8 +++ b/.flake8 @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2020 Google LLC +# 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. @@ -18,8 +18,6 @@ [flake8] ignore = E203, E231, E266, E501, W503 exclude = - third_party/** - # Exclude generated code. **/proto/** **/gapic/** diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml new file mode 100644 index 00000000000..a3da1b0d4cd --- /dev/null +++ b/.github/.OwlBot.lock.yaml @@ -0,0 +1,17 @@ +# 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. +docker: + image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest + digest: sha256:3e3800bb100af5d7f9e810d48212b37812c1856d20ffeafb99ebe66461b61fc7 +# created: 2023-08-02T10:53:29.114535628Z diff --git a/.github/.OwlBot.yaml b/.github/.OwlBot.yaml new file mode 100644 index 00000000000..c379bd3092d --- /dev/null +++ b/.github/.OwlBot.yaml @@ -0,0 +1,18 @@ +# Copyright 2021 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. + +docker: + image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest + +begin-after-commit-hash: 92006bb3cdc84677aa93c7f5235424ec2b157146 diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index cd904459a91..7686a50da62 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -4,6 +4,7 @@ # 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/yoshi-python @googleapis/api-bigquery-dataframe are the default owners for changes in this repo * @googleapis/yoshi-python @googleapis/api-bigquery-dataframe diff --git a/.github/header-checker-lint.yml b/.github/header-checker-lint.yml index 3058bec338c..6fe78aa7987 100644 --- a/.github/header-checker-lint.yml +++ b/.github/header-checker-lint.yml @@ -2,14 +2,14 @@ "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", + "ts", + "js", + "java", + "sh", + "Dockerfile", + "yaml", "py", "html", "txt" ] -} +} \ No newline at end of file diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index 5004e872b4c..b7dae76ba3b 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -15,7 +15,7 @@ branchProtectionRules: - 'unit (3.10)' - 'unit (3.11)' - 'cover' - # TODO(tswast): add Kokoro once we've enabled it + - 'Kokoro presubmit' permissionRules: - team: actools-python permission: admin diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index d2aee5b7d8e..16d5a9e90f6 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -12,7 +12,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.8" - name: Install nox run: | python -m pip install --upgrade setuptools pip wheel diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index bb268fe6c66..465199fc9a9 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -41,7 +41,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.8" - name: Install coverage run: | python -m pip install --upgrade setuptools pip wheel @@ -54,4 +54,4 @@ jobs: - name: Report coverage results run: | coverage combine .coverage-results/.coverage* - coverage report --show-missing --fail-under=40 + coverage report --show-missing --fail-under=35 diff --git a/.gitignore b/.gitignore index 2cb117c23ac..b4243ced74e 100644 --- a/.gitignore +++ b/.gitignore @@ -26,12 +26,9 @@ pip-log.txt # Unit test / coverage reports .coverage -.coverage.* .nox .cache -.mypy_cache .pytest_cache -.pytype # Mac @@ -53,7 +50,6 @@ docs.metadata # Virtual environment env/ -venv/ # Test logs coverage.xml @@ -65,6 +61,3 @@ system_tests/local_test_setup # Make sure a generated file isn't accidentally committed. pylintrc pylintrc.test - -# Notebook scratch -.ipynb_checkpoints diff --git a/.kokoro/build.sh b/.kokoro/build.sh index 402ac0eb8c9..f91c541c6c9 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -16,13 +16,13 @@ set -eo pipefail if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then - PROJECT_SCM="github" + PROJECT_SCM="github/python-bigquery-dataframes" else - PROJECT_SCM="git" + PROJECT_SCM="git/bigframes" fi if [[ -z "${PROJECT_ROOT:-}" ]]; then - PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes" + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}" fi cd "${PROJECT_ROOT}" diff --git a/.kokoro/continuous/continuous.cfg b/.kokoro/continuous/continuous.cfg index 18a4c35325b..8f43917d92f 100644 --- a/.kokoro/continuous/continuous.cfg +++ b/.kokoro/continuous/continuous.cfg @@ -1 +1 @@ -# Format: //devtools/kokoro/config/proto/build.proto +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file 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/docker/docs/Dockerfile b/.kokoro/docker/docs/Dockerfile index e8ee8191ee1..8e39a2cc438 100644 --- a/.kokoro/docker/docs/Dockerfile +++ b/.kokoro/docker/docs/Dockerfile @@ -80,4 +80,4 @@ RUN wget -O /tmp/get-pip.py 'https://bootstrap.pypa.io/get-pip.py' \ # Test pip RUN python3 -m pip -CMD ["python3.9"] +CMD ["python3.8"] diff --git a/.kokoro/docs/common.cfg b/.kokoro/docs/common.cfg index bd73988540c..e6e409f29c2 100644 --- a/.kokoro/docs/common.cfg +++ b/.kokoro/docs/common.cfg @@ -11,7 +11,7 @@ action { gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" # Use the trampoline script to run in docker. -build_file: "bigframes/.kokoro/trampoline_v2.sh" +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" # Configure the docker image for kokoro-trampoline. env_vars: { @@ -20,12 +20,18 @@ env_vars: { } env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: ".kokoro/publish-docs.sh" + value: "github/python-bigquery-dataframes/.kokoro/publish-docs.sh" +} + +env_vars: { + key: "STAGING_BUCKET" + value: "docs-staging" } env_vars: { key: "V2_STAGING_BUCKET" - # Push google cloud library docs to the Cloud RAD bucket `docs-staging-v2` + # Push non-cloud library docs to `docs-staging-v2-staging` instead of the + # Cloud RAD bucket `docs-staging-v2` value: "docs-staging-v2" } @@ -58,4 +64,4 @@ before_action { keyname: "docuploader_service_account" } } -} +} \ No newline at end of file diff --git a/.kokoro/docs/docs-presubmit.cfg b/.kokoro/docs/docs-presubmit.cfg index 1d0dc4b4991..805cfd162b2 100644 --- a/.kokoro/docs/docs-presubmit.cfg +++ b/.kokoro/docs/docs-presubmit.cfg @@ -1,5 +1,10 @@ # Format: //devtools/kokoro/config/proto/build.proto +env_vars: { + key: "STAGING_BUCKET" + value: "gcloud-python-test" +} + env_vars: { key: "V2_STAGING_BUCKET" value: "gcloud-python-test" @@ -13,11 +18,11 @@ env_vars: { env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: ".kokoro/build.sh" + value: "github/python-bigquery-dataframes/.kokoro/build.sh" } # Only run this nox session. env_vars: { key: "NOX_SESSION" - value: "docfx" + value: "docs docfx" } diff --git a/.kokoro/docs/docs.cfg b/.kokoro/docs/docs.cfg index 18a4c35325b..8f43917d92f 100644 --- a/.kokoro/docs/docs.cfg +++ b/.kokoro/docs/docs.cfg @@ -1 +1 @@ -# Format: //devtools/kokoro/config/proto/build.proto +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file diff --git a/.kokoro/populate-secrets.sh b/.kokoro/populate-secrets.sh new file mode 100755 index 00000000000..6f3972140e8 --- /dev/null +++ b/.kokoro/populate-secrets.sh @@ -0,0 +1,43 @@ +#!/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 +# +# 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 index 5d40578ac79..97e0651aa92 100644 --- a/.kokoro/presubmit/common.cfg +++ b/.kokoro/presubmit/common.cfg @@ -7,4 +7,4 @@ action { } } -build_file: "bigframes/.kokoro/build.sh" +build_file: "python-bigquery-dataframes/.kokoro/build.sh" 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.cfg b/.kokoro/presubmit/presubmit.cfg index 18a4c35325b..8f43917d92f 100644 --- a/.kokoro/presubmit/presubmit.cfg +++ b/.kokoro/presubmit/presubmit.cfg @@ -1 +1 @@ -# Format: //devtools/kokoro/config/proto/build.proto +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file diff --git a/.kokoro/publish-docs.sh b/.kokoro/publish-docs.sh index 3fc3905ac1c..7700c90ee92 100755 --- a/.kokoro/publish-docs.sh +++ b/.kokoro/publish-docs.sh @@ -24,6 +24,25 @@ export PATH="${HOME}/.local/bin:${PATH}" python3 -m pip install --require-hashes -r .kokoro/requirements.txt python3 -m nox --version +# build docs +nox -s docs + +# create metadata +python3 -m docuploader create-metadata \ + --name=$(jq --raw-output '.name // empty' .repo-metadata.json) \ + --version=$(python3 setup.py --version) \ + --language=$(jq --raw-output '.language // empty' .repo-metadata.json) \ + --distribution-name=$(python3 setup.py --name) \ + --product-page=$(jq --raw-output '.product_documentation // empty' .repo-metadata.json) \ + --github-repository=$(jq --raw-output '.repo // empty' .repo-metadata.json) \ + --issue-tracker=$(jq --raw-output '.issue_tracker // empty' .repo-metadata.json) + +cat docs.metadata + +# upload docs +python3 -m docuploader upload docs/_build/html --metadata-file docs.metadata --staging-bucket "${STAGING_BUCKET}" + + # docfx yaml files nox -s docfx diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh index 582808a15ce..488dbb9e130 100755 --- a/.kokoro/release-nightly.sh +++ b/.kokoro/release-nightly.sh @@ -35,13 +35,13 @@ while [ $# -gt 0 ] ; do done if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then - PROJECT_SCM="github" + PROJECT_SCM="github/python-bigquery-dataframes" else - PROJECT_SCM="git" + PROJECT_SCM="git/bigframes" fi if [ -z "${PROJECT_ROOT:-}" ]; then - PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes" + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}" fi # Move into the package, build the distribution and upload to shared bucket. @@ -158,49 +158,6 @@ sed -i -e "s/$RELEASE_VERSION/$BIGFRAMES_VERSION/g" bigframes/version.py mv -f ${SETUP_CFG_BKP} setup.cfg rm -f ${THIRD_PARTY_NOTICES_FILE} -# Keep this last so as not to block the release on PDF docs build. -pdf_docs () { - sudo apt update - sudo apt install -y texlive texlive-latex-extra latexmk - - pushd "${PROJECT_ROOT}/docs" - make latexpdf - - cp "_build/latex/bigframes.pdf" "_build/latex/bigframes-${RELEASE_VERSION}.pdf" - cp "_build/latex/bigframes.pdf" "_build/latex/bigframes-latest.pdf" - - if ! [ ${DRY_RUN} ]; then - for gcs_path in gs://vertex_sdk_private_releases/bigframe/ \ - gs://dl-platform-colab/bigframes/ \ - gs://bigframes-wheels/; - do - gsutil cp -v "_build/latex/bigframes-*.pdf" ${gcs_path} - done - fi - - popd -} - -pdf_docs - -# Copy html docs to GCS from where it can be deployed to anywhere else -gcs_docs () { - docs_gcs_bucket=gs://bigframes-docs - docs_local_html_folder=docs/_build/html - if [ ! -d ${docs_local_html_folder} ]; then - python3.10 -m nox -s docs - fi - - if ! [ ${DRY_RUN} ]; then - gsutil -m cp -v -r ${docs_local_html_folder} ${docs_gcs_bucket}/${GIT_HASH} - - # Copy the script to refresh firebase docs website from GCS to GCS itself - gsutil -m cp -v scripts/update_firebase_docs_site.sh ${docs_gcs_bucket} - fi -} - -gcs_docs - if ! [ ${DRY_RUN} ]; then # Copy docs and wheels to Google Drive python3.10 scripts/upload_to_google_drive.py diff --git a/.kokoro/release.sh b/.kokoro/release.sh index 2bce143e1c9..320ac51271e 100755 --- a/.kokoro/release.sh +++ b/.kokoro/release.sh @@ -16,7 +16,7 @@ set -eo pipefail # Start the releasetool reporter -python3 -m pip install --require-hashes -r github/bigframes/.kokoro/requirements.txt +python3 -m pip install --require-hashes -r github/python-bigquery-dataframes/.kokoro/requirements.txt python3 -m releasetool publish-reporter-script > /tmp/publisher-script; source /tmp/publisher-script # Disable buffering, so that the logs stream through. @@ -24,6 +24,6 @@ export PYTHONUNBUFFERED=1 # Move into the package, build the distribution and upload. TWINE_PASSWORD=$(cat "${KOKORO_KEYSTORE_DIR}/73713_google-cloud-pypi-token-keystore-1") -cd github/bigframes +cd github/python-bigquery-dataframes python3 setup.py sdist bdist_wheel twine upload --username __token__ --password "${TWINE_PASSWORD}" dist/* diff --git a/.kokoro/release/common.cfg b/.kokoro/release/common.cfg index 021198fb269..a0c39946cf1 100644 --- a/.kokoro/release/common.cfg +++ b/.kokoro/release/common.cfg @@ -11,7 +11,7 @@ action { gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" # Use the trampoline script to run in docker. -build_file: "bigframes/.kokoro/trampoline_v2.sh" +build_file: "python-bigquery-dataframes/.kokoro/trampoline.sh" # Configure the docker image for kokoro-trampoline. env_vars: { @@ -20,7 +20,7 @@ env_vars: { } env_vars: { key: "TRAMPOLINE_BUILD_FILE" - value: "github/bigframes/.kokoro/release.sh" + value: "github/python-bigquery-dataframes/.kokoro/release.sh" } # Fetch PyPI password @@ -43,7 +43,7 @@ env_vars: { # what we published, which we can use to generate SBOMs and attestations. action { define_artifacts { - regex: "github/bigframes/**/*.tar.gz" - strip_prefix: "github/bigframes" + regex: "github/python-bigquery-dataframes/**/*.tar.gz" + strip_prefix: "github/python-bigquery-dataframes" } } diff --git a/.kokoro/requirements.in b/.kokoro/requirements.in index 9fd9e127baf..ec867d9fd65 100644 --- a/.kokoro/requirements.in +++ b/.kokoro/requirements.in @@ -8,6 +8,3 @@ setuptools nox>=2022.11.21 # required to remove dependency on py charset-normalizer<3 click<8.1.0 -sphinx==4.5.0 # docs PDF see noxfile for HTML docs build -alabaster # docs PDF -recommonmark # docs PDF diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index fd0e0b9d489..029bd342de9 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -1,38 +1,28 @@ # -# This file is autogenerated by pip-compile with Python 3.10 +# This file is autogenerated by pip-compile with Python 3.9 # by the following command: # # pip-compile --allow-unsafe --generate-hashes requirements.in # -alabaster==0.7.13 \ - --hash=sha256:1ee19aca801bbabb5ba3f5f258e4422dfa86f82f3e9cefb0859b283cdd7f62a3 \ - --hash=sha256:a27a4a084d5e690e16e01e03ad2b2e552c61a65469419b907243193de1a84ae2 - # via - # -r requirements.in - # sphinx -argcomplete==3.1.1 \ - --hash=sha256:35fa893a88deea85ea7b20d241100e64516d6af6d7b0ae2bed1d263d26f70948 \ - --hash=sha256:6c4c563f14f01440aaffa3eae13441c5db2357b5eec639abe7c0b15334627dff +argcomplete==2.0.0 \ + --hash=sha256:6372ad78c89d662035101418ae253668445b391755cfe94ea52f1b9d22425b20 \ + --hash=sha256:cffa11ea77999bb0dd27bb25ff6dc142a6796142f68d45b1a26b11f58724561e # via nox -attrs==23.1.0 \ - --hash=sha256:1f28b4522cdc2fb4256ac1a020c78acf9cba2c6b461ccd2c126f3aa8e8335d04 \ - --hash=sha256:6279836d581513a26f1bf235f9acd333bc9115683f14f7e8fae46c98fc50e015 +attrs==22.1.0 \ + --hash=sha256:29adc2665447e5191d0e7c568fde78b21f9672d344281d0c6e1ab085429b22b6 \ + --hash=sha256:86efa402f67bf2df34f51a335487cf46b1ec130d02b8d39fd248abfd30da551c # via gcp-releasetool -babel==2.12.1 \ - --hash=sha256:b4246fb7677d3b98f501a39d43396d3cafdc8eadb045f4a31be01863f655c610 \ - --hash=sha256:cc2d99999cd01d44420ae725a21c9e3711b3aadc7976d6147f622d8581963455 - # via sphinx -bleach==6.0.0 \ - --hash=sha256:1a1a85c1595e07d8db14c5f09f09e6433502c51c595970edc090551f0db99414 \ - --hash=sha256:33c16e3353dbd13028ab4799a0f89a83f113405c766e9c122df8a06f5b85b3f4 +bleach==5.0.1 \ + --hash=sha256:085f7f33c15bd408dd9b17a4ad77c577db66d76203e5984b1bd59baeee948b2a \ + --hash=sha256:0d03255c47eb9bd2f26aa9bb7f2107732e7e8fe195ca2f64709fcf3b0a4a085c # via readme-renderer -cachetools==5.3.1 \ - --hash=sha256:95ef631eeaea14ba2e36f06437f36463aac3a096799e876ee55e5cdccb102590 \ - --hash=sha256:dce83f2d9b4e1f732a8cd44af8e8fab2dbe46201467fc98b3ef8f269092bf62b +cachetools==5.2.0 \ + --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \ + --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db # via google-auth -certifi==2023.5.7 \ - --hash=sha256:0f0d56dc5a6ad56fd4ba36484d6cc34451e1c6548c61daad8c320169f91eddc7 \ - --hash=sha256:c6c2e98f5c7869efca1f8916fed228dd91539f9f1b444c314c06eef02980c716 +certifi==2023.7.22 \ + --hash=sha256:539cc1d13202e33ca466e88b2807e29f4c13049d6d87031a3c110744495cb082 \ + --hash=sha256:92d6037539857d8206b8f6ae472e8b77db8058fec5937a1ef3f54304089edbb9 # via requests cffi==1.15.1 \ --hash=sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5 \ @@ -122,7 +112,7 @@ colorlog==6.7.0 \ commonmark==0.9.1 \ --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 - # via recommonmark + # via rich cryptography==41.0.3 \ --hash=sha256:0d09fb5356f975974dbcb595ad2d178305e5050656affb7890a1583f5e02a306 \ --hash=sha256:23c2d778cf829f7d0ae180600b17e9fceea3c2ef8b31a99e3c694cbbf3a24b84 \ @@ -154,34 +144,31 @@ distlib==0.3.6 \ --hash=sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46 \ --hash=sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e # via virtualenv -docutils==0.17.1 \ - --hash=sha256:686577d2e4c32380bb50cbb22f575ed742d58168cee37e99117a854bcd88f125 \ - --hash=sha256:cf316c8370a737a022b72b56874f6602acf974a37a9fba42ec2876387549fc61 - # via - # readme-renderer - # recommonmark - # sphinx -filelock==3.12.2 \ - --hash=sha256:002740518d8aa59a26b0c76e10fb8c6e15eae825d34b6fdf670333fd7b938d81 \ - --hash=sha256:cbb791cdea2a72f23da6ac5b5269ab0a0d161e9ef0100e653b69049a7706d1ec +docutils==0.19 \ + --hash=sha256:33995a6753c30b7f577febfc2c50411fec6aac7f7ffeb7c4cfe5991072dcf9e6 \ + --hash=sha256:5e1de4d849fee02c63b040a4a3fd567f4ab104defd8a5511fbbc24a8a017efbc + # via readme-renderer +filelock==3.8.0 \ + --hash=sha256:55447caa666f2198c5b6b13a26d2084d26fa5b115c00d065664b2124680c4edc \ + --hash=sha256:617eb4e5eedc82fc5f47b6d61e4d11cb837c56cb4544e39081099fa17ad109d4 # via virtualenv -gcp-docuploader==0.6.5 \ - --hash=sha256:30221d4ac3e5a2b9c69aa52fdbef68cc3f27d0e6d0d90e220fc024584b8d2318 \ - --hash=sha256:b7458ef93f605b9d46a4bf3a8dc1755dad1f31d030c8679edf304e343b347eea +gcp-docuploader==0.6.4 \ + --hash=sha256:01486419e24633af78fd0167db74a2763974765ee8078ca6eb6964d0ebd388af \ + --hash=sha256:70861190c123d907b3b067da896265ead2eeb9263969d6955c9e0bb091b5ccbf # via -r requirements.in -gcp-releasetool==1.14.0 \ - --hash=sha256:a39708206838477bf42804a9dd7b549131ff3fc08c6abb0dd6939e274bd6dfd4 \ - --hash=sha256:aa163ec8322045c9d46e03f8f7500974568bd5d8d322b6a2d492068a5bdf57bf +gcp-releasetool==1.10.5 \ + --hash=sha256:174b7b102d704b254f2a26a3eda2c684fd3543320ec239baf771542a2e58e109 \ + --hash=sha256:e29d29927fe2ca493105a82958c6873bb2b90d503acac56be2c229e74de0eec9 # via -r requirements.in -google-api-core==2.11.1 \ - --hash=sha256:25d29e05a0058ed5f19c61c0a78b1b53adea4d9364b464d014fbda941f6d1c9a \ - --hash=sha256:d92a5a92dc36dd4f4b9ee4e55528a90e432b059f93aee6ad857f9de8cc7ae94a +google-api-core==2.10.2 \ + --hash=sha256:10c06f7739fe57781f87523375e8e1a3a4674bf6392cd6131a3222182b971320 \ + --hash=sha256:34f24bd1d5f72a8c4519773d99ca6bf080a6c4e041b4e9f024fe230191dda62e # via # google-cloud-core # google-cloud-storage -google-auth==2.20.0 \ - --hash=sha256:030af34138909ccde0fbce611afc178f1d65d32fbff281f25738b1fe1c6f3eaa \ - --hash=sha256:23b7b0950fcda519bfb6692bf0d5289d2ea49fc143717cc7188458ec620e63fa +google-auth==2.14.1 \ + --hash=sha256:ccaa901f31ad5cbb562615eb8b664b3dd0bf5404a67618e642307f00613eda4d \ + --hash=sha256:f5d8701633bebc12e0deea4df8abd8aff31c28b355360597f7f2ee60f2e4d016 # via # gcp-releasetool # google-api-core @@ -191,9 +178,9 @@ google-cloud-core==2.3.2 \ --hash=sha256:8417acf6466be2fa85123441696c4badda48db314c607cf1e5d543fa8bdc22fe \ --hash=sha256:b9529ee7047fd8d4bf4a2182de619154240df17fbe60ead399078c1ae152af9a # via google-cloud-storage -google-cloud-storage==2.9.0 \ - --hash=sha256:83a90447f23d5edd045e0037982c270302e3aeb45fc1288d2c2ca713d27bad94 \ - --hash=sha256:9b6ae7b509fc294bdacb84d0f3ea8e20e2c54a8b4bbe39c5707635fec214eff3 +google-cloud-storage==2.6.0 \ + --hash=sha256:104ca28ae61243b637f2f01455cc8a05e8f15a2a18ced96cb587241cdd3820f5 \ + --hash=sha256:4ad0415ff61abdd8bb2ae81c1f8f7ec7d91a1011613f2db87c614c550f97bfe9 # via gcp-docuploader google-crc32c==1.5.0 \ --hash=sha256:024894d9d3cfbc5943f8f230e23950cd4906b2fe004c72e29b209420a1e6b05a \ @@ -265,25 +252,21 @@ google-crc32c==1.5.0 \ --hash=sha256:fd8536e902db7e365f49e7d9029283403974ccf29b13fc7028b97e2295b33556 \ --hash=sha256:fe70e325aa68fa4b5edf7d1a4b6f691eb04bbccac0ace68e34820d283b5f80d4 # via google-resumable-media -google-resumable-media==2.5.0 \ - --hash=sha256:218931e8e2b2a73a58eb354a288e03a0fd5fb1c4583261ac6e4c078666468c93 \ - --hash=sha256:da1bd943e2e114a56d85d6848497ebf9be6a14d3db23e9fc57581e7c3e8170ec +google-resumable-media==2.4.0 \ + --hash=sha256:2aa004c16d295c8f6c33b2b4788ba59d366677c0a25ae7382436cb30f776deaa \ + --hash=sha256:8d5518502f92b9ecc84ac46779bd4f09694ecb3ba38a3e7ca737a86d15cbca1f # via google-cloud-storage -googleapis-common-protos==1.59.1 \ - --hash=sha256:0cbedb6fb68f1c07e18eb4c48256320777707e7d0c55063ae56c15db3224a61e \ - --hash=sha256:b35d530fe825fb4227857bc47ad84c33c809ac96f312e13182bdeaa2abe1178a +googleapis-common-protos==1.57.0 \ + --hash=sha256:27a849d6205838fb6cc3c1c21cb9800707a661bb21c6ce7fb13e99eb1f8a0c46 \ + --hash=sha256:a9f4a1d7f6d9809657b7f1316a1aa527f6664891531bcfcc13b6696e685f443c # via google-api-core idna==3.4 \ --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \ --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2 # via requests -imagesize==1.4.1 \ - --hash=sha256:0d8d18d08f840c19d0ee7ca1fd82490fdc3729b7ac93f49870406ddde8ef8d8b \ - --hash=sha256:69150444affb9cb0d5cc5a92b3676f0b2fb7cd9ae39e947a5e11a36b4497cd4a - # via sphinx -importlib-metadata==6.7.0 \ - --hash=sha256:1aaf550d4f73e5d6783e7acb77aec43d49da8017410afae93822cc9cca98c4d4 \ - --hash=sha256:cb52082e659e97afc5dac71e79de97d8681de3aa07ff18578330904a9d18e5b5 +importlib-metadata==5.0.0 \ + --hash=sha256:da31db32b304314d044d3c12c79bd59e307889b287ad12ff387b3500835fc2ab \ + --hash=sha256:ddb0e35065e8938f867ed4928d0ae5bf2a53b7773871bfe6bcc7e4fcdc7dea43 # via # -r requirements.in # keyring @@ -301,97 +284,76 @@ jeepney==0.8.0 \ jinja2==3.1.2 \ --hash=sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852 \ --hash=sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61 - # via - # gcp-releasetool - # sphinx -keyring==24.0.0 \ - --hash=sha256:4e87665a19c514c7edada8b15015cf89bd99b8d7edabc5c43cca77166fa8dfad \ - --hash=sha256:770f609eed2a16c65a6349f3ba1545d00c73f9fed4254c13766c674fe6d0d22b + # via gcp-releasetool +keyring==23.11.0 \ + --hash=sha256:3dd30011d555f1345dec2c262f0153f2f0ca6bca041fb1dc4588349bb4c0ac1e \ + --hash=sha256:ad192263e2cdd5f12875dedc2da13534359a7e760e77f8d04b50968a821c2361 # via # gcp-releasetool # twine -markdown-it-py==3.0.0 \ - --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ - --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb - # via rich -markupsafe==2.1.3 \ - --hash=sha256:05fb21170423db021895e1ea1e1f3ab3adb85d1c2333cbc2310f2a26bc77272e \ - --hash=sha256:0a4e4a1aff6c7ac4cd55792abf96c915634c2b97e3cc1c7129578aa68ebd754e \ - --hash=sha256:10bbfe99883db80bdbaff2dcf681dfc6533a614f700da1287707e8a5d78a8431 \ - --hash=sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686 \ - --hash=sha256:1577735524cdad32f9f694208aa75e422adba74f1baee7551620e43a3141f559 \ - --hash=sha256:1b40069d487e7edb2676d3fbdb2b0829ffa2cd63a2ec26c4938b2d34391b4ecc \ - --hash=sha256:282c2cb35b5b673bbcadb33a585408104df04f14b2d9b01d4c345a3b92861c2c \ - --hash=sha256:2c1b19b3aaacc6e57b7e25710ff571c24d6c3613a45e905b1fde04d691b98ee0 \ - --hash=sha256:2ef12179d3a291be237280175b542c07a36e7f60718296278d8593d21ca937d4 \ - --hash=sha256:338ae27d6b8745585f87218a3f23f1512dbf52c26c28e322dbe54bcede54ccb9 \ - --hash=sha256:3c0fae6c3be832a0a0473ac912810b2877c8cb9d76ca48de1ed31e1c68386575 \ - --hash=sha256:3fd4abcb888d15a94f32b75d8fd18ee162ca0c064f35b11134be77050296d6ba \ - --hash=sha256:42de32b22b6b804f42c5d98be4f7e5e977ecdd9ee9b660fda1a3edf03b11792d \ - --hash=sha256:504b320cd4b7eff6f968eddf81127112db685e81f7e36e75f9f84f0df46041c3 \ - --hash=sha256:525808b8019e36eb524b8c68acdd63a37e75714eac50e988180b169d64480a00 \ - --hash=sha256:56d9f2ecac662ca1611d183feb03a3fa4406469dafe241673d521dd5ae92a155 \ - --hash=sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac \ - --hash=sha256:65c1a9bcdadc6c28eecee2c119465aebff8f7a584dd719facdd9e825ec61ab52 \ - --hash=sha256:68e78619a61ecf91e76aa3e6e8e33fc4894a2bebe93410754bd28fce0a8a4f9f \ - --hash=sha256:69c0f17e9f5a7afdf2cc9fb2d1ce6aabdb3bafb7f38017c0b77862bcec2bbad8 \ - --hash=sha256:6b2b56950d93e41f33b4223ead100ea0fe11f8e6ee5f641eb753ce4b77a7042b \ - --hash=sha256:787003c0ddb00500e49a10f2844fac87aa6ce977b90b0feaaf9de23c22508b24 \ - --hash=sha256:7ef3cb2ebbf91e330e3bb937efada0edd9003683db6b57bb108c4001f37a02ea \ - --hash=sha256:8023faf4e01efadfa183e863fefde0046de576c6f14659e8782065bcece22198 \ - --hash=sha256:8758846a7e80910096950b67071243da3e5a20ed2546e6392603c096778d48e0 \ - --hash=sha256:8afafd99945ead6e075b973fefa56379c5b5c53fd8937dad92c662da5d8fd5ee \ - --hash=sha256:8c41976a29d078bb235fea9b2ecd3da465df42a562910f9022f1a03107bd02be \ - --hash=sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2 \ - --hash=sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707 \ - --hash=sha256:962f82a3086483f5e5f64dbad880d31038b698494799b097bc59c2edf392fce6 \ - --hash=sha256:9dcdfd0eaf283af041973bff14a2e143b8bd64e069f4c383416ecd79a81aab58 \ - --hash=sha256:aa7bd130efab1c280bed0f45501b7c8795f9fdbeb02e965371bbef3523627779 \ - --hash=sha256:ab4a0df41e7c16a1392727727e7998a467472d0ad65f3ad5e6e765015df08636 \ - --hash=sha256:ad9e82fb8f09ade1c3e1b996a6337afac2b8b9e365f926f5a61aacc71adc5b3c \ - --hash=sha256:af598ed32d6ae86f1b747b82783958b1a4ab8f617b06fe68795c7f026abbdcad \ - --hash=sha256:b076b6226fb84157e3f7c971a47ff3a679d837cf338547532ab866c57930dbee \ - --hash=sha256:b7ff0f54cb4ff66dd38bebd335a38e2c22c41a8ee45aa608efc890ac3e3931bc \ - --hash=sha256:bfce63a9e7834b12b87c64d6b155fdd9b3b96191b6bd334bf37db7ff1fe457f2 \ - --hash=sha256:c011a4149cfbcf9f03994ec2edffcb8b1dc2d2aede7ca243746df97a5d41ce48 \ - --hash=sha256:c9c804664ebe8f83a211cace637506669e7890fec1b4195b505c214e50dd4eb7 \ - --hash=sha256:ca379055a47383d02a5400cb0d110cef0a776fc644cda797db0c5696cfd7e18e \ - --hash=sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b \ - --hash=sha256:cd0f502fe016460680cd20aaa5a76d241d6f35a1c3350c474bac1273803893fa \ - --hash=sha256:ceb01949af7121f9fc39f7d27f91be8546f3fb112c608bc4029aef0bab86a2a5 \ - --hash=sha256:d080e0a5eb2529460b30190fcfcc4199bd7f827663f858a226a81bc27beaa97e \ - --hash=sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb \ - --hash=sha256:df0be2b576a7abbf737b1575f048c23fb1d769f267ec4358296f31c2479db8f9 \ - --hash=sha256:e09031c87a1e51556fdcb46e5bd4f59dfb743061cf93c4d6831bf894f125eb57 \ - --hash=sha256:e4dd52d80b8c83fdce44e12478ad2e85c64ea965e75d66dbeafb0a3e77308fcc \ - --hash=sha256:fec21693218efe39aa7f8599346e90c705afa52c5b31ae019b2e57e8f6542bb2 +markupsafe==2.1.1 \ + --hash=sha256:0212a68688482dc52b2d45013df70d169f542b7394fc744c02a57374a4207003 \ + --hash=sha256:089cf3dbf0cd6c100f02945abeb18484bd1ee57a079aefd52cffd17fba910b88 \ + --hash=sha256:10c1bfff05d95783da83491be968e8fe789263689c02724e0c691933c52994f5 \ + --hash=sha256:33b74d289bd2f5e527beadcaa3f401e0df0a89927c1559c8566c066fa4248ab7 \ + --hash=sha256:3799351e2336dc91ea70b034983ee71cf2f9533cdff7c14c90ea126bfd95d65a \ + --hash=sha256:3ce11ee3f23f79dbd06fb3d63e2f6af7b12db1d46932fe7bd8afa259a5996603 \ + --hash=sha256:421be9fbf0ffe9ffd7a378aafebbf6f4602d564d34be190fc19a193232fd12b1 \ + --hash=sha256:43093fb83d8343aac0b1baa75516da6092f58f41200907ef92448ecab8825135 \ + --hash=sha256:46d00d6cfecdde84d40e572d63735ef81423ad31184100411e6e3388d405e247 \ + --hash=sha256:4a33dea2b688b3190ee12bd7cfa29d39c9ed176bda40bfa11099a3ce5d3a7ac6 \ + --hash=sha256:4b9fe39a2ccc108a4accc2676e77da025ce383c108593d65cc909add5c3bd601 \ + --hash=sha256:56442863ed2b06d19c37f94d999035e15ee982988920e12a5b4ba29b62ad1f77 \ + --hash=sha256:671cd1187ed5e62818414afe79ed29da836dde67166a9fac6d435873c44fdd02 \ + --hash=sha256:694deca8d702d5db21ec83983ce0bb4b26a578e71fbdbd4fdcd387daa90e4d5e \ + --hash=sha256:6a074d34ee7a5ce3effbc526b7083ec9731bb3cbf921bbe1d3005d4d2bdb3a63 \ + --hash=sha256:6d0072fea50feec76a4c418096652f2c3238eaa014b2f94aeb1d56a66b41403f \ + --hash=sha256:6fbf47b5d3728c6aea2abb0589b5d30459e369baa772e0f37a0320185e87c980 \ + --hash=sha256:7f91197cc9e48f989d12e4e6fbc46495c446636dfc81b9ccf50bb0ec74b91d4b \ + --hash=sha256:86b1f75c4e7c2ac2ccdaec2b9022845dbb81880ca318bb7a0a01fbf7813e3812 \ + --hash=sha256:8dc1c72a69aa7e082593c4a203dcf94ddb74bb5c8a731e4e1eb68d031e8498ff \ + --hash=sha256:8e3dcf21f367459434c18e71b2a9532d96547aef8a871872a5bd69a715c15f96 \ + --hash=sha256:8e576a51ad59e4bfaac456023a78f6b5e6e7651dcd383bcc3e18d06f9b55d6d1 \ + --hash=sha256:96e37a3dc86e80bf81758c152fe66dbf60ed5eca3d26305edf01892257049925 \ + --hash=sha256:97a68e6ada378df82bc9f16b800ab77cbf4b2fada0081794318520138c088e4a \ + --hash=sha256:99a2a507ed3ac881b975a2976d59f38c19386d128e7a9a18b7df6fff1fd4c1d6 \ + --hash=sha256:a49907dd8420c5685cfa064a1335b6754b74541bbb3706c259c02ed65b644b3e \ + --hash=sha256:b09bf97215625a311f669476f44b8b318b075847b49316d3e28c08e41a7a573f \ + --hash=sha256:b7bd98b796e2b6553da7225aeb61f447f80a1ca64f41d83612e6139ca5213aa4 \ + --hash=sha256:b87db4360013327109564f0e591bd2a3b318547bcef31b468a92ee504d07ae4f \ + --hash=sha256:bcb3ed405ed3222f9904899563d6fc492ff75cce56cba05e32eff40e6acbeaa3 \ + --hash=sha256:d4306c36ca495956b6d568d276ac11fdd9c30a36f1b6eb928070dc5360b22e1c \ + --hash=sha256:d5ee4f386140395a2c818d149221149c54849dfcfcb9f1debfe07a8b8bd63f9a \ + --hash=sha256:dda30ba7e87fbbb7eab1ec9f58678558fd9a6b8b853530e176eabd064da81417 \ + --hash=sha256:e04e26803c9c3851c931eac40c695602c6295b8d432cbe78609649ad9bd2da8a \ + --hash=sha256:e1c0b87e09fa55a220f058d1d49d3fb8df88fbfab58558f1198e08c1e1de842a \ + --hash=sha256:e72591e9ecd94d7feb70c1cbd7be7b3ebea3f548870aa91e2732960fa4d57a37 \ + --hash=sha256:e8c843bbcda3a2f1e3c2ab25913c80a3c5376cd00c6e8c4a86a89a28c8dc5452 \ + --hash=sha256:efc1913fd2ca4f334418481c7e595c00aad186563bbc1ec76067848c7ca0a933 \ + --hash=sha256:f121a1420d4e173a5d96e47e9a0c0dcff965afdf1626d28de1460815f7c4ee7a \ + --hash=sha256:fc7b548b17d238737688817ab67deebb30e8073c95749d55538ed473130ec0c7 # via jinja2 -mdurl==0.1.2 \ - --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \ - --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba - # via markdown-it-py -more-itertools==9.1.0 \ - --hash=sha256:cabaa341ad0389ea83c17a94566a53ae4c9d07349861ecb14dc6d0345cf9ac5d \ - --hash=sha256:d2bc7f02446e86a68911e58ded76d6561eea00cddfb2a91e7019bbb586c799f3 +more-itertools==9.0.0 \ + --hash=sha256:250e83d7e81d0c87ca6bd942e6aeab8cc9daa6096d12c5308f3f92fa5e5c1f41 \ + --hash=sha256:5a6257e40878ef0520b1803990e3e22303a41b5714006c32a3fd8304b26ea1ab # via jaraco-classes -nox==2023.4.22 \ - --hash=sha256:0b1adc619c58ab4fa57d6ab2e7823fe47a32e70202f287d78474adcc7bda1891 \ - --hash=sha256:46c0560b0dc609d7d967dc99e22cb463d3c4caf54a5fda735d6c11b5177e3a9f +nox==2022.11.21 \ + --hash=sha256:0e41a990e290e274cb205a976c4c97ee3c5234441a8132c8c3fd9ea3c22149eb \ + --hash=sha256:e21c31de0711d1274ca585a2c5fde36b1aa962005ba8e9322bf5eeed16dcd684 # via -r requirements.in -packaging==23.1 \ - --hash=sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61 \ - --hash=sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f +packaging==21.3 \ + --hash=sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb \ + --hash=sha256:ef103e05f519cdc783ae24ea4e2e0f508a9c99b2d4969652eed6a2e1ea5bd522 # via # gcp-releasetool # nox - # sphinx -pkginfo==1.9.6 \ - --hash=sha256:4b7a555a6d5a22169fcc9cf7bfd78d296b0361adad412a346c1226849af5e546 \ - --hash=sha256:8fd5896e8718a4372f0ea9cc9d96f6417c9b986e23a4d116dda26b62cc29d046 +pkginfo==1.8.3 \ + --hash=sha256:848865108ec99d4901b2f7e84058b6e7660aae8ae10164e015a6dcf5b242a594 \ + --hash=sha256:a84da4318dd86f870a9447a8c98340aa06216bfc6f2b7bdc4b8766984ae1867c # via twine -platformdirs==3.7.0 \ - --hash=sha256:87fbf6473e87c078d536980ba970a472422e94f17b752cfad17024c18876d481 \ - --hash=sha256:cfd065ba43133ff103ab3bd10aecb095c2a0035fcd1f07217c9376900d94ba07 +platformdirs==2.5.4 \ + --hash=sha256:1006647646d80f16130f052404c6b901e80ee4ed6bef6792e1f238a8969106f7 \ + --hash=sha256:af0276409f9a02373d540bf8480021a048711d572745aef4b7842dad245eba10 # via virtualenv protobuf==3.20.3 \ --hash=sha256:03038ac1cfbc41aa21f6afcbcd357281d7521b4157926f30ebecc8d4ea59dcb7 \ @@ -420,32 +382,34 @@ protobuf==3.20.3 \ # gcp-docuploader # gcp-releasetool # google-api-core - # googleapis-common-protos -pyasn1==0.5.0 \ - --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \ - --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde +pyasn1==0.4.8 \ + --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \ + --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba # via # pyasn1-modules # rsa -pyasn1-modules==0.3.0 \ - --hash=sha256:5bd01446b736eb9d31512a30d46c1ac3395d676c6f3cafa4c03eb54b9925631c \ - --hash=sha256:d3ccd6ed470d9ffbc716be08bd90efbd44d0734bc9303818f7336070984a162d +pyasn1-modules==0.2.8 \ + --hash=sha256:905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e \ + --hash=sha256:a50b808ffeb97cb3601dd25981f6b016cbb3d31fbf57a8b8a87428e6158d0c74 # via google-auth pycparser==2.21 \ --hash=sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9 \ --hash=sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206 # via cffi -pygments==2.15.1 \ - --hash=sha256:8ace4d3c1dd481894b2005f560ead0f9f19ee64fe983366be1a21e171d12775c \ - --hash=sha256:db2db3deb4b4179f399a09054b023b6a586b76499d36965813c71aa8ed7b5fd1 +pygments==2.15.0 \ + --hash=sha256:77a3299119af881904cd5ecd1ac6a66214b6e9bed1f2db16993b54adede64094 \ + --hash=sha256:f7e36cffc4c517fbc252861b9a6e4644ca0e5abadf9a113c72d1358ad09b9500 # via # readme-renderer # rich - # sphinx -pyjwt==2.7.0 \ - --hash=sha256:ba2b425b15ad5ef12f200dc67dd56af4e26de2331f965c5439994dad075876e1 \ - --hash=sha256:bd6ca4a3c4285c1a2d4349e5a035fdf8fb94e04ccd0fcbe6ba289dae9cc3e074 +pyjwt==2.6.0 \ + --hash=sha256:69285c7e31fc44f68a1feb309e948e0df53259d579295e6cfe2b1792329f05fd \ + --hash=sha256:d83c3d892a77bbb74d3e1a2cfa90afaadb60945205d1095d9221f04466f64c14 # via gcp-releasetool +pyparsing==3.0.9 \ + --hash=sha256:2b020ecf7d21b687f219b71ecad3631f644a47f01403fa1d1036b0c6416d70fb \ + --hash=sha256:5026bae9a10eeaefb61dab2f09052b9f4307d44aee4eda64b309723d8d206bbc + # via packaging pyperclip==1.8.2 \ --hash=sha256:105254a8b04934f0bc84e9c24eb360a591aaf6535c9def5f29d92af107a9bf57 # via gcp-releasetool @@ -453,14 +417,10 @@ python-dateutil==2.8.2 \ --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \ --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9 # via gcp-releasetool -readme-renderer==40.0 \ - --hash=sha256:9f77b519d96d03d7d7dce44977ba543090a14397c4f60de5b6eb5b8048110aa4 \ - --hash=sha256:e18feb2a1e7706f2865b81ebb460056d93fb29d69daa10b223c00faa7bd9a00a +readme-renderer==37.3 \ + --hash=sha256:cd653186dfc73055656f090f227f5cb22a046d7f71a841dfa305f55c9a513273 \ + --hash=sha256:f67a16caedfa71eef48a31b39708637a6f4664c4394801a7b0d6432d13907343 # via twine -recommonmark==0.7.1 \ - --hash=sha256:1b1db69af0231efce3fa21b94ff627ea33dee7079a01dd0a7f8482c3da148b3f \ - --hash=sha256:bdb4db649f2222dcd8d2d844f0006b958d627f732415d399791ee436a3686d67 - # via -r requirements.in requests==2.31.0 \ --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1 @@ -469,19 +429,18 @@ requests==2.31.0 \ # google-api-core # google-cloud-storage # requests-toolbelt - # sphinx # twine -requests-toolbelt==1.0.0 \ - --hash=sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6 \ - --hash=sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06 +requests-toolbelt==0.10.1 \ + --hash=sha256:18565aa58116d9951ac39baa288d3adb5b3ff975c4f25eee78555d89e8f247f7 \ + --hash=sha256:62e09f7ff5ccbda92772a29f394a49c3ad6cb181d568b1337626b2abb628a63d # via twine rfc3986==2.0.0 \ --hash=sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd \ --hash=sha256:97aacf9dbd4bfd829baad6e6309fa6573aaf1be3f6fa735c8ab05e46cecb261c # via twine -rich==13.4.2 \ - --hash=sha256:8f87bc7ee54675732fa66a05ebfe489e27264caeeff3728c945d25971b6485ec \ - --hash=sha256:d653d6bccede5844304c605d5aac802c7cf9621efd700b46c7ec2b51ea914898 +rich==12.6.0 \ + --hash=sha256:a4eb26484f2c82589bd9a17c73d32a010b1e29d89f1604cd9bf3a2097b81bb5e \ + --hash=sha256:ba3a3775974105c221d31141f2c116f4fd65c5ceb0698657a11e9f295ec93fd0 # via twine rsa==4.9 \ --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ @@ -499,74 +458,39 @@ six==1.16.0 \ # gcp-docuploader # google-auth # python-dateutil -snowballstemmer==2.2.0 \ - --hash=sha256:09b16deb8547d3412ad7b590689584cd0fe25ec8db3be37788be3810cbf19cb1 \ - --hash=sha256:c8e1716e83cc398ae16824e5572ae04e0d9fc2c6b985fb0f900f5f0c96ecba1a - # via sphinx -sphinx==4.5.0 \ - --hash=sha256:7bf8ca9637a4ee15af412d1a1d9689fec70523a68ca9bb9127c2f3eeb344e2e6 \ - --hash=sha256:ebf612653238bcc8f4359627a9b7ce44ede6fdd75d9d30f68255c7383d3a6226 - # via - # -r requirements.in - # recommonmark -sphinxcontrib-applehelp==1.0.4 \ - --hash=sha256:29d341f67fb0f6f586b23ad80e072c8e6ad0b48417db2bde114a4c9746feb228 \ - --hash=sha256:828f867945bbe39817c210a1abfd1bc4895c8b73fcaade56d45357a348a07d7e - # via sphinx -sphinxcontrib-devhelp==1.0.2 \ - --hash=sha256:8165223f9a335cc1af7ffe1ed31d2871f325254c0423bc0c4c7cd1c1e4734a2e \ - --hash=sha256:ff7f1afa7b9642e7060379360a67e9c41e8f3121f2ce9164266f61b9f4b338e4 - # via sphinx -sphinxcontrib-htmlhelp==2.0.1 \ - --hash=sha256:0cbdd302815330058422b98a113195c9249825d681e18f11e8b1f78a2f11efff \ - --hash=sha256:c38cb46dccf316c79de6e5515e1770414b797162b23cd3d06e67020e1d2a6903 - # via sphinx -sphinxcontrib-jsmath==1.0.1 \ - --hash=sha256:2ec2eaebfb78f3f2078e73666b1415417a116cc848b72e5172e596c871103178 \ - --hash=sha256:a9925e4a4587247ed2191a22df5f6970656cb8ca2bd6284309578f2153e0c4b8 - # via sphinx -sphinxcontrib-qthelp==1.0.3 \ - --hash=sha256:4c33767ee058b70dba89a6fc5c1892c0d57a54be67ddd3e7875a18d14cba5a72 \ - --hash=sha256:bd9fc24bcb748a8d51fd4ecaade681350aa63009a347a8c14e637895444dfab6 - # via sphinx -sphinxcontrib-serializinghtml==1.1.5 \ - --hash=sha256:352a9a00ae864471d3a7ead8d7d79f5fc0b57e8b3f95e9867eb9eb28999b92fd \ - --hash=sha256:aa5f6de5dfdf809ef505c4895e51ef5c9eac17d0f287933eb49ec495280b6952 - # via sphinx -twine==4.0.2 \ - --hash=sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8 \ - --hash=sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8 +twine==4.0.1 \ + --hash=sha256:42026c18e394eac3e06693ee52010baa5313e4811d5a11050e7d48436cf41b9e \ + --hash=sha256:96b1cf12f7ae611a4a40b6ae8e9570215daff0611828f5fe1f37a16255ab24a0 # via -r requirements.in -typing-extensions==4.6.3 \ - --hash=sha256:88a4153d8505aabbb4e13aacb7c486c2b4a33ca3b3f807914a9b4c844c471c26 \ - --hash=sha256:d91d5919357fe7f681a9f2b5b4cb2a5f1ef0a1e9f59c4d8ff0d3491e05c0ffd5 +typing-extensions==4.4.0 \ + --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \ + --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e # via -r requirements.in -urllib3==1.26.16 \ - --hash=sha256:8d36afa7616d8ab714608411b4a3b13e58f463aee519024578e062e141dce20f \ - --hash=sha256:8f135f6502756bde6b2a9b28989df5fbe87c9970cecaa69041edcce7f0589b14 +urllib3==1.26.12 \ + --hash=sha256:3fa96cf423e6987997fc326ae8df396db2a8b7c667747d47ddd8ecba91f4a74e \ + --hash=sha256:b930dd878d5a8afb066a637fbb35144fe7901e3b209d1cd4f524bd0e9deee997 # via - # google-auth # requests # twine -virtualenv==20.23.1 \ - --hash=sha256:34da10f14fea9be20e0fd7f04aba9732f84e593dac291b757ce42e3368a39419 \ - --hash=sha256:8ff19a38c1021c742148edc4f81cb43d7f8c6816d2ede2ab72af5b84c749ade1 +virtualenv==20.16.7 \ + --hash=sha256:8691e3ff9387f743e00f6bb20f70121f5e4f596cae754531f2b3b3a1b1ac696e \ + --hash=sha256:efd66b00386fdb7dbe4822d172303f40cd05e50e01740b19ea42425cbe653e29 # via nox webencodings==0.5.1 \ --hash=sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78 \ --hash=sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923 # via bleach -wheel==0.40.0 \ - --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \ - --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247 +wheel==0.38.4 \ + --hash=sha256:965f5259b566725405b05e7cf774052044b1ed30119b5d586b2703aafe8719ac \ + --hash=sha256:b60533f3f5d530e971d6737ca6d58681ee434818fab630c83a734bb10c083ce8 # via -r requirements.in -zipp==3.15.0 \ - --hash=sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b \ - --hash=sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556 +zipp==3.10.0 \ + --hash=sha256:4fcb6f278987a6605757302a6e40e896257570d11c51628968ccb2a47e80c6c1 \ + --hash=sha256:7a7262fd930bd3e36c50b9a64897aec3fafff3dfdeec9623ae22b40e93f99bb8 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: -setuptools==68.0.0 \ - --hash=sha256:11e52c67415a381d10d6b462ced9cfb97066179f0e871399e006c4ab101fc85f \ - --hash=sha256:baf1fdb41c6da4cd2eae722e135500da913332ab3f2f5c7d33af9b492acb5235 +setuptools==65.5.1 \ + --hash=sha256:d0b9a8433464d5800cbe05094acf5c6d52a91bfac9b52bcfc4d41382be5d5d31 \ + --hash=sha256:e197a19aa8ec9722928f2206f8de752def0e4c9fc6953527360d1c36d94ddb2f # via -r requirements.in 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.7/common.cfg b/.kokoro/samples/python3.7/common.cfg new file mode 100644 index 00000000000..09d7af02ba9 --- /dev/null +++ b/.kokoro/samples/python3.7/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.7" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-py37" +} + +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.7/continuous.cfg b/.kokoro/samples/python3.7/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.7/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.7/periodic-head.cfg b/.kokoro/samples/python3.7/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.7/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.7/periodic.cfg b/.kokoro/samples/python3.7/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.7/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.7/presubmit.cfg b/.kokoro/samples/python3.7/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.7/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.8/common.cfg b/.kokoro/samples/python3.8/common.cfg new file mode 100644 index 00000000000..976d9ce8c5c --- /dev/null +++ b/.kokoro/samples/python3.8/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.8" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-py38" +} + +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.8/continuous.cfg b/.kokoro/samples/python3.8/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.8/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.8/periodic-head.cfg b/.kokoro/samples/python3.8/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.8/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.8/periodic.cfg b/.kokoro/samples/python3.8/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.8/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.8/presubmit.cfg b/.kokoro/samples/python3.8/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.8/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.9/common.cfg b/.kokoro/samples/python3.9/common.cfg new file mode 100644 index 00000000000..603cfffa280 --- /dev/null +++ b/.kokoro/samples/python3.9/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.9" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-py39" +} + +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.9/continuous.cfg b/.kokoro/samples/python3.9/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.9/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.9/periodic-head.cfg b/.kokoro/samples/python3.9/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.9/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.9/periodic.cfg b/.kokoro/samples/python3.9/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.9/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.9/presubmit.cfg b/.kokoro/samples/python3.9/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.9/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..63ac41dfae1 --- /dev/null +++ b/.kokoro/test-samples-against-head.sh @@ -0,0 +1,26 @@ +#!/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. + +# 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..5a0f5fab6a8 --- /dev/null +++ b/.kokoro/test-samples-impl.sh @@ -0,0 +1,102 @@ +#!/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. + + +# `-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 +python3.9 -m pip install --upgrade --quiet nox + +# 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.9 -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..50b35a48c19 --- /dev/null +++ b/.kokoro/test-samples.sh @@ -0,0 +1,44 @@ +#!/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. + +# 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..d85b1f26769 --- /dev/null +++ b/.kokoro/trampoline.sh @@ -0,0 +1,28 @@ +#!/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 +# +# 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 old mode 100644 new mode 100755 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 07d7e6b7445..6e0fd8b98fb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,4 @@ -# Copyright 2021 Google LLC +# 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. @@ -29,11 +29,11 @@ repos: - repo: https://github.com/psf/black rev: 22.3.0 hooks: - - id: black + - id: black - repo: https://github.com/pycqa/flake8 - rev: 3.9.2 + rev: 6.1.0 hooks: - - id: flake8 + - id: flake8 - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.1.1 hooks: diff --git a/.trampolinerc b/.trampolinerc index 4c16e33fee0..a7dfeb42c6d 100644 --- a/.trampolinerc +++ b/.trampolinerc @@ -20,6 +20,9 @@ required_envvars+=( pass_down_envvars+=( "NOX_SESSION" ############### + # Docs builds + ############### + "STAGING_BUCKET" "V2_STAGING_BUCKET" ################## # Samples builds diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md index 51d6abc4d3d..039f4368120 100644 --- a/CODE_OF_CONDUCT.md +++ b/CODE_OF_CONDUCT.md @@ -1,3 +1,4 @@ + # Code of Conduct ## Our Pledge @@ -91,4 +92,4 @@ harassment or threats to anyone's safety, we may take action without notice. 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 +https://www.contributor-covenant.org/version/1/4/code-of-conduct.html \ No newline at end of file diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst new file mode 100644 index 00000000000..3933152cf78 --- /dev/null +++ b/CONTRIBUTING.rst @@ -0,0 +1,277 @@ +.. Generated by synthtool. DO NOT EDIT! +############ +Contributing +############ + +#. **Please sign one of the contributor license agreements below.** +#. Fork the repo, develop and test your code changes, add docs. +#. Make sure that your commit messages clearly describe the changes. +#. Send a pull request. (Please Read: `Faster Pull Request Reviews`_) + +.. _Faster Pull Request Reviews: https://github.com/kubernetes/community/blob/master/contributors/guide/pull-requests.md#best-practices-for-faster-reviews + +.. contents:: Here are some guidelines for hacking on the Google Cloud Client libraries. + +*************** +Adding Features +*************** + +In order to add a feature: + +- The feature must be documented in both the API and narrative + documentation. + +- The feature must work fully on the following CPython versions: + 3.9, 3.10 and 3.11 on both UNIX and Windows. + +- The feature must not add unnecessary dependencies (where + "unnecessary" is of course subjective, but new dependencies should + be discussed). + +**************************** +Using a Development Checkout +**************************** + +You'll have to create a development environment using a Git checkout: + +- While logged into your GitHub account, navigate to the + ``python-bigquery-dataframes`` `repo`_ on GitHub. + +- Fork and clone the ``python-bigquery-dataframes`` repository to your GitHub account by + clicking the "Fork" button. + +- Clone your fork of ``python-bigquery-dataframes`` from your GitHub account to your local + computer, substituting your account username and specifying the destination + as ``hack-on-python-bigquery-dataframes``. E.g.:: + + $ cd ${HOME} + $ git clone git@github.com:USERNAME/python-bigquery-dataframes.git hack-on-python-bigquery-dataframes + $ cd hack-on-python-bigquery-dataframes + # Configure remotes such that you can pull changes from the googleapis/python-bigquery-dataframes + # repository into your local repository. + $ git remote add upstream git@github.com:googleapis/python-bigquery-dataframes.git + # fetch and merge changes from upstream into main + $ git fetch upstream + $ git merge upstream/main + +Now your local repo is set up such that you will push changes to your GitHub +repo, from which you can submit a pull request. + +To work on the codebase and run the tests, we recommend using ``nox``, +but you can also use a ``virtualenv`` of your own creation. + +.. _repo: https://github.com/googleapis/python-bigquery-dataframes + +Using ``nox`` +============= + +We use `nox `__ to instrument our tests. + +- To test your changes, run unit tests with ``nox``:: + $ nox -s unit + +- To run a single unit test:: + + $ nox -s unit-3.11 -- -k + + + .. note:: + + The unit tests and system tests are described in the + ``noxfile.py`` files in each directory. + +.. nox: https://pypi.org/project/nox/ + +***************************************** +I'm getting weird errors... Can you help? +***************************************** + +If the error mentions ``Python.h`` not being found, +install ``python-dev`` and try again. +On Debian/Ubuntu:: + + $ sudo apt-get install python-dev + +************ +Coding Style +************ +- We use the automatic code formatter ``black``. You can run it using + the nox session ``blacken``. This will eliminate many lint errors. Run via:: + + $ nox -s blacken + +- PEP8 compliance is required, with exceptions defined in the linter configuration. + If you have ``nox`` installed, you can test that you have not introduced + any non-compliant code via:: + + $ nox -s lint + +- In order to make ``nox -s lint`` run faster, you can set some environment + variables:: + + export GOOGLE_CLOUD_TESTING_REMOTE="upstream" + export GOOGLE_CLOUD_TESTING_BRANCH="main" + + By doing this, you are specifying the location of the most up-to-date + version of ``python-bigquery-dataframes``. The + remote name ``upstream`` should point to the official ``googleapis`` + checkout and the branch should be the default branch on that remote (``main``). + +- This repository contains configuration for the + `pre-commit `__ tool, which automates checking + our linters during a commit. If you have it installed on your ``$PATH``, + you can enable enforcing those checks via: + +.. code-block:: bash + + $ pre-commit install + pre-commit installed at .git/hooks/pre-commit + +Exceptions to PEP8: + +- Many unit tests use a helper method, ``_call_fut`` ("FUT" is short for + "Function-Under-Test"), which is PEP8-incompliant, but more readable. + Some also use a local variable, ``MUT`` (short for "Module-Under-Test"). + +******************** +Running System Tests +******************** + +- To run system tests, you can execute:: + + # Run all system tests + $ nox -s system + + # Run a single system test + $ nox -s system-3.11 -- -k + + + .. note:: + + System tests are only configured to run under Python 3.9 and 3.11. + For expediency, we do not run them in older versions of Python 3. + + This alone will not run the tests. You'll need to change some local + auth settings and change some configuration in your project to + run all the tests. + +- System tests will be run against an actual project. You should use local credentials from gcloud when possible. See `Best practices for application authentication `__. Some tests require a service account. For those tests see `Authenticating as a service account `__. + +************* +Test Coverage +************* + +- The codebase *must* have 100% test statement coverage after each commit. + You can test coverage via ``nox -s cover``. + +****************************************************** +Documentation Coverage and Building HTML Documentation +****************************************************** + +If you fix a bug, and the bug requires an API or behavior modification, all +documentation in this package which references that API or behavior must be +changed to reflect the bug fix, ideally in the same commit that fixes the bug +or adds the feature. + +Build the docs via: + + $ nox -s docs + +************************* +Samples and code snippets +************************* + +Code samples and snippets live in the `samples/` catalogue. Feel free to +provide more examples, but make sure to write tests for those examples. +Each folder containing example code requires its own `noxfile.py` script +which automates testing. If you decide to create a new folder, you can +base it on the `samples/snippets` folder (providing `noxfile.py` and +the requirements files). + +The tests will run against a real Google Cloud Project, so you should +configure them just like the System Tests. + +- To run sample tests, you can execute:: + + # Run all tests in a folder + $ cd samples/snippets + $ nox -s py-3.8 + + # Run a single sample test + $ cd samples/snippets + $ nox -s py-3.8 -- -k + +******************************************** +Note About ``README`` as it pertains to PyPI +******************************************** + +The `description on PyPI`_ for the project comes directly from the +``README``. Due to the reStructuredText (``rst``) parser used by +PyPI, relative links which will work on GitHub (e.g. ``CONTRIBUTING.rst`` +instead of +``https://github.com/googleapis/python-bigquery-dataframes/blob/main/CONTRIBUTING.rst``) +may cause problems creating links or rendering the description. + +.. _description on PyPI: https://pypi.org/project/bigframes + + +************************* +Supported Python Versions +************************* + +We support: + +- `Python 3.9`_ +- `Python 3.10`_ +- `Python 3.11`_ + +.. _Python 3.9: https://docs.python.org/3.9/ +.. _Python 3.10: https://docs.python.org/3.10/ +.. _Python 3.11: https://docs.python.org/3.11/ + + +Supported versions can be found in our ``noxfile.py`` `config`_. + +.. _config: https://github.com/googleapis/python-bigquery-dataframes/blob/main/noxfile.py + + +We also explicitly decided to support Python 3 beginning with version 3.9. +Reasons for this include: + +- Encouraging use of newest versions of Python 3 +- Taking the lead of `prominent`_ open-source `projects`_ +- `Unicode literal support`_ which allows for a cleaner codebase that + works in both Python 2 and Python 3 + +.. _prominent: https://docs.djangoproject.com/en/1.9/faq/install/#what-python-version-can-i-use-with-django +.. _projects: http://flask.pocoo.org/docs/0.10/python3/ +.. _Unicode literal support: https://www.python.org/dev/peps/pep-0414/ + +********** +Versioning +********** + +This library follows `Semantic Versioning`_. + +.. _Semantic Versioning: http://semver.org/ + +Some packages are currently in major version zero (``0.y.z``), which means that +anything may change at any time and the public API should not be considered +stable. + +****************************** +Contributor License Agreements +****************************** + +Before we can accept your pull requests you'll need to sign a Contributor +License Agreement (CLA): + +- **If you are an individual writing original source code** and **you own the + intellectual property**, then you'll need to sign an + `individual CLA `__. +- **If you work for a company that wants to allow you to contribute your work**, + then you'll need to sign a + `corporate CLA `__. + +You can sign these electronically (just scroll to the bottom). After that, +we'll be able to accept your pull requests. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000000..b422266a96a --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- +# +# 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. + +# Generated by synthtool. DO NOT EDIT! +include README.rst LICENSE +recursive-include third_party * +recursive-include bigframes *.json *.proto py.typed +recursive-include tests * +global-exclude *.py[co] +global-exclude __pycache__ + +# Exclude scripts for samples readmegen +prune scripts/readme-gen diff --git a/docs/_static/custom.css b/docs/_static/custom.css index c663a6ef337..b0a295464b2 100644 --- a/docs/_static/custom.css +++ b/docs/_static/custom.css @@ -1,3 +1,8 @@ +div#python2-eol { + border-color: red; + border-width: medium; +} + /* Ensure minimum width for 'Parameters' / 'Returns' column */ dl.field-list > dt { min-width: 100px diff --git a/docs/_templates/layout.html b/docs/_templates/layout.html index bfdcc4759c5..6316a537f72 100644 --- a/docs/_templates/layout.html +++ b/docs/_templates/layout.html @@ -20,6 +20,11 @@ {% endblock %}
+ {% block body %} {% endblock %}
diff --git a/docs/conf.py b/docs/conf.py index 3ab5be51035..af8c5efda89 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# 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. @@ -25,8 +25,8 @@ # serve to show the default. import os +import shlex import sys -from typing import Any, Dict # If extensions (or modules to document with autodoc) are in another directory, # add these directories to sys.path here. If the directory is relative to the @@ -80,8 +80,8 @@ root_doc = "index" # General information about the project. -project = "BigQuery DataFrames" -copyright = "2022-2023 Google LLC" +project = "bigframes" +copyright = "2019, Google" author = "Google APIs" # The version info for the project you're documenting, acts as replacement for @@ -98,7 +98,7 @@ # # This is also used if you do content translation via gettext catalogs. # Usually you set "language" from the command line for these cases. -language = "en" +language = None # There are two options for replacing |today|: either, you set today to some # non-false value, then it is used: @@ -154,10 +154,10 @@ # further. For a list of options available for each theme, see the # documentation. html_theme_options = { - "description": "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine.", - # "github_user": "googleapis", - # "github_repo": "python-bigquery-storage", - # "github_banner": True, + "description": "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine", + "github_user": "googleapis", + "github_repo": "python-bigquery-dataframes", + "github_banner": True, "font_family": "'Roboto', Georgia, sans", "head_font_family": "'Roboto', Georgia, serif", "code_font_family": "'Roboto Mono', 'Consolas', monospace", @@ -185,7 +185,7 @@ # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -# html_static_path = ["_static"] +html_static_path = ["_static"] # Add any extra paths that contain custom files (such as robots.txt or # .htaccess) here, relative to this directory. These files are copied @@ -264,13 +264,15 @@ # -- Options for LaTeX output --------------------------------------------- -latex_elements: Dict[str, Any] = { - # Avoid "too deeply nested" error by using enumitem package. - # See: https://stackoverflow.com/a/28454426/101923 - "preamble": r""" -\usepackage{enumitem} -\setlistdepth{99} -""" +latex_elements = { + # The paper size ('letterpaper' or 'a4paper'). + #'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + #'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + #'preamble': '', + # Latex figure (float) alignment + #'figure_align': 'htbp', } # Grouping the document tree into LaTeX files. List of tuples @@ -366,6 +368,11 @@ "grpc": ("https://grpc.github.io/grpc/python/", None), "proto-plus": ("https://proto-plus-python.readthedocs.io/en/latest/", None), "protobuf": ("https://googleapis.dev/python/protobuf/latest/", None), + "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), + "pydata-google-auth": ( + "https://pydata-google-auth.readthedocs.io/en/latest/", + None, + ), } diff --git a/noxfile.py b/noxfile.py index 1ceca6831b2..2355a9b27ba 100644 --- a/noxfile.py +++ b/noxfile.py @@ -31,7 +31,7 @@ BLACK_VERSION = "black==22.3.0" ISORT_VERSION = "isort==5.12.0" SPHINX_VERSION = "sphinx==4.5.0" -LINT_PATHS = ["docs", "bigframes", "tests", "noxfile.py", "setup.py"] +LINT_PATHS = ["docs", "bigframes", "tests", "third_party", "noxfile.py", "setup.py"] DEFAULT_PYTHON_VERSION = "3.10" @@ -108,6 +108,7 @@ def lint(session): "--check", *LINT_PATHS, ) + # TODO(tswast): lint all LINT_PATHS session.run("flake8", "bigframes", "tests") diff --git a/owlbot.py b/owlbot.py index 4ba7d14eb5e..be30eea5c2c 100644 --- a/owlbot.py +++ b/owlbot.py @@ -15,6 +15,7 @@ """This script is used to synthesize generated parts of this library.""" import pathlib +import re from synthtool import gcp import synthtool as s @@ -27,11 +28,10 @@ # ---------------------------------------------------------------------------- # Add templated files # ---------------------------------------------------------------------------- - templated_files = common.py_library( unit_test_python_versions=["3.9", "3.10", "3.11"], system_test_python_versions=["3.9", "3.11"], - cov_level=40, + cov_level=35, intersphinx_dependencies={ "pandas": "https://pandas.pydata.org/pandas-docs/stable/", "pydata-google-auth": "https://pydata-google-auth.readthedocs.io/en/latest/", @@ -40,11 +40,17 @@ s.move( templated_files, excludes=[ - # Multi-processing note isn't relevant, as pandas_gbq is responsible for + # Multi-processing note isn't relevant, as bigframes is responsible for # creating clients, not the end user. "docs/multiprocessing.rst", "noxfile.py", + ".pre-commit-config.yaml", "README.rst", + ".github/release-trigger.yml", + # BigQuery DataFrames manages its own Kokoro cluster for presubmit & continuous tests. + ".kokoro/build.sh", + ".kokoro/continuous/common.cfg", + ".kokoro/presubmit/common.cfg", ], ) @@ -52,6 +58,46 @@ # Fixup files # ---------------------------------------------------------------------------- +# Make sure build includes all necessary files. +s.replace( + ["MANIFEST.in"], + re.escape("recursive-include google"), + "recursive-include third_party *\nrecursive-include bigframes", +) + +# Even though BigQuery DataFrames isn't technically a client library, we are +# opting into Cloud RAD for docs hosting. +s.replace( + [".kokoro/docs/common.cfg"], + re.escape('value: "docs-staging-v2-staging"'), + 'value: "docs-staging-v2"', +) + +# Use a custom table of contents since the default one isn't organized well +# enough for the number of classes we have. +s.replace( + [".kokoro/publish-docs.sh"], + ( + re.escape("# upload docs") + + "\n" + + re.escape( + 'python3 -m docuploader upload docs/_build/html/docfx_yaml --metadata-file docs.metadata --destination-prefix docfx --staging-bucket "${V2_STAGING_BUCKET}"' + ) + ), + ( + "# Replace toc.yml template file\n" + + "mv docs/templates/toc.yml docs/_build/html/docfx_yaml/toc.yml\n\n" + + "# upload docs\n" + + 'python3 -m docuploader upload docs/_build/html/docfx_yaml --metadata-file docs.metadata --destination-prefix docfx --staging-bucket "${V2_STAGING_BUCKET}"' + ), +) + +# Fixup the documentation. +s.replace( + ["docs/conf.py"], + re.escape("Google Cloud Client Libraries for bigframes"), + "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine", +) # ---------------------------------------------------------------------------- # Samples templates @@ -63,6 +109,6 @@ # Final cleanup # ---------------------------------------------------------------------------- -s.shell.run(["nox", "-s", "blacken"], hide_output=False) +s.shell.run(["nox", "-s", "format"], hide_output=False) for noxfile in REPO_ROOT.glob("samples/**/noxfile.py"): s.shell.run(["nox", "-s", "blacken"], cwd=noxfile.parent, hide_output=False) diff --git a/renovate.json b/renovate.json new file mode 100644 index 00000000000..39b2a0ec929 --- /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"], + "pip_requirements": { + "fileMatch": ["requirements-test.txt", "samples/[\\S/]*constraints.txt", "samples/[\\S/]*constraints-test.txt"] + } +} diff --git a/scripts/decrypt-secrets.sh b/scripts/decrypt-secrets.sh new file mode 100755 index 00000000000..0018b421ddf --- /dev/null +++ b/scripts/decrypt-secrets.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +# Copyright 2023 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..1acc119835b --- /dev/null +++ b/scripts/readme-gen/readme_gen.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python + +# 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. + +"""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/scripts/upload_to_google_drive.py b/scripts/upload_to_google_drive.py index e5791513595..dcdc9168ba2 100644 --- a/scripts/upload_to_google_drive.py +++ b/scripts/upload_to_google_drive.py @@ -41,12 +41,9 @@ wheel_id = "15fZ1DkrFDk4ibMNTzms4akpxmf2pzeAR" wheel_path = next(iter((repo_root / "dist").glob("bigframes-*.whl"))) -pdf_id = "1agYjxmPLrxelsaHI-lc41QHcgnQYemcX" -pdf_path = repo_root / "docs" / "_build" / "latex" / "bigframes-latest.pdf" - uploads = ( (wheel_id, wheel_path, "application/octet-stream"), - (pdf_id, pdf_path, "application/pdf"), + # (pdf_id, pdf_path, "application/pdf"), ) upload_template = ( diff --git a/setup.cfg b/setup.cfg index 8bd749387ed..05235008950 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2020 Google LLC +# 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. @@ -17,17 +17,3 @@ # Generated by synthtool. DO NOT EDIT! [bdist_wheel] universal = 1 - -[pytype] -python_version = 3.9 -inputs = - google/cloud/ -exclude = - tests/ - google/cloud/bigquery_v2/ # Legacy proto-based types. -output = .pytype/ -disable = - # There's some issue with finding some pyi files, thus disabling. - # The issue https://github.com/google/pytype/issues/150 is closed, but the - # error still occurs for some reason. - pyi-error diff --git a/testing/.gitignore b/testing/.gitignore new file mode 100644 index 00000000000..b05fbd63088 --- /dev/null +++ b/testing/.gitignore @@ -0,0 +1,3 @@ +test-env.sh +service-account.json +client-secrets.json \ No newline at end of file From 7c6b0dd2f99139c8830e762201a45b28486532ff Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Sat, 16 Sep 2023 12:27:51 +0000 Subject: [PATCH 19/80] feat: add `DataFrame.combine` and `DataFrame.combine_first` (#27) feat: add `DataFrame.skew` and `GroupBy.skew` test: remove unneeded mock perf: `bigframes-api` label to I/O query jobs fix: `remote_function` uses same credentials as other APIs test: BQML golden SQL unit tests feat: add `DataFrame.pct_change` and `Series.pct_change` test: disable `remote_function` reuse in tests test: fix flaky repr_cache tests test: add unit tests for private `ArrayValue` class feat: add `DataFrame.to_dict`, `to_excel`, `to_latex`, `to_records`, `to_string`, `to_markdown`, `to_pickle`, `to_orc` fix: use for literals `Int64Dtype` in `cut` feat: add `DataFrame.nlargest`, `nsmallest` chore: refactor PCA tests feat: add `bfill` and `ffill` to `DataFrame` and `Series` feat: add `reindex_like` to `DataFrame` and `Series` fix: use lowercase strings for parameter literals in `bigframes.ml` (**breaking change**) feat: support `DataFrame.loc[bool_series, column] = scalar` fix: support column joins with "None indexer" docs: document region logic in README feat: add partial support for `Sereies.replace` fix: add type hints to models test: add more unit tests for internal `ArrayValue` feat: add `filter` and `reindex` to `Series` and `DataFrame` docs: document possible parameter values for PaLM2TextGenerator test: mark generate_text test as flaky feat: support a persistent `name` in `remote_function` fix: raise error when ARIMAPlus is used with Pipeline feat: add `swaplevel` to `DataFrame` and `Series` feat: add `axis` parameter to `droplevel` and `reorder_levels` docs: fix OneHotEncoder sample fix: remove `transforms` parameter in `model.fit` (**breaking change**) feat: add `diff` method to `DataFrame` and `GroupBy` --- .kokoro/docs/docs-presubmit-gerrit.cfg | 23 + .kokoro/presubmit/e2e-gerrit.cfg | 7 + .kokoro/presubmit/presubmit-gerrit.cfg | 1 + OWNERS | 1 + README.rst | 11 +- bigframes/core/__init__.py | 2 +- bigframes/core/block_transforms.py | 176 +++++++ bigframes/core/blocks.py | 16 +- bigframes/core/groupby/__init__.py | 44 +- bigframes/core/indexers.py | 42 +- bigframes/core/indexes/index.py | 20 + bigframes/core/utils.py | 4 +- bigframes/dataframe.py | 472 ++++++++++++++++- bigframes/dtypes.py | 2 +- bigframes/ml/base.py | 32 +- bigframes/ml/cluster.py | 6 +- bigframes/ml/compose.py | 2 +- bigframes/ml/decomposition.py | 6 +- bigframes/ml/ensemble.py | 64 +-- bigframes/ml/forecasting.py | 4 +- bigframes/ml/linear_model.py | 11 +- bigframes/ml/llm.py | 8 +- bigframes/ml/pipeline.py | 16 +- bigframes/ml/preprocessing.py | 4 +- bigframes/operations/aggregations.py | 70 ++- bigframes/pandas/__init__.py | 2 + bigframes/remote_function.py | 192 ++++--- bigframes/series.py | 218 +++++--- bigframes/session.py | 107 +++- noxfile.py | 1 + samples/snippets/remote_function.py | 24 +- setup.py | 3 +- testing/constraints-3.9.txt | 5 +- tests/system/large/ml/test_ensemble.py | 10 +- tests/system/large/test_remote_function.py | 296 ++++++++++- tests/system/small/ml/test_core.py | 2 + tests/system/small/ml/test_decomposition.py | 29 +- tests/system/small/test_dataframe.py | 489 ++++++++++++++++++ tests/system/small/test_groupby.py | 26 + tests/system/small/test_ipython.py | 5 +- tests/system/small/test_multiindex.py | 162 +++++- tests/system/small/test_pandas.py | 14 + tests/system/small/test_series.py | 179 ++++++- tests/unit/ml/test_golden_sql.py | 47 ++ tests/unit/ml/test_pipeline.py | 44 +- tests/unit/test_core.py | 169 ++++++ .../ibis/backends/bigquery/registry.py | 12 + .../ibis/expr/operations/__init__.py | 1 + .../ibis/expr/operations/analytic.py | 26 + .../bigframes_vendored/pandas/core/frame.py | 481 ++++++++++++++++- .../bigframes_vendored/pandas/core/generic.py | 89 ++++ .../pandas/core/groupby/__init__.py | 20 + .../pandas/core/indexes/base.py | 10 + .../bigframes_vendored/pandas/core/series.py | 124 ++++- .../bigframes_vendored/sklearn/base.py | 1 + .../sklearn/cluster/_kmeans.py | 5 - .../sklearn/decomposition/_pca.py | 4 - .../sklearn/ensemble/_forest.py | 4 - .../sklearn/linear_model/_base.py | 5 - .../sklearn/linear_model/_logistic.py | 5 - .../sklearn/preprocessing/_data.py | 31 +- .../sklearn/preprocessing/_encoder.py | 14 +- .../bigframes_vendored/xgboost/sklearn.py | 4 - 63 files changed, 3487 insertions(+), 417 deletions(-) create mode 100644 .kokoro/docs/docs-presubmit-gerrit.cfg create mode 100644 .kokoro/presubmit/e2e-gerrit.cfg create mode 100644 .kokoro/presubmit/presubmit-gerrit.cfg create mode 100644 tests/unit/ml/test_golden_sql.py create mode 100644 third_party/bigframes_vendored/ibis/expr/operations/analytic.py diff --git a/.kokoro/docs/docs-presubmit-gerrit.cfg b/.kokoro/docs/docs-presubmit-gerrit.cfg new file mode 100644 index 00000000000..1d0dc4b4991 --- /dev/null +++ b/.kokoro/docs/docs-presubmit-gerrit.cfg @@ -0,0 +1,23 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "V2_STAGING_BUCKET" + value: "gcloud-python-test" +} + +# We only upload the image in the main `docs` build. +env_vars: { + key: "TRAMPOLINE_IMAGE_UPLOAD" + value: "false" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: ".kokoro/build.sh" +} + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "docfx" +} diff --git a/.kokoro/presubmit/e2e-gerrit.cfg b/.kokoro/presubmit/e2e-gerrit.cfg new file mode 100644 index 00000000000..d875f360603 --- /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 samples" +} 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/OWNERS b/OWNERS index 672da38afa5..f86ad551efa 100644 --- a/OWNERS +++ b/OWNERS @@ -1,3 +1,4 @@ +ashleyxu@google.com bmil@google.com chelsealin@google.com garrettwu@google.com diff --git a/README.rst b/README.rst index 935c54cc8b5..23aea446ff7 100644 --- a/README.rst +++ b/README.rst @@ -57,9 +57,13 @@ internally to manage metadata on the service side. This session is tied to a BigQuery DataFrames uses the US multi-region as the default location, but you can use ``session_options.location`` to set a different location. Every query in a session is executed in the location where the session was created. +BigQuery DataFrames +auto-populates ``bf.options.bigquery.location`` if the user starts with +``read_gbq/read_gbq_table/read_gbq_query()`` and specifies a table, either +directly or in a SQL statement. If you want to reset the location of the created DataFrame or Series objects, -can reset the session by executing ``bigframes.pandas.reset_session()``. +you can reset the session by executing ``bigframes.pandas.reset_session()``. After that, you can reuse ``bigframes.pandas.options.bigquery.location`` to specify another location. @@ -68,6 +72,11 @@ specify another location. querying is not in the US multi-region. If you try to read a table from another location, you get a NotFound exception. +Project +------- +If ``bf.options.bigquery.project`` is not set, the ``$GOOGLE_CLOUD_PROJECT`` +environment variable is used, which is set in the notebook runtime serving the +BigQuery Studio/Vertex Notebooks. ML Capabilities --------------- diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 7086269af9b..27fe4a4fe6c 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -518,8 +518,8 @@ def aggregate( """ Apply aggregations to the expression. Arguments: - by_column_id: column id of the aggregation key, this is preserved through the transform aggregations: input_column_id, operation, output_column_id tuples + by_column_id: column id of the aggregation key, this is preserved through the transform dropna: whether null keys should be dropped """ table = self.to_ibis_expr(ordering_mode="unordered") diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index abf8b887d82..5dcd9fe7538 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -117,6 +117,25 @@ def value_counts( return block.select_column(count_id).with_column_labels(["count"]) +def pct_change(block: blocks.Block, periods: int = 1) -> blocks.Block: + column_labels = block.column_labels + window_spec = core.WindowSpec( + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + + original_columns = block.value_columns + block, shift_columns = block.multi_apply_window_op( + original_columns, agg_ops.ShiftOp(periods), window_spec=window_spec + ) + result_ids = [] + for original_col, shifted_col in zip(original_columns, shift_columns): + block, change_id = block.apply_binary_op(original_col, shifted_col, ops.sub_op) + block, pct_change_id = block.apply_binary_op(change_id, shifted_col, ops.div_op) + result_ids.append(pct_change_id) + return block.select_columns(result_ids).with_column_labels(column_labels) + + def rank( block: blocks.Block, method: str = "average", @@ -229,3 +248,160 @@ def dropna(block: blocks.Block, how: typing.Literal["all", "any"] = "any"): filtered_block = filtered_block.filter(predicate) filtered_block = filtered_block.select_columns(block.value_columns) return filtered_block + + +def nsmallest( + block: blocks.Block, + n: int, + column_ids: typing.Sequence[str], + keep: str, +) -> blocks.Block: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + if keep == "last": + block = block.reversed() + order_refs = [ + ordering.OrderingColumnReference( + col_id, direction=ordering.OrderingDirection.ASC + ) + for col_id in column_ids + ] + block = block.order_by(order_refs, stable=True) + if keep in ("first", "last"): + return block.slice(0, n) + else: # keep == "all": + block, counter = block.apply_window_op( + column_ids[0], + agg_ops.rank_op, + window_spec=core.WindowSpec(ordering=order_refs), + ) + block, condition = block.apply_unary_op( + counter, ops.partial_right(ops.le_op, n) + ) + block = block.filter(condition) + return block.drop_columns([counter, condition]) + + +def nlargest( + block: blocks.Block, + n: int, + column_ids: typing.Sequence[str], + keep: str, +) -> blocks.Block: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + if keep == "last": + block = block.reversed() + order_refs = [ + ordering.OrderingColumnReference( + col_id, direction=ordering.OrderingDirection.DESC + ) + for col_id in column_ids + ] + block = block.order_by(order_refs, stable=True) + if keep in ("first", "last"): + return block.slice(0, n) + else: # keep == "all": + block, counter = block.apply_window_op( + column_ids[0], + agg_ops.rank_op, + window_spec=core.WindowSpec(ordering=order_refs), + ) + block, condition = block.apply_unary_op( + counter, ops.partial_right(ops.le_op, n) + ) + block = block.filter(condition) + return block.drop_columns([counter, condition]) + + +def skew( + block: blocks.Block, + skew_column_ids: typing.Sequence[str], + grouping_column_ids: typing.Sequence[str] = (), +) -> blocks.Block: + + original_columns = skew_column_ids + column_labels = block.select_columns(original_columns).column_labels + + block, delta3_ids = _mean_delta_to_power( + block, 3, original_columns, grouping_column_ids + ) + # counts, moment3 for each column + aggregations = [] + for i, col in enumerate(original_columns): + count_agg = (col, agg_ops.count_op) + moment3_agg = (delta3_ids[i], agg_ops.mean_op) + variance_agg = (col, agg_ops.PopVarOp()) + aggregations.extend([count_agg, moment3_agg, variance_agg]) + + block, agg_ids = block.aggregate( + by_column_ids=grouping_column_ids, aggregations=aggregations + ) + + skew_ids = [] + for i, col in enumerate(original_columns): + # Corresponds to order of aggregations in preceding loop + count_id, moment3_id, var_id = agg_ids[i * 3 : (i * 3) + 3] + block, skew_id = _skew_from_moments_and_count( + block, count_id, moment3_id, var_id + ) + skew_ids.append(skew_id) + + block = block.select_columns(skew_ids).with_column_labels(column_labels) + if not grouping_column_ids: + # When ungrouped, stack everything into single column so can be returned as series + block = block.stack() + block = block.drop_levels([block.index_columns[0]]) + return block + + +def _mean_delta_to_power( + block: blocks.Block, + n_power, + column_ids: typing.Sequence[str], + grouping_column_ids: typing.Sequence[str], +) -> typing.Tuple[blocks.Block, typing.Sequence[str]]: + """Calculate (x-mean(x))^n. Useful for calculating moment statistics such as skew and kurtosis.""" + window = core.WindowSpec(grouping_keys=grouping_column_ids) + block, mean_ids = block.multi_apply_window_op(column_ids, agg_ops.mean_op, window) + delta_ids = [] + cube_op = ops.partial_right(ops.pow_op, n_power) + for val_id, mean_val_id in zip(column_ids, mean_ids): + block, delta_id = block.apply_binary_op(val_id, mean_val_id, ops.sub_op) + block, delta_power_id = block.apply_unary_op(delta_id, cube_op) + block = block.drop_columns(delta_id) + delta_ids.append(delta_power_id) + return block, delta_ids + + +def _skew_from_moments_and_count( + block: blocks.Block, count_id: str, moment3_id: str, var_id: str +) -> typing.Tuple[blocks.Block, str]: + # Calculate skew using count, third moment and population variance + # See G1 estimator: + # https://en.wikipedia.org/wiki/Skewness#Sample_skewness + block, denominator_id = block.apply_unary_op( + var_id, ops.partial_right(ops.pow_op, 3 / 2) + ) + block, base_id = block.apply_binary_op(moment3_id, denominator_id, ops.div_op) + block, countminus1_id = block.apply_unary_op( + count_id, ops.partial_right(ops.sub_op, 1) + ) + block, countminus2_id = block.apply_unary_op( + count_id, ops.partial_right(ops.sub_op, 2) + ) + block, adjustment_id = block.apply_binary_op(count_id, countminus1_id, ops.mul_op) + block, adjustment_id = block.apply_unary_op( + adjustment_id, ops.partial_right(ops.pow_op, 1 / 2) + ) + block, adjustment_id = block.apply_binary_op( + adjustment_id, countminus2_id, ops.div_op + ) + block, skew_id = block.apply_binary_op(base_id, adjustment_id, ops.mul_op) + + # Need to produce NA if have less than 3 data points + block, na_cond_id = block.apply_unary_op(count_id, ops.partial_right(ops.ge_op, 3)) + block, skew_id = block.apply_binary_op( + skew_id, na_cond_id, ops.partial_arg3(ops.where_op, None) + ) + return block, skew_id diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 482cfd0141e..5b414252ee9 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -709,8 +709,9 @@ def multi_apply_window_op( window_spec: core.WindowSpec, *, skip_null_groups: bool = False, - ) -> Block: + ) -> typing.Tuple[Block, typing.Sequence[str]]: block = self + result_ids = [] for i, col_id in enumerate(columns): label = self.col_id_to_label[col_id] block, result_id = block.apply_window_op( @@ -721,9 +722,8 @@ def multi_apply_window_op( result_label=label, skip_null_groups=skip_null_groups, ) - block = block.copy_values(result_id, col_id) - block = block.drop_columns([result_id]) - return block + result_ids.append(result_id) + return block, result_ids def multi_apply_unary_op( self, @@ -1123,7 +1123,9 @@ def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]: ) def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: - axis_number = bigframes.core.utils.get_axis_number(axis) + axis_number = bigframes.core.utils.get_axis_number( + "rows" if (axis is None) else axis + ) if axis_number == 0: expr = self._expr for index_col in self._index_columns: @@ -1140,7 +1142,9 @@ def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block: return self.rename(columns=lambda label: f"{prefix}{label}") def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block: - axis_number = bigframes.core.utils.get_axis_number(axis) + axis_number = bigframes.core.utils.get_axis_number( + "rows" if (axis is None) else axis + ) if axis_number == 0: expr = self._expr for index_col in self._index_columns: diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 589c5c251ca..810e145d33f 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -20,6 +20,7 @@ import bigframes.constants as constants import bigframes.core as core +import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks import bigframes.core.ordering as order import bigframes.core.utils as utils @@ -145,6 +146,16 @@ def var( self._raise_on_non_numeric("var") return self._aggregate_all(agg_ops.var_op, numeric_only=True) + def skew( + self, + *, + numeric_only: bool = False, + ) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("skew") + block = block_ops.skew(self._block, self._selected_cols, self._by_col_ids) + return df.DataFrame(block) + def all(self) -> df.DataFrame: return self._aggregate_all(agg_ops.all_op) @@ -168,6 +179,22 @@ def cummax(self, *args, numeric_only: bool = False, **kwargs) -> df.DataFrame: def cumprod(self, *args, **kwargs) -> df.DataFrame: return self._apply_window_op(agg_ops.product_op, numeric_only=True) + def shift(self, periods=1) -> series.Series: + window = core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.ShiftOp(periods), window=window) + + def diff(self, periods=1) -> series.Series: + window = core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.DiffOp(periods), window=window) + def agg(self, func=None, **kwargs) -> df.DataFrame: if func: if isinstance(func, str): @@ -323,10 +350,10 @@ def _apply_window_op( grouping_keys=self._by_col_ids, following=0 ) columns = self._aggregated_columns(numeric_only=numeric_only) - block = self._block.multi_apply_window_op( + block, result_ids = self._block.multi_apply_window_op( columns, op, window_spec=window_spec, skip_null_groups=self._dropna ) - block = block.select_columns(columns) + block = block.select_columns(result_ids) return df.DataFrame(block) def _resolve_label(self, label: blocks.Label) -> str: @@ -391,6 +418,10 @@ def std(self, *args, **kwargs) -> series.Series: def var(self, *args, **kwargs) -> series.Series: return self._aggregate(agg_ops.var_op) + def skew(self, *args, **kwargs) -> series.Series: + block = block_ops.skew(self._block, [self._value_column], self._by_col_ids) + return series.Series(block) + def prod(self, *args) -> series.Series: return self._aggregate(agg_ops.product_op) @@ -459,8 +490,13 @@ def shift(self, periods=1) -> series.Series: ) return self._apply_window_op(agg_ops.ShiftOp(periods), window=window) - def diff(self) -> series.Series: - return self._ungroup() - self.shift(1) + def diff(self, periods=1) -> series.Series: + window = core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.DiffOp(periods), window=window) def rolling(self, window: int, min_periods=None) -> windows.Window: # To get n size window, need current row and n-1 preceding rows. diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 46091f211ac..28bce05338c 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -145,23 +145,41 @@ def __setitem__( value: bigframes.dataframe.SingleItemValue, ): if ( - not isinstance(key, tuple) - or len(key) != 2 - or not isinstance(key[0], slice) - or (key[0].start is not None and key[0].start != 0) - or (key[0].step is not None and key[0].step != 1) - or key[0].stop is not None + isinstance(key, tuple) + and len(key) == 2 + and isinstance(key[0], slice) + and (key[0].start is None or key[0].start == 0) + and (key[0].step is None or key[0].step == 1) + and key[0].stop is None ): + # TODO(swast): Support setting multiple columns with key[1] as a list + # of labels and value as a DataFrame. + df = self._dataframe.assign(**{key[1]: value}) + self._dataframe._set_block(df._get_block()) + elif ( + isinstance(key, tuple) + and len(key) == 2 + and isinstance(key[0], bigframes.series.Series) + and key[0].dtype == "boolean" + ) and pd.api.types.is_scalar(value): + new_column = key[0].map({True: value, False: None}) + try: + original_column = self._dataframe[key[1]] + except KeyError: + self._dataframe[key[1]] = new_column + return + try: + self._dataframe[key[1]] = new_column.fillna(original_column) + except ibis.common.exceptions.IbisTypeError: + raise TypeError( + f"Cannot assign scalar of type {type(value)} to column of type {original_column.dtype}, or index type of series argument does not match dataframe." + ) + else: raise NotImplementedError( - "Only setting a column by DataFrame.loc[:, 'column'] is supported." + "Only DataFrame.loc[:, 'column'] and DataFrame.loc[bool series, 'column'] = Scalar are supported." f"{constants.FEEDBACK_LINK}" ) - # TODO(swast): Support setting multiple columns with key[1] as a list - # of labels and value as a DataFrame. - df = self._dataframe.assign(**{key[1]: value}) - self._dataframe._set_block(df._get_block()) - class ILocDataFrameIndexer: def __init__(self, dataframe: bigframes.dataframe.DataFrame): diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 04b9a36b64a..748a68c9447 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -53,6 +53,10 @@ def names(self) -> typing.Sequence[blocks.Label]: def names(self, values: typing.Sequence[blocks.Label]): return self._data._set_block(self._data._get_block().with_index_labels(values)) + @property + def nlevels(self) -> int: + return len(self._data._get_block().index_columns) + @property def shape(self) -> typing.Tuple[int]: return (self._data._get_block().shape[0],) @@ -97,6 +101,22 @@ def is_monotonic_decreasing(self) -> bool: ), ) + @property + def is_unique(self) -> bool: + # TODO: Cache this at block level + # Avoid circular imports + import bigframes.core.block_transforms as block_ops + import bigframes.dataframe as df + + duplicates_block, _ = block_ops.indicate_duplicates( + self._data._get_block(), self._data._get_block().index_columns + ) + duplicates_block = duplicates_block.with_column_labels( + ["values", "is_duplicate"] + ) + duplicates_df = df.DataFrame(duplicates_block) + return not duplicates_df["is_duplicate"].any() + def __getitem__(self, key: int) -> typing.Any: if isinstance(key, int): result_pd_df, _ = self._data._get_block().slice(key, key + 1, 1).to_pandas() diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py index 1c0a2a1a81e..75175690ce0 100644 --- a/bigframes/core/utils.py +++ b/bigframes/core/utils.py @@ -23,8 +23,8 @@ UNNAMED_INDEX_ID = "bigframes_unnamed_index" -def get_axis_number(axis: typing.Union[str, int, None]) -> typing.Literal[0, 1]: - if axis in {0, "index", "rows", None}: +def get_axis_number(axis: typing.Union[str, int]) -> typing.Literal[0, 1]: + if axis in {0, "index", "rows"}: return 0 elif axis in {1, "columns"}: return 1 diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 31777f3face..d65d4ce344e 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -455,7 +455,7 @@ def __getattr__(self, key: str): raise AttributeError(key) def __repr__(self) -> str: - """Converts a DataFrame to a string. Calls compute. + """Converts a DataFrame to a string. Calls to_pandas. Only represents the first `bigframes.options.display.max_rows`. """ @@ -532,13 +532,14 @@ def _apply_binop( other: float | int | bigframes.series.Series | DataFrame, op, axis: str | int = "columns", + how: str = "outer", ): if isinstance(other, (float, int)): return self._apply_scalar_binop(other, op) elif isinstance(other, bigframes.series.Series): - return self._apply_series_binop(other, op, axis=axis) + return self._apply_series_binop(other, op, axis=axis, how=how) elif isinstance(other, DataFrame): - return self._apply_dataframe_binop(other, op) + return self._apply_dataframe_binop(other, op, how=how) raise NotImplementedError( f"binary operation is not implemented on the second operand of type {type(other).__name__}." f"{constants.FEEDBACK_LINK}" @@ -559,6 +560,7 @@ def _apply_series_binop( other: bigframes.series.Series, op: ops.BinaryOp, axis: str | int = "columns", + how: str = "outer", ) -> DataFrame: if axis not in ("columns", "index", 0, 1): raise ValueError(f"Invalid input: axis {axis}.") @@ -569,7 +571,7 @@ def _apply_series_binop( ) joined_index, (get_column_left, get_column_right) = self._block.index.join( - other._block.index, how="outer" + other._block.index, how=how ) series_column_id = other._value.get_name() @@ -591,22 +593,27 @@ def _apply_series_binop( return DataFrame(block) def _apply_dataframe_binop( - self, - other: DataFrame, - op: ops.BinaryOp, + self, other: DataFrame, op: ops.BinaryOp, how: str = "outer" ) -> DataFrame: # Join rows joined_index, (get_column_left, get_column_right) = self._block.index.join( - other._block.index, how="outer" + other._block.index, how=how ) # join columns schema + # indexers will be none for exact match columns, lcol_indexer, rcol_indexer = self.columns.join( - other.columns, how="outer", return_indexers=True + other.columns, how=how, return_indexers=True ) binop_result_ids = [] block = joined_index._block - for left_index, right_index in zip(lcol_indexer, rcol_indexer): + + column_indices = zip( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)), + rcol_indexer if (lcol_indexer is not None) else range(len(columns)), + ) + + for left_index, right_index in column_indices: if left_index >= 0 and right_index >= 0: # -1 indices indicate missing left_col_id = self._block.value_columns[left_index] right_col_id = other._block.value_columns[right_index] @@ -617,13 +624,19 @@ def _apply_dataframe_binop( ) binop_result_ids.append(result_col_id) elif left_index >= 0: - dtype = self.dtypes[left_index] - block, null_col_id = block.create_constant(None, dtype=dtype) - binop_result_ids.append(null_col_id) + left_col_id = self._block.value_columns[left_index] + block, result_col_id = block.apply_unary_op( + get_column_left(left_col_id), + ops.partial_right(op, None), + ) + binop_result_ids.append(result_col_id) elif right_index >= 0: - dtype = other.dtypes[right_index] - block, null_col_id = block.create_constant(None, dtype=dtype) - binop_result_ids.append(null_col_id) + right_col_id = other._block.value_columns[right_index] + block, result_col_id = block.apply_unary_op( + get_column_right(right_col_id), + ops.partial_left(op, None), + ) + binop_result_ids.append(result_col_id) else: # Should not be possible raise ValueError("No right or left index.") @@ -759,6 +772,75 @@ def rpow( __rpow__ = rpow + def combine( + self, + other: DataFrame, + func: typing.Callable[ + [bigframes.series.Series, bigframes.series.Series], bigframes.series.Series + ], + fill_value=None, + overwrite: bool = True, + ) -> DataFrame: + # Join rows + joined_index, (get_column_left, get_column_right) = self._block.index.join( + other._block.index, how="outer" + ) + columns, lcol_indexer, rcol_indexer = self.columns.join( + other.columns, how="outer", return_indexers=True + ) + + column_indices = zip( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)), + rcol_indexer if (lcol_indexer is not None) else range(len(columns)), + ) + + block = joined_index._block + results = [] + for left_index, right_index in column_indices: + if left_index >= 0 and right_index >= 0: # -1 indices indicate missing + left_col_id = get_column_left(self._block.value_columns[left_index]) + right_col_id = get_column_right(other._block.value_columns[right_index]) + left_series = bigframes.series.Series(block.select_column(left_col_id)) + right_series = bigframes.series.Series( + block.select_column(right_col_id) + ) + if fill_value is not None: + left_series = left_series.fillna(fill_value) + right_series = right_series.fillna(fill_value) + results.append(func(left_series, right_series)) + elif left_index >= 0: + # Does not exist in other + if overwrite: + dtype = self.dtypes[left_index] + block, null_col_id = block.create_constant(None, dtype=dtype) + result = bigframes.series.Series(block.select_column(null_col_id)) + results.append(result) + else: + left_col_id = get_column_left(self._block.value_columns[left_index]) + result = bigframes.series.Series(block.select_column(left_col_id)) + if fill_value is not None: + result = result.fillna(fill_value) + results.append(result) + elif right_index >= 0: + right_col_id = get_column_right(other._block.value_columns[right_index]) + result = bigframes.series.Series(block.select_column(right_col_id)) + if fill_value is not None: + result = result.fillna(fill_value) + results.append(result) + else: + # Should not be possible + raise ValueError("No right or left index.") + + if all([isinstance(val, bigframes.series.Series) for val in results]): + import bigframes.core.reshape as rs + + return rs.concat(results, axis=1) + else: + raise ValueError("'func' must return Series") + + def combine_first(self, other: DataFrame): + return self._apply_dataframe_binop(other, ops.fillna_op) + def to_pandas( self, max_download_size: Optional[int] = None, @@ -810,6 +892,28 @@ def head(self, n: int = 5) -> DataFrame: def tail(self, n: int = 5) -> DataFrame: return typing.cast(DataFrame, self.iloc[-n:]) + def nlargest( + self, + n: int, + columns: typing.Union[blocks.Label, typing.Sequence[blocks.Label]], + keep: str = "first", + ) -> DataFrame: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + column_ids = self._sql_names(columns) + return DataFrame(block_ops.nlargest(self._block, n, column_ids, keep=keep)) + + def nsmallest( + self, + n: int, + columns: typing.Union[blocks.Label, typing.Sequence[blocks.Label]], + keep: str = "first", + ) -> DataFrame: + if keep not in ("first", "last", "all"): + raise ValueError("'keep must be one of 'first', 'last', or 'all'") + column_ids = self._sql_names(columns) + return DataFrame(block_ops.nsmallest(self._block, n, column_ids, keep=keep)) + def drop( self, labels: typing.Any = None, @@ -852,13 +956,50 @@ def drop( raise ValueError("Must specify 'labels' or 'index'/'columns") return DataFrame(block) - def droplevel(self, level: LevelsType): - resolved_level_ids = self._resolve_levels(level) - return DataFrame(self._block.drop_levels(resolved_level_ids)) + def droplevel(self, level: LevelsType, axis: int | str = 0): + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + resolved_level_ids = self._resolve_levels(level) + return DataFrame(self._block.drop_levels(resolved_level_ids)) + else: + if isinstance(self.columns, pandas.MultiIndex): + new_df = self.copy() + new_df.columns = self.columns.droplevel(level) + return new_df + else: + raise ValueError("Columns must be a multiindex to drop levels.") + + def swaplevel(self, i: int = -2, j: int = -1, axis: int | str = 0): + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + level_i = self._block.index_columns[i] + level_j = self._block.index_columns[j] + mapping = {level_i: level_j, level_j: level_i} + reordering = [ + mapping.get(index_id, index_id) + for index_id in self._block.index_columns + ] + return DataFrame(self._block.reorder_levels(reordering)) + else: + if isinstance(self.columns, pandas.MultiIndex): + new_df = self.copy() + new_df.columns = self.columns.swaplevel(i, j) + return new_df + else: + raise ValueError("Columns must be a multiindex to reorder levels.") - def reorder_levels(self, order: LevelsType): - resolved_level_ids = self._resolve_levels(order) - return DataFrame(self._block.reorder_levels(resolved_level_ids)) + def reorder_levels(self, order: LevelsType, axis: int | str = 0): + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + resolved_level_ids = self._resolve_levels(order) + return DataFrame(self._block.reorder_levels(resolved_level_ids)) + else: + if isinstance(self.columns, pandas.MultiIndex): + new_df = self.copy() + new_df.columns = self.columns.reorder_levels(order) + return new_df + else: + raise ValueError("Columns must be a multiindex to reorder levels.") def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]: if utils.is_list_like(level): @@ -1096,8 +1237,177 @@ def add_suffix(self, suffix: str, axis: int | str | None = None) -> DataFrame: axis = 1 if axis is None else axis return DataFrame(self._get_block().add_suffix(suffix, axis)) + def filter( + self, + items: typing.Optional[typing.Iterable] = None, + like: typing.Optional[str] = None, + regex: typing.Optional[str] = None, + axis: int | str | None = None, + ) -> DataFrame: + if sum([(items is not None), (like is not None), (regex is not None)]) != 1: + raise ValueError( + "Need to provide exactly one of 'items', 'like', or 'regex'" + ) + axis_n = utils.get_axis_number(axis) if (axis is not None) else 1 + if axis_n == 0: # row labels + return self._filter_rows(items, like, regex) + else: # column labels + return self._filter_columns(items, like, regex) + + def _filter_rows( + self, + items: typing.Optional[typing.Iterable] = None, + like: typing.Optional[str] = None, + regex: typing.Optional[str] = None, + ) -> DataFrame: + if len(self._block.index_columns) > 1: + raise NotImplementedError( + "Method filter does not support rows multiindex. {constants.FEEDBACK_LINK}" + ) + if (like is not None) or (regex is not None): + block = self._block + block, label_string_id = block.apply_unary_op( + self._block.index_columns[0], + ops.AsTypeOp(pandas.StringDtype(storage="pyarrow")), + ) + if like is not None: + block, mask_id = block.apply_unary_op( + label_string_id, ops.ContainsStringOp(pat=like) + ) + else: # regex + assert regex is not None + block, mask_id = block.apply_unary_op( + label_string_id, ops.ContainsRegexOp(pat=regex) + ) + + block = block.filter(mask_id) + block = block.select_columns(self._block.value_columns) + return DataFrame(block) + elif items is not None: + # Behavior matches pandas 2.1+, older pandas versions would reindex + block = self._block + block, mask_id = block.apply_unary_op( + self._block.index_columns[0], ops.IsInOp(values=list(items)) + ) + block = block.filter(mask_id) + block = block.select_columns(self._block.value_columns) + return DataFrame(block) + else: + raise ValueError("Need to provide 'items', 'like', or 'regex'") + + def _filter_columns( + self, + items: typing.Optional[typing.Iterable] = None, + like: typing.Optional[str] = None, + regex: typing.Optional[str] = None, + ) -> DataFrame: + if (like is not None) or (regex is not None): + + def label_filter(label): + label_str = label if isinstance(label, str) else str(label) + if like: + return like in label_str + else: # regex + return re.match(regex, label_str) is not None + + cols = [ + col_id + for col_id, label in zip(self._block.value_columns, self.columns) + if label_filter(label) + ] + return DataFrame(self._block.select_columns(cols)) + if items is not None: + # Behavior matches pandas 2.1+, older pandas versions would reorder using order of items + new_columns = self.columns.intersection(pandas.Index(items)) + return self.reindex(columns=new_columns) + else: + raise ValueError("Need to provide 'items', 'like', or 'regex'") + + def reindex( + self, + labels=None, + *, + index=None, + columns=None, + axis: typing.Optional[typing.Union[str, int]] = None, + validate: typing.Optional[bool] = None, + ): + if labels: + if index or columns: + raise ValueError("Cannot specify both 'labels' and 'index'/'columns") + axis_n = utils.get_axis_number(axis) if (axis is not None) else 0 + if axis_n == 0: + index = labels + else: + columns = labels + if (index is not None) and (columns is not None): + return self._reindex_columns(columns)._reindex_rows( + index, validate=validate or False + ) + if index is not None: + return self._reindex_rows(index, validate=validate or False) + if columns is not None: + return self._reindex_columns(columns) + + def _reindex_rows( + self, + index, + *, + validate: typing.Optional[bool] = None, + ): + if validate and not self.index.is_unique: + raise ValueError("Original index must be unique to reindex") + keep_original_names = False + if isinstance(index, indexes.Index): + new_indexer = DataFrame(data=index._data._get_block())[[]] + else: + if not isinstance(index, pandas.Index): + keep_original_names = True + index = pandas.Index(index) + if index.nlevels != self.index.nlevels: + raise NotImplementedError( + "Cannot reindex with index with different nlevels" + ) + new_indexer = DataFrame(index=index)[[]] + # multiindex join is senstive to index names, so we will set all these + result = new_indexer.rename_axis(range(new_indexer.index.nlevels)).join( + self.rename_axis(range(self.index.nlevels)), + how="left", + ) + # and then reset the names after the join + return result.rename_axis( + self.index.names if keep_original_names else index.names + ) + + def _reindex_columns(self, columns): + block = self._block + new_column_index, indexer = self.columns.reindex(columns) + result_cols = [] + for label, index in zip(columns, indexer): + if index >= 0: + result_cols.append(self._block.value_columns[index]) + else: + block, null_col = block.create_constant( + pandas.NA, label, dtype=pandas.Float64Dtype() + ) + result_cols.append(null_col) + result_df = DataFrame(block.select_columns(result_cols)) + result_df.columns = new_column_index + return result_df + + def reindex_like(self, other: DataFrame, *, validate: typing.Optional[bool] = None): + return self.reindex(index=other.index, columns=other.columns, validate=validate) + def fillna(self, value=None) -> DataFrame: - return self._apply_binop(value, ops.fillna_op) + return self._apply_binop(value, ops.fillna_op, how="left") + + def ffill(self, *, limit: typing.Optional[int] = None) -> DataFrame: + window = bigframes.core.WindowSpec(preceding=limit, following=0) + return self._apply_window_op(agg_ops.LastNonNullOp(), window) + + def bfill(self, *, limit: typing.Optional[int] = None) -> DataFrame: + window = bigframes.core.WindowSpec(preceding=0, following=limit) + return self._apply_window_op(agg_ops.FirstNonNullOp(), window) def isin(self, values) -> DataFrame: if utils.is_dict_like(values): @@ -1309,6 +1619,14 @@ def describe(self) -> DataFrame: ) return typing.cast(DataFrame, result) + def skew(self, *, numeric_only: bool = False): + if not numeric_only: + frame = self._raise_on_non_numeric("skew") + else: + frame = self._drop_non_numeric() + result_block = block_ops.skew(frame._block, frame._block.value_columns) + return bigframes.series.Series(result_block) + def pivot( self, *, @@ -1702,17 +2020,29 @@ def shift(self, periods: int = 1) -> DataFrame: ) return self._apply_window_op(agg_ops.ShiftOp(periods), window) + def diff(self, periods: int = 1) -> DataFrame: + window = bigframes.core.WindowSpec( + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.DiffOp(periods), window) + + def pct_change(self, periods: int = 1) -> DataFrame: + # Future versions of pandas will not perfrom ffill automatically + df = self.ffill() + return DataFrame(block_ops.pct_change(df._block, periods=periods)) + def _apply_window_op( self, op: agg_ops.WindowOp, window_spec: bigframes.core.WindowSpec, ): - block = self._block.multi_apply_window_op( + block, result_ids = self._block.multi_apply_window_op( self._block.value_columns, op, window_spec=window_spec, ) - return DataFrame(block) + return DataFrame(block.select_columns(result_ids)) def sample( self, @@ -1875,6 +2205,98 @@ def to_parquet(self, path: str, *, index: bool = True) -> None: _, query_job = self._block.expr._session._start_query(export_data_statement) self._set_internal_query_job(query_job) + def to_dict( + self, + orient: Literal[ + "dict", "list", "series", "split", "tight", "records", "index" + ] = "dict", + into: type[dict] = dict, + **kwargs, + ) -> dict | list[dict]: + return self.to_pandas().to_dict(orient, into, **kwargs) # type: ignore + + def to_excel(self, excel_writer, sheet_name: str = "Sheet1", **kwargs) -> None: + return self.to_pandas().to_excel(excel_writer, sheet_name, **kwargs) + + def to_latex( + self, + buf=None, + columns: Sequence | None = None, + header: bool | Sequence[str] = True, + index: bool = True, + **kwargs, + ) -> str | None: + return self.to_pandas().to_latex( + buf, columns=columns, header=header, index=index, **kwargs # type: ignore + ) + + def to_records( + self, index: bool = True, column_dtypes=None, index_dtypes=None + ) -> numpy.recarray: + return self.to_pandas().to_records(index, column_dtypes, index_dtypes) + + def to_string( + self, + buf=None, + columns: Sequence[str] | None = None, + col_space=None, + header: bool | Sequence[str] = True, + index: bool = True, + na_rep: str = "NaN", + formatters=None, + float_format=None, + sparsify: bool | None = None, + index_names: bool = True, + justify: str | None = None, + max_rows: int | None = None, + max_cols: int | None = None, + show_dimensions: bool = False, + decimal: str = ".", + line_width: int | None = None, + min_rows: int | None = None, + max_colwidth: int | None = None, + encoding: str | None = None, + ) -> str | None: + return self.to_pandas().to_string( + buf, + columns, # type: ignore + col_space, + header, # type: ignore + index, + na_rep, + formatters, + float_format, + sparsify, + index_names, + justify, + max_rows, + max_cols, + show_dimensions, + decimal, + line_width, + min_rows, + max_colwidth, + encoding, + ) + + def to_markdown( + self, + buf=None, + mode: str = "wt", + index: bool = True, + **kwargs, + ) -> str | None: + return self.to_pandas().to_markdown(buf, mode, index, **kwargs) # type: ignore + + def to_pickle(self, path, **kwargs) -> None: + return self.to_pandas().to_pickle(path, **kwargs) + + def to_orc(self, path=None, **kwargs) -> bytes | None: + as_pandas = self.to_pandas() + # to_orc only works with default index + as_pandas_default_index = as_pandas.reset_index() + return as_pandas_default_index.to_orc(path, **kwargs) + def _apply_unary_op(self, operation: ops.UnaryOp) -> DataFrame: block = self._block.multi_apply_unary_op(self._block.value_columns, operation) return DataFrame(block) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 95cf737b2e7..af3209b0e1c 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -157,7 +157,7 @@ def ibis_dtype_to_bigframes_dtype( return IBIS_TO_BIGFRAMES[ibis_dtype] else: raise ValueError( - f"Unexpected Ibis data type {type(ibis_dtype)}. {constants.FEEDBACK_LINK}" + f"Unexpected Ibis data type {ibis_dtype}. {constants.FEEDBACK_LINK}" ) diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py index 9f9d9f85d0e..f07274f8fcb 100644 --- a/bigframes/ml/base.py +++ b/bigframes/ml/base.py @@ -133,7 +133,7 @@ class TrainablePredictor(Predictor): Also the predictor can be attached to a pipeline with transformers.""" @abc.abstractmethod - def fit(self, X, y, transforms): + def _fit(self, X, y, transforms=None): pass @abc.abstractmethod @@ -146,6 +146,36 @@ def to_gbq(self, model_name, replace): pass +class SupervisedTrainablePredictor(TrainablePredictor): + """A BigQuery DataFrames ML Supervised Model base class that can be used to fit and predict outputs. + + Need to provide both X and y in supervised tasks.""" + + _T = TypeVar("_T", bound="SupervisedTrainablePredictor") + + def fit( + self: _T, + X: Union[bpd.DataFrame, bpd.Series], + y: Union[bpd.DataFrame, bpd.Series], + ) -> _T: + return self._fit(X, y) + + +class UnsupervisedTrainablePredictor(TrainablePredictor): + """A BigQuery DataFrames ML Unsupervised Model base class that can be used to fit and predict outputs. + + Only need to provide both X (y is optional and ignored) in unsupervised tasks.""" + + _T = TypeVar("_T", bound="UnsupervisedTrainablePredictor") + + def fit( + self: _T, + X: Union[bpd.DataFrame, bpd.Series], + y: Optional[Union[bpd.DataFrame, bpd.Series]] = None, + ) -> _T: + return self._fit(X, y) + + class Transformer(BaseEstimator): """A BigQuery DataFrames Transformer base class that transforms data. diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 2501d2b21f1..14cce2879e9 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -28,13 +28,13 @@ class KMeans( + base.UnsupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans, - base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans.__doc__ - def __init__(self, n_clusters=8): + def __init__(self, n_clusters: int = 8): self.n_clusters = n_clusters self._bqml_model: Optional[core.BqmlModel] = None @@ -58,7 +58,7 @@ def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: """The model options as they will be set for BQML""" return {"model_type": "KMEANS", "num_clusters": self.n_clusters} - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y=None, # ignored diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index df01303ffac..a3d3503ad0a 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -33,8 +33,8 @@ class ColumnTransformer( - third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, base.Transformer, + third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, ): __doc__ = ( third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer.__doc__ diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 75b57f2e54f..0cfe3b3ddf9 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -28,12 +28,12 @@ class PCA( + base.UnsupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.decomposition._pca.PCA, - base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.sklearn.decomposition._pca.PCA.__doc__ - def __init__(self, n_components=3): + def __init__(self, n_components: int = 3): self.n_components = n_components self._bqml_model: Optional[core.BqmlModel] = None @@ -52,7 +52,7 @@ def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> PCA: new_pca._bqml_model = core.BqmlModel(session, model) return new_pca - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y=None, diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index 56a0cc3d949..142edaa00f6 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -48,8 +48,8 @@ class XGBRegressor( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor, - base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor.__doc__ @@ -57,22 +57,22 @@ def __init__( self, num_parallel_tree: int = 1, booster: Literal["gbtree", "dart"] = "gbtree", - dart_normalized_type: Literal["TREE", "FOREST"] = "TREE", + dart_normalized_type: Literal["tree", "forest"] = "tree", tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", min_tree_child_weight: int = 1, - colsample_bytree=1.0, - colsample_bylevel=1.0, - colsample_bynode=1.0, - gamma=0.0, + colsample_bytree: float = 1.0, + colsample_bylevel: float = 1.0, + colsample_bynode: float = 1.0, + gamma: float = 0.0, max_depth: int = 6, - subsample=1.0, - reg_alpha=0.0, - reg_lambda=1.0, - early_stop=True, - learning_rate=0.3, + subsample: float = 1.0, + reg_alpha: float = 0.0, + reg_lambda: float = 1.0, + early_stop: float = True, + learning_rate: float = 0.3, max_iterations: int = 20, - min_rel_progress=0.01, - enable_global_explain=False, + min_rel_progress: float = 0.01, + enable_global_explain: bool = False, xgboost_version: Literal["0.9", "1.1"] = "0.9", ): self.num_parallel_tree = num_parallel_tree @@ -143,7 +143,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: "xgboost_version": self.xgboost_version, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], @@ -211,8 +211,8 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor: class XGBClassifier( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier, - base.TrainablePredictor, ): __doc__ = third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier.__doc__ @@ -221,22 +221,22 @@ def __init__( self, num_parallel_tree: int = 1, booster: Literal["gbtree", "dart"] = "gbtree", - dart_normalized_type: Literal["TREE", "FOREST"] = "TREE", + dart_normalized_type: Literal["tree", "forest"] = "tree", tree_method: Literal["auto", "exact", "approx", "hist"] = "auto", min_tree_child_weight: int = 1, - colsample_bytree=1.0, - colsample_bylevel=1.0, - colsample_bynode=1.0, - gamma=0.0, + colsample_bytree: float = 1.0, + colsample_bylevel: float = 1.0, + colsample_bynode: float = 1.0, + gamma: float = 0.0, max_depth: int = 6, - subsample=1.0, - reg_alpha=0.0, - reg_lambda=1.0, - early_stop=True, - learning_rate=0.3, + subsample: float = 1.0, + reg_alpha: float = 0.0, + reg_lambda: float = 1.0, + early_stop: bool = True, + learning_rate: float = 0.3, max_iterations: int = 20, - min_rel_progress=0.01, - enable_global_explain=False, + min_rel_progress: float = 0.01, + enable_global_explain: bool = False, xgboost_version: Literal["0.9", "1.1"] = "0.9", ): self.num_parallel_tree = num_parallel_tree @@ -307,7 +307,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: "xgboost_version": self.xgboost_version, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], @@ -374,8 +374,8 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier: class RandomForestRegressor( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor, - base.TrainablePredictor, ): __doc__ = ( @@ -461,7 +461,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: "xgboost_version": self.xgboost_version, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], @@ -542,8 +542,8 @@ def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegresso class RandomForestClassifier( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier, - base.TrainablePredictor, ): __doc__ = ( @@ -629,7 +629,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: "xgboost_version": self.xgboost_version, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index b7e0553ecb7..22d81294fc3 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -27,7 +27,7 @@ _PREDICT_OUTPUT_COLUMNS = ["forecast_timestamp", "forecast_value"] -class ARIMAPlus(base.TrainablePredictor): +class ARIMAPlus(base.SupervisedTrainablePredictor): """Time Series ARIMA Plus model.""" def __init__(self): @@ -48,7 +48,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: """The model options as they will be set for BQML.""" return {"model_type": "ARIMA_PLUS"} - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index 0b18db93154..1606a15d730 100644 --- a/bigframes/ml/linear_model.py +++ b/bigframes/ml/linear_model.py @@ -30,8 +30,8 @@ class LinearRegression( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression, - base.TrainablePredictor, ): __doc__ = ( third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression.__doc__ @@ -39,7 +39,7 @@ class LinearRegression( def __init__( self, - fit_intercept=True, + fit_intercept: bool = True, ): self.fit_intercept = fit_intercept self._bqml_model: Optional[core.BqmlModel] = None @@ -71,7 +71,7 @@ def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: "fit_intercept": self.fit_intercept, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], @@ -136,8 +136,8 @@ def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression: class LogisticRegression( + base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression, - base.TrainablePredictor, ): __doc__ = ( third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression.__doc__ @@ -189,12 +189,13 @@ def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: # "class_weights": self.class_weights, } - def fit( + def _fit( self, X: Union[bpd.DataFrame, bpd.Series], y: Union[bpd.DataFrame, bpd.Series], transforms: Optional[List[str]] = None, ) -> LogisticRegression: + """Fit model with transforms.""" X, y = utils.convert_to_dataframe(X, y) self._bqml_model = core.create_bqml_model( diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 345e3deb728..973fbf2ad96 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -100,26 +100,26 @@ def predict( Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that expect a true or correct response, while higher temperatures can lead to more diverse or unexpected results. A temperature of 0 is deterministic: the highest probability token is always selected. For most use cases, try starting with a temperature of 0.2. - Default 0. + Default 0. Possible values [0.0, 1.0]. max_output_tokens (int, default 128): Maximum number of tokens that can be generated in the response. Specify a lower value for shorter responses and a higher value for longer responses. A token may be smaller than a word. A token is approximately four characters. 100 tokens correspond to roughly 60-80 words. - Default 128. + Default 128. Possible values [1, 1024]. top_k (int, default 40): Top-k changes how the model selects tokens for output. A top-k of 1 means the selected token is the most probable among all tokens in the model’s vocabulary (also called greedy decoding), while a top-k of 3 means that the next token is selected from among the 3 most probable tokens (using temperature). For each token selection step, the top K tokens with the highest probabilities are sampled. Then tokens are further filtered based on topP with the final token selected using temperature sampling. Specify a lower value for less random responses and a higher value for more random responses. - Default 40. + Default 40. Possible values [1, 40]. top_p (float, default 0.95):: Top-p changes how the model selects tokens for output. Tokens are selected from most K (see topK parameter) probable to least until the sum of their probabilities equals the top-p value. For example, if tokens A, B, and C have a probability of 0.3, 0.2, and 0.1 and the top-p value is 0.5, then the model will select either A or B as the next token (using temperature) and not consider C at all. Specify a lower value for less random responses and a higher value for more random responses. - Default 0.95. + Default 0.95. Possible values [0.0, 1.0]. Returns: diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index bfd03925265..bff0bf36ad4 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -24,14 +24,14 @@ import bigframes import bigframes.constants as constants -from bigframes.ml import base, compose, loader, preprocessing, utils +from bigframes.ml import base, compose, forecasting, loader, preprocessing, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.pipeline class Pipeline( - third_party.bigframes_vendored.sklearn.pipeline.Pipeline, base.BaseEstimator, + third_party.bigframes_vendored.sklearn.pipeline.Pipeline, ): __doc__ = third_party.bigframes_vendored.sklearn.pipeline.Pipeline.__doc__ @@ -55,7 +55,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): self._transform = transform else: raise NotImplementedError( - f"Transform {transform} is not yet supported by Pipeline. {constants.FEEDBACK_LINK}" + f"Transformer type {type(transform)} is not yet supported by Pipeline. {constants.FEEDBACK_LINK}" ) if not isinstance( @@ -63,7 +63,13 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): base.TrainablePredictor, ): raise NotImplementedError( - f"Estimator {estimator} is not supported by Pipeline. {constants.FEEDBACK_LINK}" + f"Estimator type {type(estimator)} is not supported by Pipeline. {constants.FEEDBACK_LINK}" + ) + + # BQML doesn't support ARIMA_PLUS with transformers. b/298676367 + if isinstance(estimator, forecasting.ARIMAPlus): + raise NotImplementedError( + f"Estimator type {type(estimator)} is not supported by Pipeline. {constants.FEEDBACK_LINK}" ) self._transform = transform @@ -92,7 +98,7 @@ def fit( (y,) = utils.convert_to_dataframe(y) transform_sqls.extend(y.columns.tolist()) - self._estimator.fit(X=X, y=y, transforms=transform_sqls) + self._estimator._fit(X=X, y=y, transforms=transform_sqls) return self def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index ee46a370520..8add7bdd767 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -29,8 +29,8 @@ class StandardScaler( - third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, ): __doc__ = ( third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler.__doc__ @@ -105,8 +105,8 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: class OneHotEncoder( - third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, ): # BQML max value https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder#syntax TOP_K_DEFAULT = 1000000 diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py index 874c264194c..23271e82205 100644 --- a/bigframes/operations/aggregations.py +++ b/bigframes/operations/aggregations.py @@ -19,8 +19,10 @@ import ibis import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types +from pandas import Int64Dtype import bigframes.constants as constants +import bigframes.dtypes as dtypes import third_party.bigframes_vendored.ibis.expr.operations as vendored_ibis_ops @@ -202,6 +204,16 @@ def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: ) +class PopVarOp(AggregateOp): + name = "popvar" + + @numeric_op + def _as_ibis(self, x: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present( + typing.cast(ibis_types.NumericColumn, x).var(how="pop"), window + ) + + class CountOp(AggregateOp): name = "count" @@ -217,16 +229,20 @@ def skips_nulls(self): class CutOp(WindowOp): def __init__(self, bins: int): - self._bins = bins + self._bins_ibis = dtypes.literal_to_ibis_scalar(bins, force_dtype=Int64Dtype()) + self._bins_int = bins def _as_ibis(self, x: ibis_types.Column, window=None): col_min = _apply_window_if_present(x.min(), window) col_max = _apply_window_if_present(x.max(), window) - bin_width = (col_max - col_min) / self._bins + bin_width = (col_max - col_min) / self._bins_ibis out = ibis.case() - for bin in range(self._bins - 1): - out = out.when(x <= (col_min + (bin + 1) * bin_width), bin) - out = out.when(x.notnull(), self._bins - 1) + for this_bin in range(self._bins_int - 1): + out = out.when( + x <= (col_min + (this_bin + 1) * bin_width), + dtypes.literal_to_ibis_scalar(this_bin, force_dtype=Int64Dtype()), + ) + out = out.when(x.notnull(), self._bins_ibis - 1) return out.end() @property @@ -305,6 +321,28 @@ def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: return _apply_window_if_present(column.first(), window) +class FirstNonNullOp(WindowOp): + @property + def skips_nulls(self): + return False + + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present( + vendored_ibis_ops.FirstNonNullValue(column).to_expr(), window # type: ignore + ) + + +class LastNonNullOp(WindowOp): + @property + def skips_nulls(self): + return False + + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + return _apply_window_if_present( + vendored_ibis_ops.LastNonNullValue(column).to_expr(), window # type: ignore + ) + + class ShiftOp(WindowOp): def __init__(self, periods: int): self._periods = periods @@ -321,6 +359,28 @@ def skips_nulls(self): return False +class DiffOp(WindowOp): + def __init__(self, periods: int): + self._periods = periods + + def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value: + shifted = ShiftOp(self._periods)._as_ibis(column, window) + if column.type().is_boolean(): + return typing.cast(ibis_types.BooleanColumn, column) != typing.cast( + ibis_types.BooleanColumn, shifted + ) + elif column.type().is_numeric(): + return typing.cast(ibis_types.NumericColumn, column) - typing.cast( + ibis_types.NumericColumn, shifted + ) + else: + raise TypeError(f"Cannot perform diff on type{column.type()}") + + @property + def skips_nulls(self): + return False + + class AllOp(AggregateOp): def _as_ibis( self, column: ibis_types.Column, window=None diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index e1a23e67a1c..cb278345907 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -399,6 +399,7 @@ def remote_function( dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, + name: Optional[str] = None, ): return global_session.with_default_session( bigframes.session.Session.remote_function, @@ -407,6 +408,7 @@ def remote_function( dataset=dataset, bigquery_connection=bigquery_connection, reuse=reuse, + name=name, ) diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 6932e5b5805..6fc2f8e59fb 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -28,6 +28,8 @@ import textwrap from typing import List, NamedTuple, Optional, Sequence, TYPE_CHECKING +import requests + if TYPE_CHECKING: from bigframes.session import Session @@ -99,7 +101,7 @@ def get_remote_function_locations(bq_location): def _get_hash(def_): - "Get hash of a function." + "Get hash (32 digits alphanumeric) of a function." def_repr = cloudpickle.dumps(def_, protocol=_pickle_protocol_version) return hashlib.md5(def_repr).hexdigest() @@ -128,7 +130,7 @@ class IbisSignature(NamedTuple): def get_cloud_function_name(def_, uniq_suffix=None): - """Get the name of the cloud function.""" + "Get a name for the cloud function for the given user defined function." cf_name = _get_hash(def_) cf_name = f"bigframes-{cf_name}" # for identification if uniq_suffix: @@ -137,7 +139,7 @@ def get_cloud_function_name(def_, uniq_suffix=None): def get_remote_function_name(def_, uniq_suffix=None): - """Get the name for the BQ remote function.""" + "Get a name for the BQ remote function for the given user defined function." bq_rf_name = _get_hash(def_) bq_rf_name = f"bigframes_{bq_rf_name}" # for identification if uniq_suffix: @@ -206,9 +208,15 @@ def create_bq_remote_function( query_job.result() # Wait for the job to complete. logger.info(f"Created remote function {query_job.ddl_target_routine}") + def get_cloud_function_fully_qualified_parent(self): + "Get the fully qualilfied parent for a cloud function." + return self._cloud_functions_client.common_location_path( + self._gcp_project_id, self._cloud_function_region + ) + def get_cloud_function_fully_qualified_name(self, name): "Get the fully qualilfied name for a cloud function." - return "projects/{}/locations/{}/functions/{}".format( + return self._cloud_functions_client.function_path( self._gcp_project_id, self._cloud_function_region, name ) @@ -319,6 +327,7 @@ def create_cloud_function(self, def_, cf_name): # Build and deploy folder structure containing cloud function with tempfile.TemporaryDirectory() as dir: entry_point = self.generate_cloud_function_code(def_, dir) + archive_path = shutil.make_archive(dir, "zip", dir) # We are creating cloud function source code from the currently running # python version. Use the same version to deploy. This is necessary @@ -331,50 +340,56 @@ def create_cloud_function(self, def_, cf_name): sys.version_info.major, sys.version_info.minor ) - # deploy/redeploy the cloud function - # TODO(shobs): Figure out a way to skip this step if a cloud function - # already exists with the same name and source code - command = ( - "gcloud functions deploy" - + f" {cf_name} --gen2" - + f" --runtime={python_version}" - + f" --project={self._gcp_project_id}" - + f" --region={self._cloud_function_region}" - + f" --source={dir}" - + f" --entry-point={entry_point}" - + " --trigger-http" + # Determine an upload URL for user code + upload_url_request = functions_v2.GenerateUploadUrlRequest() + upload_url_request.parent = self.get_cloud_function_fully_qualified_parent() + upload_url_response = self._cloud_functions_client.generate_upload_url( + request=upload_url_request ) - # If the cloud function is being created for the first time, then let's - # make it not allow unauthenticated calls. If it was previously created - # then this invocation will update it, in which case do not touch that - # aspect and let the previous policy hold. The reason we do this is to - # avoid an IAM permission needed to update the invocation policy. - # For example, when a cloud function is being created for the first - # time, i.e. - # $ gcloud functions deploy python-foo-http --gen2 --runtime=python310 - # --region=us-central1 - # --source=/source/code/dir - # --entry-point=foo_http - # --trigger-http - # --no-allow-unauthenticated - # It works. When an invocation of the same command is done for the - # second time, it may run into an error like: - # ERROR: (gcloud.functions.deploy) PERMISSION_DENIED: Permission - # 'run.services.setIamPolicy' denied on resource - # 'projects/my_project/locations/us-central1/services/python-foo-http' (or resource may not exist) - # But when --no-allow-unauthenticated is omitted then it goes through. - # It suggests that in the second invocation the command is trying to set - # the IAM policy of the service, and the user running BigQuery - # DataFrame may not have privilege to do so, so better avoid this - # if we can. - if self.get_cloud_function_endpoint(cf_name): - logger.info(f"Updating existing cloud function: {command}") - else: - command = f"{command} --no-allow-unauthenticated" - logger.info(f"Creating new cloud function: {command}") - - _run_system_command(command) + # Upload the code to GCS + with open(archive_path, "rb") as f: + response = requests.put( + upload_url_response.upload_url, + data=f, + headers={"content-type": "application/zip"}, + ) + if response.status_code != 200: + raise RuntimeError( + "Failed to upload user code. code={}, reason={}, text={}".format( + response.status_code, response.reason, response.text + ) + ) + + # Deploy Cloud Function + create_function_request = functions_v2.CreateFunctionRequest() + create_function_request.parent = ( + self.get_cloud_function_fully_qualified_parent() + ) + create_function_request.function_id = cf_name + function = functions_v2.Function() + function.name = self.get_cloud_function_fully_qualified_name(cf_name) + function.build_config = functions_v2.BuildConfig() + function.build_config.runtime = python_version + function.build_config.entry_point = entry_point + function.build_config.source = functions_v2.Source() + function.build_config.source.storage_source = functions_v2.StorageSource() + function.build_config.source.storage_source.bucket = ( + upload_url_response.storage_source.bucket + ) + function.build_config.source.storage_source.object_ = ( + upload_url_response.storage_source.object_ + ) + create_function_request.function = function + + # Create the cloud function and wait for it to be ready to use + operation = self._cloud_functions_client.create_function( + request=create_function_request + ) + operation.result() + + # Cleanup + os.remove(archive_path) # Fetch the endpoint of the just created function endpoint = self.get_cloud_function_endpoint(cf_name) @@ -389,23 +404,47 @@ def create_cloud_function(self, def_, cf_name): return endpoint def provision_bq_remote_function( - self, def_, input_types, output_type, uniq_suffix=None + self, + def_, + input_types, + output_type, + reuse, + name, ): """Provision a BigQuery remote function.""" - # Derive the name of the underlying cloud function and first create - # it if it does not exist + # If reuse of any existing function with the same name (indicated by the + # same hash of its source code) is not intended, then attach a unique + # suffix to the intended function name to make it unique. + uniq_suffix = None + if not reuse: + uniq_suffix = "".join( + random.choices(string.ascii_lowercase + string.digits, k=8) + ) + + # Derive the name of the cloud function underlying the intended BQ + # remote function cloud_function_name = get_cloud_function_name(def_, uniq_suffix) cf_endpoint = self.get_cloud_function_endpoint(cloud_function_name) + + # Create the cloud function if it does not exist if not cf_endpoint: - self.check_cloud_function_tools_and_permissions() cf_endpoint = self.create_cloud_function(def_, cloud_function_name) else: logger.info(f"Cloud function {cloud_function_name} already exists.") - # Derive the name of the remote function and create/replace it if needed - remote_function_name = get_remote_function_name(def_, uniq_suffix) + # Derive the name of the remote function + remote_function_name = name + if not remote_function_name: + remote_function_name = get_remote_function_name(def_, uniq_suffix) rf_endpoint, rf_conn = self.get_remote_function_specs(remote_function_name) - if rf_endpoint != cf_endpoint or rf_conn != self._bq_connection_id: + + # Create the BQ remote function in following circumstances: + # 1. It does not exist + # 2. It exists but the existing remote function has different + # configuration than intended + if not rf_endpoint or ( + rf_endpoint != cf_endpoint or rf_conn != self._bq_connection_id + ): input_args = inspect.getargs(def_.__code__).args if len(input_args) != len(input_types): raise ValueError( @@ -439,27 +478,6 @@ def get_remote_function_specs(self, remote_function_name): break return (http_endpoint, bq_connection) - def check_cloud_function_tools_and_permissions(self): - """Check if the necessary tools and permissions are in place for creating remote function""" - # gcloud CLI comes with bq CLI and they are required for creating google - # cloud function and BigQuery remote function respectively - if not shutil.which("gcloud"): - raise ValueError( - "gcloud tool not installed, install it from https://cloud.google.com/sdk/docs/install. " - f"{constants.FEEDBACK_LINK}" - ) - - # TODO(shobs): Check for permissions too - # I (shobs) tried the following method - # $ gcloud asset search-all-iam-policies \ - # --format=json \ - # --scope=projects/{gcp_project_id} \ - # --query='policy.role.permissions:cloudfunctions.functions.create' - # as a proxy to all the privilges necessary to create cloud function - # https://cloud.google.com/functions/docs/reference/iam/roles#cloudfunctions.developer - # but that itself required the runner to have the permission to enable - # `cloudasset.googleapis.com` - def remote_function_node( routine_ref: bigquery.RoutineReference, ibis_signature: IbisSignature @@ -583,6 +601,7 @@ def remote_function( dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, + name: Optional[str] = None, ): """Decorator to turn a user defined function into a BigQuery remote function. @@ -613,7 +632,7 @@ def remote_function( * BigQuery Data Editor (roles/bigquery.dataEditor) * BigQuery Connection Admin (roles/bigquery.connectionAdmin) * Cloud Functions Developer (roles/cloudfunctions.developer) - * Service Account User (roles/iam.serviceAccountUser) + * Service Account User (roles/iam.serviceAccountUser) on the service account `PROJECT_NUMBER-compute@developer.gserviceaccount.com` * Storage Object Viewer (roles/storage.objectViewer) * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) @@ -664,10 +683,16 @@ def remote_function( reuse (bool, Optional): Reuse the remote function if is already exists. `True` by default, which results in reusing an existing remote - function (if any) that was previously created for the same udf. - Setting it to false forces the creation of creating a unique remote function. + function and corresponding cloud function (if any) that was + previously created for the same udf. + Setting it to `False` forces the creation of a unique remote function. If the required remote function does not exist then it would be created irrespective of this param. + name (str, Optional): + Explicit name of the persisted BigQuery remote function. Use it with + caution, because two users working in the same project and dataset + could overwrite each other's remote functions if they use the same + persistent name. """ @@ -739,12 +764,6 @@ def remote_function( f"{constants.FEEDBACK_LINK}" ) - uniq_suffix = None - if not reuse: - uniq_suffix = "".join( - random.choices(string.ascii_lowercase + string.digits, k=8) - ) - # Check connection_id with `LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID` format. if bigquery_connection.count(".") == 1: bq_connection_location, bq_connection_id = bigquery_connection.split(".") @@ -792,8 +811,13 @@ def wrapper(f): bigquery_connection, resource_manager_client, ) + rf_name, cf_name = remote_function_client.provision_bq_remote_function( - f, ibis_signature.input_types, ibis_signature.output_type, uniq_suffix + f, + ibis_signature.input_types, + ibis_signature.output_type, + reuse, + name, ) node = remote_function_node(dataset_ref.routine(rf_name), ibis_signature) diff --git a/bigframes/series.py b/bigframes/series.py index 8eadee37ed0..12e72c58b61 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -42,6 +42,7 @@ STABLE_SORTS, ) import bigframes.core.scalar as scalars +import bigframes.core.utils as utils import bigframes.core.window import bigframes.dataframe import bigframes.dtypes @@ -310,11 +311,20 @@ def drop( block = block.drop_columns([condition_id]) return Series(block.select_column(self._value_column)) - def droplevel(self, level: LevelsType): + def droplevel(self, level: LevelsType, axis: int | str = 0): resolved_level_ids = self._resolve_levels(level) return Series(self._block.drop_levels(resolved_level_ids)) - def reorder_levels(self, order: LevelsType): + def swaplevel(self, i: int = -2, j: int = -1): + level_i = self._block.index_columns[i] + level_j = self._block.index_columns[j] + mapping = {level_i: level_j, level_j: level_i} + reordering = [ + mapping.get(index_id, index_id) for index_id in self._block.index_columns + ] + return Series(self._block.reorder_levels(reordering)) + + def reorder_levels(self, order: LevelsType, axis: int | str = 0): resolved_level_ids = self._resolve_levels(order) return Series(self._block.reorder_levels(resolved_level_ids)) @@ -352,6 +362,14 @@ def cumsum(self) -> Series: agg_ops.sum_op, bigframes.core.WindowSpec(following=0) ) + def ffill(self, *, limit: typing.Optional[int] = None) -> Series: + window = bigframes.core.WindowSpec(preceding=limit, following=0) + return self._apply_window_op(agg_ops.LastNonNullOp(), window) + + def bfill(self, *, limit: typing.Optional[int] = None) -> Series: + window = bigframes.core.WindowSpec(preceding=0, following=limit) + return self._apply_window_op(agg_ops.FirstNonNullOp(), window) + def cummax(self) -> Series: return self._apply_window_op( agg_ops.max_op, bigframes.core.WindowSpec(following=0) @@ -375,7 +393,16 @@ def shift(self, periods: int = 1) -> Series: return self._apply_window_op(agg_ops.ShiftOp(periods), window) def diff(self, periods: int = 1) -> Series: - return self - self.shift(periods=periods) + window = bigframes.core.WindowSpec( + preceding=periods if periods > 0 else None, + following=-periods if periods < 0 else None, + ) + return self._apply_window_op(agg_ops.DiffOp(periods), window) + + def pct_change(self, periods: int = 1) -> Series: + # Future versions of pandas will not perfrom ffill automatically + series = self.ffill() + return Series(block_ops.pct_change(series._block, periods=periods)) def rank( self, @@ -390,6 +417,47 @@ def rank( def fillna(self, value=None) -> Series: return self._apply_binary_op(value, ops.fillna_op) + def replace( + self, to_replace: typing.Any, value: typing.Any = None, *, regex: bool = False + ): + if regex: + if not (isinstance(to_replace, str) and isinstance(value, str)): + raise NotImplementedError( + f"replace regex mode only supports strings for 'to_replace' and 'value'. {constants.FEEDBACK_LINK}" + ) + block, result_col = self._block.apply_unary_op( + self._value_column, + ops.ReplaceRegexOp(to_replace, value), + result_label=self.name, + ) + return Series(block.select_column(result_col)) + elif utils.is_dict_like(to_replace): + raise NotImplementedError( + f"Dict 'to_replace' not supported. {constants.FEEDBACK_LINK}" + ) + elif utils.is_list_like(to_replace): + block, cond = self._block.apply_unary_op( + self._value_column, ops.IsInOp(to_replace) + ) + block, result_col = block.apply_binary_op( + cond, + self._value_column, + ops.partial_arg1(ops.where_op, value), + result_label=self.name, + ) + return Series(block.select_column(result_col)) + else: # Scalar + block, cond = self._block.apply_unary_op( + self._value_column, ops.BinopPartialLeft(ops.eq_op, to_replace) + ) + block, result_col = block.apply_binary_op( + cond, + self._value_column, + ops.partial_arg1(ops.where_op, value), + result_label=self.name, + ) + return Series(block.select_column(result_col)) + def dropna( self, *, @@ -414,52 +482,16 @@ def tail(self, n: int = 5) -> Series: def nlargest(self, n: int = 5, keep: str = "first") -> Series: if keep not in ("first", "last", "all"): raise ValueError("'keep must be one of 'first', 'last', or 'all'") - block = self._block - if keep == "last": - block = block.reversed() - ordering = ( - OrderingColumnReference( - self._value_column, direction=OrderingDirection.DESC - ), + return Series( + block_ops.nlargest(self._block, n, [self._value_column], keep=keep) ) - block = block.order_by(ordering, stable=True) - if keep in ("first", "last"): - return Series(block.slice(0, n)) - else: # keep == "all": - block, counter = block.apply_window_op( - self._value_column, - agg_ops.rank_op, - window_spec=WindowSpec(ordering=ordering), - ) - block, condition = block.apply_unary_op( - counter, ops.partial_right(ops.le_op, n) - ) - block = block.filter(condition) - block = block.select_column(self._value_column) - return Series(block) def nsmallest(self, n: int = 5, keep: str = "first") -> Series: if keep not in ("first", "last", "all"): raise ValueError("'keep must be one of 'first', 'last', or 'all'") - block = self._block - if keep == "last": - block = block.reversed() - ordering = (OrderingColumnReference(self._value_column),) - block = block.order_by(ordering, stable=True) - if keep in ("first", "last"): - return Series(block.slice(0, n)) - else: # keep == "all": - block, counter = block.apply_window_op( - self._value_column, - agg_ops.rank_op, - window_spec=WindowSpec(ordering=ordering), - ) - block, condition = block.apply_unary_op( - counter, ops.partial_right(ops.le_op, n) - ) - block = block.filter(condition) - block = block.select_column(self._value_column) - return Series(block) + return Series( + block_ops.nsmallest(self._block, n, [self._value_column], keep=keep) + ) def isin(self, values) -> "Series" | None: if not _is_list_like(values): @@ -697,13 +729,9 @@ def _central_moment(self, n: int) -> float: """Useful helper for calculating central moment statistics""" # Nth central moment is mean((x-mean(x))^n) # See: https://en.wikipedia.org/wiki/Moment_(mathematics) - mean = self.mean() - mean_deltas = self - mean - delta_power = mean_deltas - # TODO(tbergeron): Replace with pow once implemented - for i in range(1, n): - delta_power = delta_power * mean_deltas - return delta_power.mean() + mean_deltas = self - self.mean() + delta_powers = mean_deltas**n + return delta_powers.mean() def agg(self, func: str | typing.Sequence[str]) -> scalars.Scalar | Series: if _is_list_like(func): @@ -1096,6 +1124,85 @@ def add_prefix(self, prefix: str, axis: int | str | None = None) -> Series: def add_suffix(self, suffix: str, axis: int | str | None = None) -> Series: return Series(self._get_block().add_suffix(suffix)) + def filter( + self, + items: typing.Optional[typing.Iterable] = None, + like: typing.Optional[str] = None, + regex: typing.Optional[str] = None, + axis: typing.Optional[typing.Union[str, int]] = None, + ) -> Series: + if (axis is not None) and utils.get_axis_number(axis) != 0: + raise ValueError(f"Invalid axis for series: {axis}") + if sum([(items is not None), (like is not None), (regex is not None)]) != 1: + raise ValueError( + "Need to provide exactly one of 'items', 'like', or 'regex'" + ) + if len(self._block.index_columns) > 1: + raise NotImplementedError( + "Method filter does not support rows multiindex. {constants.FEEDBACK_LINK}" + ) + if (like is not None) or (regex is not None): + block = self._block + block, label_string_id = block.apply_unary_op( + self._block.index_columns[0], + ops.AsTypeOp(pandas.StringDtype(storage="pyarrow")), + ) + if like is not None: + block, mask_id = block.apply_unary_op( + label_string_id, ops.ContainsStringOp(pat=like) + ) + else: # regex + assert regex is not None + block, mask_id = block.apply_unary_op( + label_string_id, ops.ContainsRegexOp(pat=regex) + ) + + block = block.filter(mask_id) + block = block.select_columns([self._value_column]) + return Series(block) + elif items is not None: + # Behavior matches pandas 2.1+, older pandas versions would reindex + block = self._block + block, mask_id = block.apply_unary_op( + self._block.index_columns[0], ops.IsInOp(values=list(items)) + ) + block = block.filter(mask_id) + block = block.select_columns([self._value_column]) + return Series(block) + else: + raise ValueError("Need to provide 'items', 'like', or 'regex'") + + def reindex(self, index=None, *, validate: typing.Optional[bool] = None): + if validate and not self.index.is_unique: + raise ValueError("Original index must be unique to reindex") + keep_original_names = False + if isinstance(index, indexes.Index): + new_indexer = bigframes.dataframe.DataFrame(data=index._data._get_block())[ + [] + ] + else: + if not isinstance(index, pandas.Index): + keep_original_names = True + index = pandas.Index(index) + if index.nlevels != self.index.nlevels: + raise NotImplementedError( + "Cannot reindex with index with different nlevels" + ) + new_indexer = bigframes.dataframe.DataFrame(index=index)[[]] + # multiindex join is senstive to index names, so we will set all these + result = new_indexer.rename_axis(range(new_indexer.index.nlevels)).join( + self.to_frame().rename_axis(range(self.index.nlevels)), + how="left", + ) + # and then reset the names after the join + result_block = result.rename_axis( + self.index.names if keep_original_names else index.names + )._block + return Series(result_block) + + def reindex_like(self, other: Series, *, validate: typing.Optional[bool] = None): + return self.reindex(other.index, validate=validate) + def drop_duplicates(self, *, keep: str = "first") -> Series: block = block_ops.drop_duplicates(self._block, (self._value_column,), keep) return Series(block) @@ -1216,14 +1323,7 @@ def to_xarray(self): def _throw_if_index_contains_duplicates( self, error_message: typing.Optional[str] = None ) -> None: - duplicates_block, _ = block_ops.indicate_duplicates( - self._get_block(), self._get_block().index_columns - ) - duplicates_block = duplicates_block.with_column_labels( - ["values", "is_duplicate"] - ) - duplicates_df = bigframes.dataframe.DataFrame(duplicates_block) - if duplicates_df["is_duplicate"].any(): + if not self.index.is_unique: error_message = ( error_message if error_message diff --git a/bigframes/session.py b/bigframes/session.py index 17444077726..3ca79a7b531 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -424,25 +424,30 @@ def read_gbq( ) -> dataframe.DataFrame: # TODO(b/281571214): Generate prompt to show the progress of read_gbq. if _is_query(query): - return self.read_gbq_query( + return self._read_gbq_query( query, index_col=index_col, col_order=col_order, max_results=max_results, + api_name="read_gbq", ) else: # TODO(swast): Query the snapshot table but mark it as a # deterministic query so we can avoid serializing if we have a # unique index. - return self.read_gbq_table( + return self._read_gbq_table( query, index_col=index_col, col_order=col_order, max_results=max_results, + api_name="read_gbq", ) def _query_to_destination( - self, query: str, index_cols: List[str] + self, + query: str, + index_cols: List[str], + api_name: str, ) -> Tuple[Optional[bigquery.TableReference], Optional[bigquery.QueryJob]]: # If there are no index columns, then there's no reason to cache to a # (clustered) session table, as we'll just have to query it again to @@ -464,7 +469,7 @@ def _query_to_destination( # operations are as speedy as they can be. try: ibis_expr = self.ibis_client.sql(query) - return self._ibis_to_session_table(ibis_expr, index_cols), None + return self._ibis_to_session_table(ibis_expr, index_cols, api_name), None except google.api_core.exceptions.BadRequest: # Some SELECT statements still aren't compatible with CREATE TEMP # TABLE ... AS SELECT ... statements. For example, if the query has @@ -490,15 +495,33 @@ def read_gbq_query( See also: :meth:`Session.read_gbq`. """ + return self._read_gbq_query( + query=query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + api_name="read_gbq_query", + ) + + def _read_gbq_query( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + api_name: str, + ) -> dataframe.DataFrame: # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so # these docstrings are inline. - if isinstance(index_col, str): index_cols = [index_col] else: index_cols = list(index_col) - destination, query_job = self._query_to_destination(query, index_cols) + destination, query_job = self._query_to_destination( + query, index_cols, api_name="read_gbq_query" + ) # If there was no destination table, that means the query must have # been DDL or DML. Return some job metadata, instead. @@ -535,6 +558,23 @@ def read_gbq_table( See also: :meth:`Session.read_gbq`. """ + return self._read_gbq_table( + query=query, + index_col=index_col, + col_order=col_order, + max_results=max_results, + api_name="read_gbq_table", + ) + + def _read_gbq_table( + self, + query: str, + *, + index_col: Iterable[str] | str = (), + col_order: Iterable[str] = (), + max_results: Optional[int] = None, + api_name: str, + ) -> dataframe.DataFrame: if max_results and max_results <= 0: raise ValueError("`max_results` should be a positive number.") @@ -646,7 +686,8 @@ def read_gbq_table( # rows for which row numbers must be generated table_expression = table_expression.limit(max_results) table_expression, ordering = self._create_sequential_ordering( - table_expression + table=table_expression, + api_name=api_name, ) hidden_cols = ( (ordering.total_order_col.column_id,) @@ -667,6 +708,7 @@ def read_gbq_table( hidden_cols=hidden_cols, ordering=ordering, is_total_ordering=is_total_ordering, + api_name=api_name, ) def _read_gbq_with_ordering( @@ -680,6 +722,7 @@ def _read_gbq_with_ordering( hidden_cols: Iterable[str] = (), ordering: core.ExpressionOrdering, is_total_ordering: bool = False, + api_name: str, ) -> dataframe.DataFrame: """Internal helper method that loads DataFrame from Google BigQuery given an ordering column. @@ -698,6 +741,8 @@ def _read_gbq_with_ordering( Columns that should be hidden. Ordering columns may (not always) be hidden ordering: Column name to be used for ordering. If not supplied, a default ordering is generated. + api_name: + The name of the API method. Returns: A DataFrame representing results of the query or table. @@ -723,7 +768,9 @@ def _read_gbq_with_ordering( if not is_total_ordering: # Rows are not ordered, we need to generate a default ordering and materialize it table_expression, ordering = self._create_sequential_ordering( - table_expression, index_cols + table=table_expression, + index_cols=index_cols, + api_name=api_name, ) index_col_values = [table_expression[index_id] for index_id in index_cols] if not col_labels: @@ -846,6 +893,11 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame Returns: bigframes.dataframe.DataFrame: The BigQuery DataFrame. """ + return self._read_pandas(pandas_dataframe, "read_pandas") + + def _read_pandas( + self, pandas_dataframe: pandas.DataFrame, api_name: str + ) -> dataframe.DataFrame: col_labels, idx_labels = ( pandas_dataframe.columns.to_list(), pandas_dataframe.index.names, @@ -878,6 +930,7 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame job_config = bigquery.LoadJobConfig(schema=schema) job_config.clustering_fields = cluster_cols + job_config.labels = {"bigframes-api": api_name} load_table_destination = self._create_session_table() load_job = self.bqclient.load_table_from_dataframe( @@ -910,6 +963,7 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame hidden_cols=(ordering_col,), ordering=ordering, is_total_ordering=True, + api_name=api_name, ) return df @@ -991,6 +1045,7 @@ def read_csv( job_config.autodetect = True job_config.field_delimiter = sep job_config.encoding = encoding + job_config.labels = {"bigframes-api": "read_csv"} # We want to match pandas behavior. If header is 0, no rows should be skipped, so we # do not need to set `skip_leading_rows`. If header is None, then there is no header. @@ -1048,7 +1103,7 @@ def read_pickle( pandas_obj.name = "0" bigframes_df = self.read_pandas(pandas_obj.to_frame()) return bigframes_df[bigframes_df.columns[0]] - return self.read_pandas(pandas_obj) + return self._read_pandas(pandas_obj, "read_pickle") def read_parquet( self, @@ -1063,6 +1118,7 @@ def read_parquet( job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED job_config.source_format = bigquery.SourceFormat.PARQUET job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY + job_config.labels = {"bigframes-api": "read_parquet"} return self._read_bigquery_load_job(path, table, job_config=job_config) @@ -1109,6 +1165,7 @@ def read_json( job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY job_config.autodetect = True job_config.encoding = encoding + job_config.labels = {"bigframes-api": "read_json"} return self._read_bigquery_load_job( path_or_buf, @@ -1176,7 +1233,10 @@ def _create_session_table(self) -> bigquery.TableReference: return dataset.table(table_name) def _create_sequential_ordering( - self, table: ibis_types.Table, index_cols: Iterable[str] = () + self, + table: ibis_types.Table, + index_cols: Iterable[str] = (), + api_name: str = "", ) -> Tuple[ibis_types.Table, core.ExpressionOrdering]: # Since this might also be used as the index, don't use the default # "ordering ID" name. @@ -1188,6 +1248,7 @@ def _create_sequential_ordering( table_ref = self._ibis_to_session_table( table, cluster_cols=list(index_cols) + [default_ordering_name], + api_name=api_name, ) table = self.ibis_client.sql(f"SELECT * FROM `{table_ref.table_id}`") ordering_reference = core.OrderingColumnReference(default_ordering_name) @@ -1199,7 +1260,10 @@ def _create_sequential_ordering( return table, ordering def _ibis_to_session_table( - self, table: ibis_types.Table, cluster_cols: Iterable[str] + self, + table: ibis_types.Table, + cluster_cols: Iterable[str], + api_name: str, ) -> bigquery.TableReference: clusterable_cols = [ col for col in cluster_cols if _can_cluster(table[col].type()) @@ -1207,10 +1271,14 @@ def _ibis_to_session_table( return self._query_to_session_table( self.ibis_client.compile(table), cluster_cols=clusterable_cols, + api_name=api_name, ) def _query_to_session_table( - self, query_text: str, cluster_cols: Iterable[str] + self, + query_text: str, + cluster_cols: Iterable[str], + api_name: str, ) -> bigquery.TableReference: if len(list(cluster_cols)) > _MAX_CLUSTER_COLUMNS: raise ValueError( @@ -1236,6 +1304,7 @@ def _query_to_session_table( # otherwise we get `BadRequest: 400 OPTIONS on temporary tables are not # supported`. job_config.labels = {"source": "bigquery-dataframes-temp"} + job_config.labels["bigframes-api"] = api_name try: self._start_query( @@ -1253,6 +1322,7 @@ def remote_function( dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, + name: Optional[str] = None, ): """Decorator to turn a user defined function into a BigQuery remote function. @@ -1280,7 +1350,7 @@ def remote_function( * BigQuery Data Editor (roles/bigquery.dataEditor) * BigQuery Connection Admin (roles/bigquery.connectionAdmin) * Cloud Functions Developer (roles/cloudfunctions.developer) - * Service Account User (roles/iam.serviceAccountUser) + * Service Account User (roles/iam.serviceAccountUser) on the service account `PROJECT_NUMBER-compute@developer.gserviceaccount.com` * Storage Object Viewer (roles/storage.objectViewer) * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.) @@ -1311,10 +1381,16 @@ def remote_function( reuse (bool, Optional): Reuse the remote function if already exists. `True` by default, which will result in reusing an existing remote - function (if any) that was previously created for the same udf. - Setting it to false would force creating a unique remote function. + function and corresponding cloud function (if any) that was + previously created for the same udf. + Setting it to `False` would force creating a unique remote function. If the required remote function does not exist then it would be created irrespective of this param. + name (str, Optional): + Explicit name of the persisted BigQuery remote function. Use it with + caution, because two users working in the same project and dataset + could overwrite each other's remote functions if they use the same + persistent name. Returns: callable: A remote function object pointing to the cloud assets created in the background to support the remote execution. The cloud assets can be @@ -1331,6 +1407,7 @@ def remote_function( dataset=dataset, bigquery_connection=bigquery_connection, reuse=reuse, + name=name, ) def read_gbq_function( diff --git a/noxfile.py b/noxfile.py index 2355a9b27ba..033bbfefe40 100644 --- a/noxfile.py +++ b/noxfile.py @@ -42,6 +42,7 @@ "pytest", "pytest-cov", "pytest-asyncio", + "pytest-mock", ] UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = [] diff --git a/samples/snippets/remote_function.py b/samples/snippets/remote_function.py index 37972672c32..9998a23eb2c 100644 --- a/samples/snippets/remote_function.py +++ b/samples/snippets/remote_function.py @@ -39,11 +39,19 @@ def run_remote_function_and_read_gbq_function(project_id: str): # 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 would be using a pre-created connection named - # `bigframes-rf-conn`. 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([float], str, bigquery_connection="bigframes-rf-conn") + # `bigframes-rf-conn`. 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( + [float], + str, + bigquery_connection="bigframes-rf-conn", + reuse=False, + ) def get_bucket(num): if not num: return "NA" @@ -80,9 +88,11 @@ def get_bucket(num): # 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 decorated it + # code instead. Let's define another scalar custom function and decorate it # as a remote function - @bpd.remote_function([str], str, bigquery_connection="bigframes-rf-conn") + @bpd.remote_function( + [str], str, bigquery_connection="bigframes-rf-conn", reuse=False + ) def get_hash(input): import hashlib diff --git a/setup.py b/setup.py index 139873e6fc9..69b71c88f13 100644 --- a/setup.py +++ b/setup.py @@ -47,6 +47,7 @@ "ibis-framework[bigquery] >=6.0.0,<=6.1.0", "pandas >=1.5.0", "pydata-google-auth >=1.8.2", + "requests >=2.27.1", "scikit-learn >=1.2.2", "sqlalchemy >=1.4,<3.0", "ipywidgets >=7.7.1", @@ -58,7 +59,7 @@ "pandas-gbq >=0.19.0", ], # Packages required for basic development flow. - "dev": ["pytest", "pre-commit", "nox", "google-cloud-testutils"], + "dev": ["pytest", "pytest-mock", "pre-commit", "nox", "google-cloud-testutils"], } extras["all"] = list(sorted(frozenset(itertools.chain.from_iterable(extras.values())))) diff --git a/testing/constraints-3.9.txt b/testing/constraints-3.9.txt index 523256ee835..cd69d45dc92 100644 --- a/testing/constraints-3.9.txt +++ b/testing/constraints-3.9.txt @@ -9,7 +9,7 @@ cachetools==5.3.0 certifi==2022.12.7 cffi==1.15.1 cfgv==3.3.1 -charset-normalizer==3.1.0 +charset-normalizer==2.0.0 click==8.1.3 cloudpickle==2.0.0 colorlog==6.7.0 @@ -90,13 +90,14 @@ pyperclip==1.8.2 pytest==7.2.2 pytest-asyncio==0.21.0 pytest-cov==4.0.0 +pytest-mock==3.11.1 pytest-retry==1.1.0 pytest-xdist==3.2.1 python-dateutil==2.8.2 pytz==2023.3 PyYAML==6.0 readme-renderer==37.3 -requests==2.28.2 +requests==2.27.1 requests-oauthlib==1.3.1 requests-toolbelt==0.10.1 rfc3986==2.0.0 diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py index 9b2872d6737..a8613dfeb9b 100644 --- a/tests/system/large/ml/test_ensemble.py +++ b/tests/system/large/ml/test_ensemble.py @@ -70,7 +70,7 @@ def test_xgbregressor_dart_booster_multiple_params( ): model = bigframes.ml.ensemble.XGBRegressor( booster="dart", - tree_method="AUTO", + tree_method="auto", min_tree_child_weight=2, colsample_bytree=0.95, colsample_bylevel=0.95, @@ -121,7 +121,7 @@ def test_xgbregressor_dart_booster_multiple_params( in reloaded_model._bqml_model.model_name ) assert reloaded_model.booster == "DART" - assert reloaded_model.dart_normalized_type == "TREE" + assert reloaded_model.dart_normalized_type == "tree" assert reloaded_model.tree_method == "AUTO" assert reloaded_model.colsample_bytree == 0.95 assert reloaded_model.colsample_bylevel == 0.95 @@ -185,7 +185,7 @@ def test_xgbclassifier_dart_booster_multiple_params( ): model = bigframes.ml.ensemble.XGBClassifier( booster="dart", - tree_method="AUTO", + tree_method="auto", min_tree_child_weight=2, colsample_bytree=0.95, colsample_bylevel=0.95, @@ -235,7 +235,7 @@ def test_xgbclassifier_dart_booster_multiple_params( in reloaded_model._bqml_model.model_name ) assert reloaded_model.booster == "DART" - assert reloaded_model.dart_normalized_type == "TREE" + assert reloaded_model.dart_normalized_type == "tree" assert reloaded_model.tree_method == "AUTO" assert reloaded_model.colsample_bytree == 0.95 assert reloaded_model.colsample_bylevel == 0.95 @@ -297,7 +297,7 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset @pytest.mark.flaky(retries=2, delay=120) def test_randomforestregressor_multiple_params(penguins_df_default_index, dataset_id): model = bigframes.ml.ensemble.RandomForestRegressor( - tree_method="AUTO", + tree_method="auto", min_tree_child_weight=2, colsample_bytree=0.95, colsample_bylevel=0.95, diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index 8033f79c47b..2f231f40c95 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -61,16 +61,32 @@ def get_remote_function_endpoints(bigquery_client, dataset_id): return endpoints -def get_cloud_functions(functions_client, project, location, name_prefix="bigframes-"): +def get_cloud_functions( + functions_client, project, location, name=None, name_prefix=None +): """Get the cloud functions in the given project and location.""" + + assert ( + not name or not name_prefix + ), f"At most one of the {name.__name__} or {name_prefix.__name__} can be passed." + _, location = get_remote_function_locations(location) parent = f"projects/{project}/locations/{location}" request = functions_v2.ListFunctionsRequest(parent=parent) page_result = functions_client.list_functions(request=request) - full_name_prefix = parent + f"/functions/{name_prefix}" for response in page_result: - if not name_prefix or response.name.startswith(full_name_prefix): - yield response + # If name is provided and it does not match then skip + if bool(name): + full_name = parent + f"/functions/{name}" + if response.name != full_name: + continue + # If name prefix is provided and it does not match then skip + elif bool(name_prefix): + full_name_prefix = parent + f"/functions/{name_prefix}" + if not response.name.startswith(full_name_prefix): + continue + + yield response def delete_cloud_function(functions_client, full_name): @@ -84,8 +100,17 @@ def cleanup_remote_function_assets( bigquery_client, functions_client, remote_udf, ignore_failures=True ): """Clean up the GCP assets behind a bigframes remote function.""" + + # Clean up BQ remote function try: bigquery_client.delete_routine(remote_udf.bigframes_remote_function) + except Exception: + # By default don't raise exception in cleanup + if not ignore_failures: + raise + + # Clean up cloud function + try: delete_cloud_function(functions_client, remote_udf.bigframes_cloud_function) except Exception: # By default don't raise exception in cleanup @@ -94,7 +119,15 @@ def cleanup_remote_function_assets( def make_uniq_udf(udf): - """Transform a udf to another with same behavior but a unique name.""" + """Transform a udf to another with same behavior but a unique name. + Use this to test remote functions with reuse=True, in which case parallel + instances of the same tests may evaluate same named cloud functions and BQ + remote functions, therefore interacting with each other and causing unwanted + failures. With this method one can transform a udf into another with the + same behavior but a different name which will remain unique for the + lifetime of one test instance. + """ + prefixer = test_utils.prefixer.Prefixer(udf.__name__, "") udf_uniq_name = prefixer.create_prefix() udf_file_name = f"{udf_uniq_name}.py" @@ -111,7 +144,18 @@ def make_uniq_udf(udf): target_code = source_code.replace(source_key, target_key, 1) f.write(target_code) spec = importlib.util.spec_from_file_location(udf_file_name, udf_file_path) - return getattr(spec.loader.load_module(), udf_uniq_name), tmpdir + udf_uniq = getattr(spec.loader.load_module(), udf_uniq_name) + + # This is a bit of a hack but we need to remove the reference to a foreign + # module, otherwise the serialization would keep the foreign module + # reference and deserialization would fail with error like following: + # ModuleNotFoundError: No module named 'add_one_2nxcmd9j' + # TODO(shobs): Figure out if there is a better way of generating the unique + # function object, but for now let's just set it to same module as the + # original udf. + udf_uniq.__module__ = udf.__module__ + + return udf_uniq, tmpdir @pytest.fixture(scope="module") @@ -136,7 +180,10 @@ def cleanup_cloud_functions(session, functions_client, dataset_id_permanent): ) delete_count = 0 for cloud_function in get_cloud_functions( - functions_client, session.bqclient.project, session.bqclient.location + functions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix="bigframes-", ): # Ignore bigframes cloud functions referred by the remote functions in # the permanent dataset @@ -524,15 +571,6 @@ def add_one(x): # Make a unique udf add_one_uniq, add_one_uniq_dir = make_uniq_udf(add_one) - # This is a bit of a hack but we need to remove the reference to a foreign - # module, otherwise the serialization would keep the foreign module - # reference and deserialization would fail with error like following: - # ModuleNotFoundError: No module named 'add_one_2nxcmd9j' - # TODO(shobs): Figure out if there is a better way of generating the unique - # function object, but for now let's just set it to same module as the - # original udf. - add_one_uniq.__module__ = add_one.__module__ - # Expected cloud function name for the unique udf add_one_uniq_cf_name = get_cloud_function_name(add_one_uniq) @@ -542,7 +580,7 @@ def add_one(x): functions_client, session.bqclient.project, session.bqclient.location, - name_prefix=add_one_uniq_cf_name, + name=add_one_uniq_cf_name, ) ) assert len(cloud_functions) == 0 @@ -563,7 +601,7 @@ def add_one(x): functions_client, session.bqclient.project, session.bqclient.location, - name_prefix=add_one_uniq_cf_name, + name=add_one_uniq_cf_name, ) ) assert len(cloud_functions) == 1 @@ -611,7 +649,7 @@ def inner_test(): functions_client, session.bqclient.project, session.bqclient.location, - name_prefix=add_one_uniq_cf_name, + name=add_one_uniq_cf_name, ) ) assert len(cloud_functions) == 0 @@ -633,7 +671,7 @@ def inner_test(): functions_client, session.bqclient.project, session.bqclient.location, - name_prefix=add_one_uniq_cf_name, + name=add_one_uniq_cf_name, ) ) assert len(cloud_functions) == 1 @@ -776,3 +814,221 @@ def test_remote_udf_lambda( cleanup_remote_function_assets( session.bqclient, functions_client, add_one_lambda_remote ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_with_explicit_name( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def square(x): + return x * x + + prefixer = test_utils.prefixer.Prefixer(square.__name__, "") + rf_name = prefixer.create_prefix() + expected_remote_function = f"{dataset_id}.{rf_name}" + + # Initially the expected BQ remote function should not exist + with pytest.raises(NotFound): + session.bqclient.get_routine(expected_remote_function) + + # Create the remote function with the name provided explicitly + square_remote = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + reuse=False, + name=rf_name, + )(square) + + # The remote function should reflect the explicitly provided name + assert square_remote.bigframes_remote_function == expected_remote_function + + # Now the expected BQ remote function should exist + session.bqclient.get_routine(expected_remote_function) + + # The behavior of the created remote function should be as expected + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_too"] + bf_result_col = bf_int64_col.apply(square_remote) + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas() + + pd_int64_col = scalars_pandas_df["int64_too"] + pd_result_col = pd_int64_col.apply(square) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col.dtype is Int64Dtype() + # pd_int64_col.apply(square).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, square_remote + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_with_explicit_name_reuse( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + dirs_to_cleanup = [] + + # Define a user code + def square(x): + return x * x + + # Make it a unique udf + square_uniq, square_uniq_dir = make_uniq_udf(square) + dirs_to_cleanup.append(square_uniq_dir) + + # Define a common routine which accepts a remote function and the + # corresponding user defined function and tests that bigframes bahavior + # on the former is in parity with the pandas behaviour on the latter + def test_internal(rf, udf): + # The behavior of the created remote function should be as expected + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_too"] + bf_result_col = bf_int64_col.apply(rf) + bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas() + + pd_int64_col = scalars_pandas_df["int64_too"] + pd_result_col = pd_int64_col.apply(udf) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col.dtype is Int64Dtype() + # pd_int64_col.apply(square).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + # Create an explicit name for the remote function + prefixer = test_utils.prefixer.Prefixer("foo", "") + rf_name = prefixer.create_prefix() + expected_remote_function = f"{dataset_id}.{rf_name}" + + # Initially the expected BQ remote function should not exist + with pytest.raises(NotFound): + session.bqclient.get_routine(expected_remote_function) + + # Create a new remote function with the name provided explicitly + square_remote1 = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + name=rf_name, + )(square_uniq) + + # The remote function should reflect the explicitly provided name + assert square_remote1.bigframes_remote_function == expected_remote_function + + # Now the expected BQ remote function should exist + routine = session.bqclient.get_routine(expected_remote_function) + square_remote1_created = routine.created + square_remote1_cf_updated = session.cloudfunctionsclient.get_function( + name=square_remote1.bigframes_cloud_function + ).update_time + + # Test pandas parity with square udf + test_internal(square_remote1, square) + + # Now Create another remote function with the same name provided + # explicitly. Since reuse is True by default, the previously created + # remote function with the same name will be reused. + square_remote2 = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + name=rf_name, + )(square_uniq) + + # The new remote function should still reflect the explicitly provided name + assert square_remote2.bigframes_remote_function == expected_remote_function + + # The expected BQ remote function should still exist + routine = session.bqclient.get_routine(expected_remote_function) + square_remote2_created = routine.created + square_remote2_cf_updated = session.cloudfunctionsclient.get_function( + name=square_remote2.bigframes_cloud_function + ).update_time + + # The new remote function should reflect that the previous BQ remote + # function and the cloud function were reused instead of creating anew + assert square_remote2_created == square_remote1_created + assert ( + square_remote2.bigframes_cloud_function + == square_remote1.bigframes_cloud_function + ) + assert square_remote2_cf_updated == square_remote1_cf_updated + + # Test again that the new remote function is actually same as the + # previous remote function + test_internal(square_remote2, square) + + # Now define a different user code + def plusone(x): + return x + 1 + + # Make it a unique udf + plusone_uniq, plusone_uniq_dir = make_uniq_udf(plusone) + dirs_to_cleanup.append(plusone_uniq_dir) + + # Now Create a third remote function with the same name provided + # explicitly. Even though reuse is True by default, the previously + # created remote function with the same name should not be reused since + # this time it is a different user code. + plusone_remote = session.remote_function( + [int], + int, + dataset_id, + bq_cf_connection, + name=rf_name, + )(plusone_uniq) + + # The new remote function should still reflect the explicitly provided name + assert plusone_remote.bigframes_remote_function == expected_remote_function + + # The expected BQ remote function should still exist + routine = session.bqclient.get_routine(expected_remote_function) + plusone_remote_created = routine.created + plusone_remote_cf_updated = session.cloudfunctionsclient.get_function( + name=plusone_remote.bigframes_cloud_function + ).update_time + + # The new remote function should reflect that the previous BQ remote + # function and the cloud function were NOT reused, instead were created + # anew + assert plusone_remote_created > square_remote2_created + assert ( + plusone_remote.bigframes_cloud_function + != square_remote2.bigframes_cloud_function + ) + assert plusone_remote_cf_updated > square_remote2_cf_updated + + # Test again that the new remote function is equivalent to the new user + # defined function + test_internal(plusone_remote, plusone) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, functions_client, square_remote1 + ) + cleanup_remote_function_assets( + session.bqclient, functions_client, square_remote2 + ) + cleanup_remote_function_assets( + session.bqclient, functions_client, plusone_remote + ) + for dir_ in dirs_to_cleanup: + shutil.rmtree(dir_) diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index 6c3e8e06f5b..ace943956fe 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -18,6 +18,7 @@ import pandas as pd import pyarrow as pa +import pytest import pytz import bigframes @@ -278,6 +279,7 @@ def test_model_predict_with_unnamed_index( ) +@pytest.mark.flaky(retries=2, delay=120) def test_model_generate_text( bqml_palm2_text_generator_model: core.BqmlModel, llm_text_df ): diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index 8df4145fcf0..c71bbbe3b0c 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -16,33 +16,14 @@ from bigframes.ml import decomposition -_PD_NEW_PENGUINS = pd.DataFrame( - { - "tag_number": [1633, 1672, 1690], - "species": [ - "Adelie Penguin (Pygoscelis adeliae)", - "Gentoo penguin (Pygoscelis papua)", - "Adelie Penguin (Pygoscelis adeliae)", - ], - "island": ["Dream", "Biscoe", "Torgersen"], - "culmen_length_mm": [37.8, 46.5, 41.1], - "culmen_depth_mm": [18.1, 14.8, 18.6], - "flipper_length_mm": [193.0, 217.0, 189.0], - "body_mass_g": [3750.0, 5200.0, 3325.0], - "sex": ["MALE", "FEMALE", "MALE"], - } -).set_index("tag_number") - -def test_pca_predict(session, penguins_pca_model: decomposition.PCA): - new_penguins = session.read_pandas(_PD_NEW_PENGUINS) - - predictions = penguins_pca_model.predict(new_penguins).to_pandas() +def test_pca_predict(penguins_pca_model, new_penguins_df): + predictions = penguins_pca_model.predict(new_penguins_df).to_pandas() expected = pd.DataFrame( { - "principal_component_1": [-1.459, 2.258, -1.685], - "principal_component_2": [-1.120, -1.351, -0.874], - "principal_component_3": [-0.646, 0.443, -0.704], + "principal_component_1": [-1.314041, -0.855813, -1.848786], + "principal_component_2": [-0.889106, -1.259753, -0.983304], + "principal_component_3": [-0.704345, 0.322555, -0.095759], }, dtype="Float64", index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 85c3cce1d77..a85777c59d6 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -13,6 +13,7 @@ # limitations under the License. import operator +import tempfile import typing from typing import Tuple @@ -137,6 +138,46 @@ def test_tail_with_custom_column_labels(scalars_df_index, scalars_pandas_df_inde pandas.testing.assert_frame_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + ("all",), + ], +) +def test_df_nlargest(scalars_df_index, scalars_pandas_df_index, keep): + bf_result = scalars_df_index.nlargest( + 3, ["bool_col", "int64_too"], keep=keep + ).to_pandas() + pd_result = scalars_pandas_df_index.nlargest( + 3, ["bool_col", "int64_too"], keep=keep + ) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + ("all",), + ], +) +def test_df_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): + bf_result = scalars_df_index.nsmallest(6, ["bool_col"], keep=keep).to_pandas() + pd_result = scalars_pandas_df_index.nsmallest(6, ["bool_col"], keep=keep) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + def test_get_column_by_attr(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs series = scalars_df.int64_col @@ -582,6 +623,22 @@ def test_df_fillna(scalars_dfs): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_df_ffill(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[["int64_col", "float64_col"]].ffill(limit=1).to_pandas() + pd_result = scalars_pandas_df[["int64_col", "float64_col"]].ffill(limit=1) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_bfill(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[["int64_col", "float64_col"]].bfill().to_pandas() + pd_result = scalars_pandas_df[["int64_col", "float64_col"]].bfill() + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + def test_df_isin_list(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs values = ["Hello, World!", 55555, 2.51, pd.NA, True] @@ -1027,6 +1084,88 @@ def test_df_notnull(scalars_dfs): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.parametrize( + ("left_labels", "right_labels", "overwrite", "fill_value"), + [ + (["a", "b", "c"], ["c", "a", "b"], True, None), + (["a", "b", "c"], ["c", "a", "b"], False, None), + (["a", "b", "c"], ["a", "b", "c"], False, 2), + ], + ids=[ + "one_one_match_overwrite", + "one_one_match_no_overwrite", + "exact_match", + ], +) +def test_combine( + scalars_df_index, + scalars_df_2_index, + scalars_pandas_df_index, + left_labels, + right_labels, + overwrite, + fill_value, +): + if pd.__version__.startswith("1."): + pytest.skip("pd.NA vs NaN not handled well in pandas 1.x.") + columns = ["int64_too", "int64_col", "float64_col"] + + bf_df_a = scalars_df_index[columns] + bf_df_a.columns = left_labels + bf_df_b = scalars_df_2_index[columns] + bf_df_b.columns = right_labels + bf_result = bf_df_a.combine( + bf_df_b, + lambda x, y: x**2 + 2 * x * y + y**2, + overwrite=overwrite, + fill_value=fill_value, + ).to_pandas() + + pd_df_a = scalars_pandas_df_index[columns] + pd_df_a.columns = left_labels + pd_df_b = scalars_pandas_df_index[columns] + pd_df_b.columns = right_labels + pd_result = pd_df_a.combine( + pd_df_b, + lambda x, y: x**2 + 2 * x * y + y**2, + overwrite=overwrite, + fill_value=fill_value, + ) + + # Some dtype inconsistency for all-NULL columns + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +def test_combine_first( + scalars_df_index, + scalars_df_2_index, + scalars_pandas_df_index, +): + if pd.__version__.startswith("1."): + pytest.skip("pd.NA vs NaN not handled well in pandas 1.x.") + columns = ["int64_too", "int64_col", "float64_col"] + + bf_df_a = scalars_df_index[columns].iloc[0:6] + bf_df_a.columns = ["a", "b", "c"] + bf_df_b = scalars_df_2_index[columns].iloc[2:8] + bf_df_b.columns = ["b", "a", "d"] + bf_result = bf_df_a.combine_first(bf_df_b).to_pandas() + + pd_df_a = scalars_pandas_df_index[columns].iloc[0:6] + pd_df_a.columns = ["a", "b", "c"] + pd_df_b = scalars_pandas_df_index[columns].iloc[2:8] + pd_df_b.columns = ["b", "a", "d"] + pd_result = pd_df_a.combine_first(pd_df_b) + + print("pandas") + print(pd_result.to_string()) + print("bigframes") + print(bf_result.to_string()) + + # Some dtype inconsistency for all-NULL columns + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + @pytest.mark.parametrize( ("op"), [ @@ -1145,11 +1284,13 @@ def test_series_binop_axis_index( (["a", "a", "b"], ["c", "c", "d"]), (["a", "b", "c"], ["c", "a", "b"]), (["a", "c", "c"], ["c", "a", "c"]), + (["a", "b", "c"], ["a", "b", "c"]), ], ids=[ "no_overlap", "one_one_match", "multi_match", + "exact_match", ], ) def test_binop_df_df_binary_op( @@ -1361,6 +1502,42 @@ def test_dataframe_general_analytic_op( ) +@pytest.mark.parametrize( + ("periods",), + [ + (1,), + (2,), + (-1,), + ], +) +def test_dataframe_diff(scalars_df_index, scalars_pandas_df_index, periods): + col_names = ["int64_too", "float64_col", "int64_col"] + bf_result = scalars_df_index[col_names].diff(periods=periods).to_pandas() + pd_result = scalars_pandas_df_index[col_names].diff(periods=periods) + pd.testing.assert_frame_equal( + pd_result, + bf_result, + ) + + +@pytest.mark.parametrize( + ("periods",), + [ + (1,), + (2,), + (-1,), + ], +) +def test_dataframe_pct_change(scalars_df_index, scalars_pandas_df_index, periods): + col_names = ["int64_too", "float64_col", "int64_col"] + bf_result = scalars_df_index[col_names].pct_change(periods=periods).to_pandas() + pd_result = scalars_pandas_df_index[col_names].pct_change(periods=periods) + pd.testing.assert_frame_equal( + pd_result, + bf_result, + ) + + def test_dataframe_agg_single_string(scalars_dfs): numeric_cols = ["int64_col", "int64_too", "float64_col"] scalars_df, scalars_pandas_df = scalars_dfs @@ -1675,6 +1852,52 @@ def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index ) +def test_loc_setitem_bool_series_scalar_new_col(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df.loc[bf_df["int64_too"] == 0, "new_col"] = 99 + pd_df.loc[pd_df["int64_too"] == 0, "new_col"] = 99 + + # pandas type difference + pd_df["new_col"] = pd_df["new_col"].astype("Float64") + + pd.testing.assert_frame_equal( + bf_df.to_pandas(), + pd_df, + ) + + +def test_loc_setitem_bool_series_scalar_existing_col(scalars_dfs): + if pd.__version__.startswith("1."): + pytest.skip("this loc overload not supported in pandas 1.x.") + + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df.loc[bf_df["int64_too"] == 1, "string_col"] = "hello" + pd_df.loc[pd_df["int64_too"] == 1, "string_col"] = "hello" + + pd.testing.assert_frame_equal( + bf_df.to_pandas(), + pd_df, + ) + + +def test_loc_setitem_bool_series_scalar_type_error(scalars_dfs): + if pd.__version__.startswith("1."): + pytest.skip("this loc overload not supported in pandas 1.x.") + + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + + with pytest.raises(TypeError): + bf_df.loc[bf_df["int64_too"] == 1, "string_col"] = 99 + with pytest.raises(TypeError): + pd_df.loc[pd_df["int64_too"] == 1, "string_col"] = 99 + + @pytest.mark.parametrize( ("op"), [ @@ -1749,6 +1972,30 @@ def test_dataframe_prod(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) +def test_df_skew_too_few_values(scalars_dfs): + columns = ["float64_col", "int64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[columns].head(2).skew().to_pandas() + pd_result = scalars_pandas_df[columns].head(2).skew() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + +def test_df_skew(scalars_dfs): + columns = ["float64_col", "int64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[columns].skew().to_pandas() + pd_result = scalars_pandas_df[columns].skew() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + @pytest.mark.parametrize( ("frac", "n", "random_state"), [ @@ -1828,6 +2075,161 @@ def test_df_add_suffix(scalars_df_index, scalars_pandas_df_index, axis): ) +def test_df_columns_filter_items(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("2.0") or pd.__version__.startswith("1."): + pytest.skip("pandas filter items behavior different pre-2.1") + bf_result = scalars_df_index.filter(items=["string_col", "int64_col"]).to_pandas() + + pd_result = scalars_pandas_df_index.filter(items=["string_col", "int64_col"]) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_columns_filter_like(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.filter(like="64_col").to_pandas() + + pd_result = scalars_pandas_df_index.filter(like="64_col") + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_columns_filter_regex(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.filter(regex="^[^_]+$").to_pandas() + + pd_result = scalars_pandas_df_index.filter(regex="^[^_]+$") + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_rows_filter_items(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("2.0") or pd.__version__.startswith("1."): + pytest.skip("pandas filter items behavior different pre-2.1") + bf_result = scalars_df_index.filter(items=[5, 1, 3], axis=0).to_pandas() + + pd_result = scalars_pandas_df_index.filter(items=[5, 1, 3], axis=0) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_rows_filter_like(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.copy().set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.copy().set_index("string_col") + + bf_result = scalars_df_index.filter(like="ello", axis=0).to_pandas() + + pd_result = scalars_pandas_df_index.filter(like="ello", axis=0) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_rows_filter_regex(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.copy().set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.copy().set_index("string_col") + + bf_result = scalars_df_index.filter(regex="^[GH].*", axis=0).to_pandas() + + pd_result = scalars_pandas_df_index.filter(regex="^[GH].*", axis=0) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_reindex_rows_list(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.reindex(index=[5, 1, 3, 99, 1]).to_pandas() + + pd_result = scalars_pandas_df_index.reindex(index=[5, 1, 3, 99, 1]) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_reindex_rows_index(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.reindex( + index=pd.Index([5, 1, 3, 99, 1], name="newname") + ).to_pandas() + + pd_result = scalars_pandas_df_index.reindex( + index=pd.Index([5, 1, 3, 99, 1], name="newname") + ) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_reindex_nonunique(scalars_df_index): + with pytest.raises(ValueError): + # int64_too is non-unique + scalars_df_index.set_index("int64_too").reindex( + index=[5, 1, 3, 99, 1], validate=True + ) + + +def test_df_reindex_columns(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.reindex( + columns=["not_a_col", "int64_col", "int64_too"] + ).to_pandas() + + pd_result = scalars_pandas_df_index.reindex( + columns=["not_a_col", "int64_col", "int64_too"] + ) + + # Pandas uses float64 as default for newly created empty column, bf uses Float64 + pd_result.not_a_col = pd_result.not_a_col.astype(pandas.Float64Dtype()) + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_reindex_like(scalars_df_index, scalars_pandas_df_index): + reindex_target_bf = scalars_df_index.reindex( + columns=["not_a_col", "int64_col", "int64_too"], index=[5, 1, 3, 99, 1] + ) + bf_result = scalars_df_index.reindex_like(reindex_target_bf).to_pandas() + + reindex_target_pd = scalars_pandas_df_index.reindex( + columns=["not_a_col", "int64_col", "int64_too"], index=[5, 1, 3, 99, 1] + ) + pd_result = scalars_pandas_df_index.reindex_like(reindex_target_pd) + + # Pandas uses float64 as default for newly created empty column, bf uses Float64 + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + # Pandas uses float64 as default for newly created empty column, bf uses Float64 + pd_result.not_a_col = pd_result.not_a_col.astype(pandas.Float64Dtype()) + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + def test_df_values(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.values @@ -2035,6 +2437,93 @@ def test_df_duplicated(scalars_df_index, scalars_pandas_df_index, keep, subset): pd.testing.assert_series_equal(pd_series, bf_series, check_dtype=False) +def test_df_to_dict(scalars_df_index, scalars_pandas_df_index): + unsupported = ["numeric_col"] # formatted differently + bf_result = scalars_df_index.drop(columns=unsupported).to_dict() + pd_result = scalars_pandas_df_index.drop(columns=unsupported).to_dict() + + assert bf_result == pd_result + + +def test_df_to_excel(scalars_df_index, scalars_pandas_df_index): + unsupported = ["timestamp_col"] + with tempfile.TemporaryFile() as bf_result_file, tempfile.TemporaryFile() as pd_result_file: + scalars_df_index.drop(columns=unsupported).to_excel(bf_result_file) + scalars_pandas_df_index.drop(columns=unsupported).to_excel(pd_result_file) + bf_result = bf_result_file.read() + pd_result = bf_result_file.read() + + assert bf_result == pd_result + + +def test_df_to_latex(scalars_df_index, scalars_pandas_df_index): + unsupported = ["numeric_col"] # formatted differently + bf_result = scalars_df_index.drop(columns=unsupported).to_latex() + pd_result = scalars_pandas_df_index.drop(columns=unsupported).to_latex() + + assert bf_result == pd_result + + +def test_df_to_records(scalars_df_index, scalars_pandas_df_index): + unsupported = ["numeric_col"] + bf_result = scalars_df_index.drop(columns=unsupported).to_records() + pd_result = scalars_pandas_df_index.drop(columns=unsupported).to_records() + + for bfi, pdi in zip(bf_result, pd_result): + for bfj, pdj in zip(bfi, pdi): + assert pd.isna(bfj) and pd.isna(pdj) or bfj == pdj + + +def test_df_to_string(scalars_df_index, scalars_pandas_df_index): + unsupported = ["numeric_col"] # formatted differently + + bf_result = scalars_df_index.drop(columns=unsupported).to_string() + pd_result = scalars_pandas_df_index.drop(columns=unsupported).to_string() + + assert bf_result == pd_result + + +def test_df_to_markdown(scalars_df_index, scalars_pandas_df_index): + # Nulls have bug from tabulate https://github.com/astanin/python-tabulate/issues/231 + bf_result = scalars_df_index.dropna().to_markdown() + pd_result = scalars_pandas_df_index.dropna().to_markdown() + + assert bf_result == pd_result + + +def test_df_to_pickle(scalars_df_index, scalars_pandas_df_index): + with tempfile.TemporaryFile() as bf_result_file, tempfile.TemporaryFile() as pd_result_file: + scalars_df_index.to_pickle(bf_result_file) + scalars_pandas_df_index.to_pickle(pd_result_file) + bf_result = bf_result_file.read() + pd_result = bf_result_file.read() + + assert bf_result == pd_result + + +def test_df_to_orc(scalars_df_index, scalars_pandas_df_index): + unsupported = [ + "numeric_col", + "bytes_col", + "date_col", + "datetime_col", + "time_col", + "timestamp_col", + "geography_col", + ] + + bf_result_file = tempfile.TemporaryFile() + pd_result_file = tempfile.TemporaryFile() + scalars_df_index.drop(columns=unsupported).to_orc(bf_result_file) + scalars_pandas_df_index.drop(columns=unsupported).reset_index().to_orc( + pd_result_file + ) + bf_result = bf_result_file.read() + pd_result = bf_result_file.read() + + assert bf_result == pd_result + + @pytest.mark.parametrize( ("subset", "normalize", "ascending", "dropna"), [ diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 987368ce771..18741468c5d 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -210,12 +210,14 @@ def test_dataframe_groupby_multi_sum( (lambda x: x.cummax(numeric_only=True)), (lambda x: x.cummin(numeric_only=True)), (lambda x: x.cumprod()), + (lambda x: x.shift(periods=2)), ], ids=[ "cumsum", "cummax", "cummin", "cumprod", + "shift", ], ) def test_dataframe_groupby_analytic( @@ -229,6 +231,30 @@ def test_dataframe_groupby_analytic( pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) +def test_series_groupby_skew(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.groupby("bool_col")["int64_too"].skew().to_pandas() + pd_result = scalars_pandas_df_index.groupby("bool_col")["int64_too"].skew() + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + +def test_dataframe_groupby_skew(scalars_df_index, scalars_pandas_df_index): + col_names = ["float64_col", "int64_col", "bool_col"] + bf_result = scalars_df_index[col_names].groupby("bool_col").skew().to_pandas() + pd_result = scalars_pandas_df_index[col_names].groupby("bool_col").skew() + + pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) + + +def test_dataframe_groupby_diff(scalars_df_index, scalars_pandas_df_index): + col_names = ["float64_col", "int64_col", "string_col"] + bf_result = scalars_df_index[col_names].groupby("string_col").diff(-1) + pd_result = scalars_pandas_df_index[col_names].groupby("string_col").diff(-1) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + def test_dataframe_groupby_getitem( scalars_df_index, scalars_pandas_df_index, diff --git a/tests/system/small/test_ipython.py b/tests/system/small/test_ipython.py index 6725805d9a7..be98ce00674 100644 --- a/tests/system/small/test_ipython.py +++ b/tests/system/small/test_ipython.py @@ -22,7 +22,8 @@ def test_repr_cache(scalars_df_index): # Make sure the df has a new block that the method return value # is not already cached. test_df = scalars_df_index.head() + test_df._block.retrieve_repr_request_results.cache_clear() results = display_formatter.format(test_df) assert results[0].keys() == {"text/plain", "text/html"} - assert test_df._block.retrieve_repr_request_results.cache_info().misses == 1 - assert test_df._block.retrieve_repr_request_results.cache_info().hits == 1 + assert test_df._block.retrieve_repr_request_results.cache_info().misses >= 1 + assert test_df._block.retrieve_repr_request_results.cache_info().hits >= 1 diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 25d1e2ad498..1e38b47b4cc 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -157,7 +157,7 @@ def test_multi_index_getitem_bool(scalars_df_index, scalars_pandas_df_index): ], ids=["level_num", "level_name", "list", "mixed_list"], ) -def test_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level): +def test_df_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level): bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) @@ -167,6 +167,26 @@ def test_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level) pandas.testing.assert_frame_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("level"), + [ + (1), + ("int64_too"), + ([0, 2]), + ([2, "bool_col"]), + ], + ids=["level_num", "level_name", "list", "mixed_list"], +) +def test_series_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, level): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame["string_col"].droplevel(level).to_pandas() + pd_result = pd_frame["string_col"].droplevel(level) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + @pytest.mark.parametrize( ("labels", "level"), [ @@ -198,7 +218,9 @@ def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, lev "num_names_mixed", ], ) -def test_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_index, order): +def test_df_multi_index_reorder_levels( + scalars_df_index, scalars_pandas_df_index, order +): bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) @@ -208,6 +230,51 @@ def test_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_index, o pandas.testing.assert_frame_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("order"), + [ + (1, 0, 2), + (["int64_col", "bool_col", "int64_too"]), + (["int64_col", "bool_col", 0]), + ], + ids=[ + "level_nums", + "level_names", + "num_names_mixed", + ], +) +def test_series_multi_index_reorder_levels( + scalars_df_index, scalars_pandas_df_index, order +): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame["string_col"].reorder_levels(order).to_pandas() + pd_result = pd_frame["string_col"].reorder_levels(order) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + +def test_df_multi_index_swaplevel(scalars_df_index, scalars_pandas_df_index): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame.swaplevel().to_pandas() + pd_result = pd_frame.swaplevel() + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_series_multi_index_swaplevel(scalars_df_index, scalars_pandas_df_index): + bf_frame = scalars_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + pd_frame = scalars_pandas_df_index.set_index(["int64_too", "bool_col", "int64_col"]) + + bf_result = bf_frame["string_col"].swaplevel(0, 2).to_pandas() + pd_result = pd_frame["string_col"].swaplevel(0, 2) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + def test_multi_index_series_groupby(scalars_df_index, scalars_pandas_df_index): bf_frame = scalars_df_index.set_index(["int64_too", "bool_col"]) bf_result = ( @@ -446,6 +513,24 @@ def test_multi_index_series_rename_dict_same_type( ) +def test_multi_index_df_reindex(scalars_df_index, scalars_pandas_df_index): + new_index = pandas.MultiIndex.from_tuples( + [(4, "Hello, World!"), (99, "some_new_string")], + names=["new_index1", "new_index2"], + ) + bf_result = ( + scalars_df_index.set_index(["rowindex_2", "string_col"]) + .reindex(index=new_index) + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index(["rowindex_2", "string_col"]).reindex( + index=new_index + ) + pandas.testing.assert_frame_equal( + bf_result, pd_result, check_dtype=False, check_index_type=False + ) + + # Column Multi-index tests @@ -722,3 +807,76 @@ def test_is_monotonic_decreasing_extra(indexes): bf_result.index.is_monotonic_decreasing == pd_result.index.is_monotonic_decreasing ) + + +def test_column_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pandas.MultiIndex.from_tuples( + zip(["a", "b", "a"], ["c", "d", "e"], ["f", "g", "f"]) + ) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.droplevel(1, axis=1).to_pandas() + pd_result = pd_df.droplevel(1, axis=1) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_column_multi_index_reindex(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "rowindex_2"] + multi_columns = pandas.MultiIndex.from_tuples(zip(["a", "b", "a"], ["a", "b", "b"])) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + new_index = pandas.MultiIndex.from_tuples( + [("z", "a"), ("a", "a")], names=["newname1", "newname2"] + ) + + bf_result = bf_df.reindex(columns=new_index).to_pandas() + + pd_result = pd_df.reindex(columns=new_index) + + # Pandas uses float64 as default for newly created empty column, bf uses Float64 + pd_result[("z", "a")] = pd_result[("z", "a")].astype(pandas.Float64Dtype()) + + pandas.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_column_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pandas.MultiIndex.from_tuples( + zip(["a", "b", "a"], ["c", "d", "e"], ["f", "g", "f"]) + ) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.reorder_levels([-2, -1, 0], axis=1).to_pandas() + pd_result = pd_df.reorder_levels([-2, -1, 0], axis=1) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_column_multi_index_swaplevel(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pandas.MultiIndex.from_tuples( + zip(["a", "b", "a"], ["c", "d", "e"], ["f", "g", "f"]) + ) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.swaplevel(-3, -1, axis=1).to_pandas() + pd_result = pd_df.swaplevel(-3, -1, axis=1) + + pandas.testing.assert_frame_equal(bf_result, pd_result) diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index e451d5c3a2c..a429c6551d6 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -209,3 +209,17 @@ def test_merge_series(scalars_dfs, merge_how): ) assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + +def test_cut(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = pd.cut(scalars_pandas_df["float64_col"], 5, labels=False) + bf_result = bpd.cut(scalars_df["float64_col"], 5, labels=False) + + # make sure the result is a supported dtype + assert bf_result.dtype == bpd.Int64Dtype() + + bf_result = bf_result.to_pandas() + pd_result = pd_result.astype("Int64") + pd.testing.assert_series_equal(bf_result, pd_result) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 07dc892ddc3..d3560540cc9 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -186,6 +186,54 @@ def test_fillna(scalars_dfs): ) +def test_series_replace_scalar_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = ( + scalars_df[col_name].replace("Hello, World!", "Howdy, Planet!").to_pandas() + ) + pd_result = scalars_pandas_df[col_name].replace("Hello, World!", "Howdy, Planet!") + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + +def test_series_replace_regex_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = ( + scalars_df[col_name].replace("^H.l", "Howdy, Planet!", regex=True).to_pandas() + ) + pd_result = scalars_pandas_df[col_name].replace( + "^H.l", "Howdy, Planet!", regex=True + ) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + +def test_series_replace_list_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + col_name = "string_col" + bf_result = ( + scalars_df[col_name] + .replace(["Hello, World!", "T"], "Howdy, Planet!") + .to_pandas() + ) + pd_result = scalars_pandas_df[col_name].replace( + ["Hello, World!", "T"], "Howdy, Planet!" + ) + + pd.testing.assert_series_equal( + pd_result, + bf_result, + ) + + @pytest.mark.parametrize( ("ignore_index",), ( @@ -759,7 +807,6 @@ def test_isin_raise_error(scalars_df_index, scalars_pandas_df_index): ) def test_isin(scalars_dfs, col_name, test_set): scalars_df, scalars_pandas_df = scalars_dfs - print(type(scalars_pandas_df["datetime_col"].iloc[0])) bf_result = scalars_df[col_name].isin(test_set).to_pandas() pd_result = scalars_pandas_df[col_name].isin(test_set).astype("boolean") pd.testing.assert_series_equal( @@ -1506,6 +1553,28 @@ def test_shift(scalars_df_index, scalars_pandas_df_index): ) +def test_series_ffill(scalars_df_index, scalars_pandas_df_index): + col_name = "numeric_col" + bf_result = scalars_df_index[col_name].ffill(limit=1).to_pandas() + pd_result = scalars_pandas_df_index[col_name].ffill(limit=1) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_bfill(scalars_df_index, scalars_pandas_df_index): + col_name = "numeric_col" + bf_result = scalars_df_index[col_name].bfill(limit=2).to_pandas() + pd_result = scalars_pandas_df_index[col_name].bfill(limit=2) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + def test_cumsum_int(scalars_df_index, scalars_pandas_df_index): if pd.__version__.startswith("1."): pytest.skip("Series.cumsum NA mask are different in pandas 1.x.") @@ -1588,7 +1657,7 @@ def test_rank_with_nulls(scalars_df_index, scalars_pandas_df_index, na_option, m ("all",), ], ) -def test_nlargest(scalars_df_index, scalars_pandas_df_index, keep): +def test_series_nlargest(scalars_df_index, scalars_pandas_df_index, keep): col_name = "bool_col" bf_result = scalars_df_index[col_name].nlargest(4, keep=keep).to_pandas() pd_result = scalars_pandas_df_index[col_name].nlargest(4, keep=keep) @@ -1622,6 +1691,25 @@ def test_diff(scalars_df_index, scalars_pandas_df_index, periods): ) +@pytest.mark.parametrize( + ("periods",), + [ + (1,), + (2,), + (-1,), + ], +) +def test_series_pct_change(scalars_df_index, scalars_pandas_df_index, periods): + bf_result = scalars_df_index["int64_col"].pct_change(periods=periods).to_pandas() + # cumsum does not behave well on nullable ints in pandas, produces object type and never ignores NA + pd_result = scalars_pandas_df_index["int64_col"].pct_change(periods=periods) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + @pytest.mark.parametrize( ("keep",), [ @@ -1630,7 +1718,7 @@ def test_diff(scalars_df_index, scalars_pandas_df_index, periods): ("all",), ], ) -def test_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): +def test_series_nsmallest(scalars_df_index, scalars_pandas_df_index, keep): col_name = "bool_col" bf_result = scalars_df_index[col_name].nsmallest(2, keep=keep).to_pandas() pd_result = scalars_pandas_df_index[col_name].nsmallest(2, keep=keep) @@ -1853,6 +1941,91 @@ def test_series_add_suffix(scalars_df_index, scalars_pandas_df_index): ) +def test_series_filter_items(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("2.0") or pd.__version__.startswith("1."): + pytest.skip("pandas filter items behavior different pre-2.1") + bf_result = scalars_df_index["float64_col"].filter(items=[5, 1, 3]).to_pandas() + + pd_result = scalars_pandas_df_index["float64_col"].filter(items=[5, 1, 3]) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_filter_like(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.copy().set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.copy().set_index("string_col") + + bf_result = scalars_df_index["float64_col"].filter(like="ello").to_pandas() + + pd_result = scalars_pandas_df_index["float64_col"].filter(like="ello") + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_filter_regex(scalars_df_index, scalars_pandas_df_index): + scalars_df_index = scalars_df_index.copy().set_index("string_col") + scalars_pandas_df_index = scalars_pandas_df_index.copy().set_index("string_col") + + bf_result = scalars_df_index["float64_col"].filter(regex="^[GH].*").to_pandas() + + pd_result = scalars_pandas_df_index["float64_col"].filter(regex="^[GH].*") + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_reindex(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["float64_col"].reindex(index=[5, 1, 3, 99, 1]).to_pandas() + ) + + pd_result = scalars_pandas_df_index["float64_col"].reindex(index=[5, 1, 3, 99, 1]) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + +def test_series_reindex_nonunique(scalars_df_index): + with pytest.raises(ValueError): + # int64_too is non-unique + scalars_df_index.set_index("int64_too")["float64_col"].reindex( + index=[5, 1, 3, 99, 1], validate=True + ) + + +def test_series_reindex_like(scalars_df_index, scalars_pandas_df_index): + bf_reindex_target = scalars_df_index["float64_col"].reindex(index=[5, 1, 3, 99, 1]) + bf_result = ( + scalars_df_index["int64_too"].reindex_like(bf_reindex_target).to_pandas() + ) + + pd_reindex_target = scalars_pandas_df_index["float64_col"].reindex( + index=[5, 1, 3, 99, 1] + ) + pd_result = scalars_pandas_df_index["int64_too"].reindex_like(pd_reindex_target) + + # Pandas uses int64 instead of Int64 (nullable) dtype. + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + def test_where_with_series(scalars_df_index, scalars_pandas_df_index): bf_result = ( scalars_df_index["int64_col"] diff --git a/tests/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py new file mode 100644 index 00000000000..8d4932a3c3c --- /dev/null +++ b/tests/unit/ml/test_golden_sql.py @@ -0,0 +1,47 @@ +# 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. + +from unittest import mock + +import pytest_mock + +import bigframes +from bigframes.ml import linear_model +import bigframes.pandas as bpd + + +def test_linear_regression_default_fit(mocker: pytest_mock.MockerFixture): + mock_session = mock.create_autospec(spec=bigframes.Session) + + mock_X = mock.create_autospec(spec=bpd.DataFrame) + mock_X._get_block().expr._session = mock_session + + mock_y = mock.create_autospec(spec=bpd.DataFrame) + mock_y.columns.tolist.return_value = ["input_label_column"] + + mock_X.join(mock_y).sql = "input_dataframe_sql" + + # return values we don't care about, but need to provide to continue the program + mock_session._start_query.return_value = (None, mock.MagicMock()) + + mocker.patch( + "bigframes.ml.core._create_temp_model_name", return_value="temp_model_name" + ) + + model = linear_model.LinearRegression() + model.fit(mock_X, mock_y) + + mock_session._start_query.assert_called_once_with( + 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=True,\n INPUT_LABEL_COLS=["input_label_column"])\nAS input_dataframe_sql' + ) diff --git a/tests/unit/ml/test_pipeline.py b/tests/unit/ml/test_pipeline.py index 27706a1a07f..ed5c621b1df 100644 --- a/tests/unit/ml/test_pipeline.py +++ b/tests/unit/ml/test_pipeline.py @@ -18,38 +18,35 @@ import sklearn.pipeline as sklearn_pipeline # type: ignore import sklearn.preprocessing as sklearn_preprocessing # type: ignore -import bigframes.ml.compose -import bigframes.ml.linear_model -import bigframes.ml.pipeline -import bigframes.ml.preprocessing +from bigframes.ml import compose, forecasting, linear_model, pipeline, preprocessing def test_pipeline_repr(): - pipeline = bigframes.ml.pipeline.Pipeline( + pl = pipeline.Pipeline( [ ( "preproc", - bigframes.ml.compose.ColumnTransformer( + compose.ColumnTransformer( [ ( "onehot", - bigframes.ml.preprocessing.OneHotEncoder(), + preprocessing.OneHotEncoder(), "species", ), ( "scale", - bigframes.ml.preprocessing.StandardScaler(), + preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), ] ), ), - ("linreg", bigframes.ml.linear_model.LinearRegression()), + ("linreg", linear_model.LinearRegression()), ] ) assert ( - pipeline.__repr__() + pl.__repr__() == """Pipeline(steps=[('preproc', ColumnTransformer(transformers=[('onehot', OneHotEncoder(), 'species'), @@ -62,29 +59,29 @@ def test_pipeline_repr(): @pytest.mark.skipif(sklearn_pipeline is None, reason="requires sklearn") def test_pipeline_repr_matches_sklearn(): - bf_pipeline = bigframes.ml.pipeline.Pipeline( + bf_pl = pipeline.Pipeline( [ ( "preproc", - bigframes.ml.compose.ColumnTransformer( + compose.ColumnTransformer( [ ( "onehot", - bigframes.ml.preprocessing.OneHotEncoder(), + preprocessing.OneHotEncoder(), "species", ), ( "scale", - bigframes.ml.preprocessing.StandardScaler(), + preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), ] ), ), - ("linreg", bigframes.ml.linear_model.LinearRegression()), + ("linreg", linear_model.LinearRegression()), ] ) - sk_pipeline = sklearn_pipeline.Pipeline( + sk_pl = sklearn_pipeline.Pipeline( [ ( "preproc", @@ -107,4 +104,17 @@ def test_pipeline_repr_matches_sklearn(): ] ) - assert bf_pipeline.__repr__() == sk_pipeline.__repr__() + assert bf_pl.__repr__() == sk_pl.__repr__() + + +def test_pipeline_arima_plus_not_implemented(): + with pytest.raises(NotImplementedError): + pipeline.Pipeline( + [ + ( + "transform", + preprocessing.StandardScaler(), + ), + ("estimator", forecasting.ARIMAPlus()), + ] + ) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index e01638e22e1..8f3e0beb0e4 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -13,9 +13,12 @@ # limitations under the License. import ibis +import ibis.expr.types as ibis_types import pandas import bigframes.core as core +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops from . import resources @@ -46,6 +49,42 @@ def test_arrayvalue_constructor_from_ibis_table_adds_all_columns(): assert len(actual.columns) == 3 +def test_arrayvalue_with_get_column_type(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + col1_type = value.get_column_type("col1") + col2_type = value.get_column_type("col2") + col3_type = value.get_column_type("col3") + assert isinstance(col1_type, pandas.Int64Dtype) + assert isinstance(col2_type, pandas.StringDtype) + assert isinstance(col3_type, pandas.Float64Dtype) + + +def test_arrayvalue_with_get_column(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + col1 = value.get_column("col1") + assert isinstance(col1, ibis_types.Value) + assert col1.get_name() == "col1" + assert col1.type().is_int64() + + def test_arrayvalue_to_ibis_expr_with_projection(): value = resources.create_arrayvalue( pandas.DataFrame( @@ -69,3 +108,133 @@ def test_arrayvalue_to_ibis_expr_with_projection(): assert actual.columns[0] == "int64_col" assert actual.columns[1] == "literals" assert actual.columns[2] == "string_col" + + +def test_arrayvalues_to_ibis_expr_with_get_column(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.get_column("col1") + assert expr.get_name() == "col1" + assert expr.type().is_int64() + + +def test_arrayvalues_to_ibis_expr_with_concat(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.concat([value]) + actual = expr.to_ibis_expr() + assert len(actual.columns) == 3 + # TODO(ashleyxu, b/299631930): test out the union expression + assert actual.columns[0] == "column_0" + assert actual.columns[1] == "column_1" + assert actual.columns[2] == "column_2" + + +def test_arrayvalues_to_ibis_expr_with_project_unary_op(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.project_unary_op("col1", ops.AsTypeOp("string")) + assert value.columns[0].type().is_int64() + assert expr.columns[0].type().is_string() + + +def test_arrayvalues_to_ibis_expr_with_project_binary_op(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": [0.2, 0.3, 0.4], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.project_binary_op("col2", "col3", ops.add_op, "col4") + assert expr.columns[3].type().is_float64() + actual = expr.to_ibis_expr() + assert len(expr.columns) == 4 + assert actual.columns[3] == "col4" + + +def test_arrayvalues_to_ibis_expr_with_project_ternary_op(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": [0.2, 0.3, 0.4], + "col3": [True, False, False], + "col4": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.project_ternary_op("col2", "col3", "col4", ops.where_op, "col5") + assert expr.columns[4].type().is_float64() + actual = expr.to_ibis_expr() + assert len(expr.columns) == 5 + assert actual.columns[4] == "col5" + + +def test_arrayvalue_to_ibis_expr_with_aggregate(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.aggregate( + aggregations=(("col1", agg_ops.sum_op, "col4"),), + by_column_ids=["col1"], + dropna=False, + ) + actual = expr.to_ibis_expr() + assert len(expr.columns) == 2 + assert actual.columns[0] == "col1" + assert actual.columns[1] == "col4" + assert expr.columns[1].type().is_int64() + + +def test_arrayvalue_to_ibis_expr_with_corr_aggregate(): + value = resources.create_arrayvalue( + pandas.DataFrame( + { + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [0.1, 0.2, 0.3], + } + ), + total_ordering_columns=["col1"], + ) + expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]) + actual = expr.to_ibis_expr() + assert len(expr.columns) == 1 + assert actual.columns[0] == "col4" + assert expr.columns[0].type().is_float64() diff --git a/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py index d209284ab7d..a4e61ca0f91 100644 --- a/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py +++ b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py @@ -12,8 +12,20 @@ def _approx_quantiles(translator, op: vendored_ibis_ops.ApproximateMultiQuantile return f"APPROX_QUANTILES({arg}, {num_bins})" +def _first_non_null_value(translator, op: vendored_ibis_ops.FirstNonNullValue): + arg = translator.translate(op.arg) + return f"FIRST_VALUE({arg} IGNORE NULLS)" + + +def _last_non_null_value(translator, op: vendored_ibis_ops.LastNonNullValue): + arg = translator.translate(op.arg) + return f"LAST_VALUE({arg} IGNORE NULLS)" + + patched_ops = { vendored_ibis_ops.ApproximateMultiQuantile: _approx_quantiles, + vendored_ibis_ops.FirstNonNullValue: _first_non_null_value, + vendored_ibis_ops.LastNonNullValue: _last_non_null_value, } OPERATION_REGISTRY.update(patched_ops) diff --git a/third_party/bigframes_vendored/ibis/expr/operations/__init__.py b/third_party/bigframes_vendored/ibis/expr/operations/__init__.py index f3ab753a3b1..1612d9c12e5 100644 --- a/third_party/bigframes_vendored/ibis/expr/operations/__init__.py +++ b/third_party/bigframes_vendored/ibis/expr/operations/__init__.py @@ -1,4 +1,5 @@ # Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/__init__.py from __future__ import annotations +from third_party.bigframes_vendored.ibis.expr.operations.analytic import * # noqa: F403 from third_party.bigframes_vendored.ibis.expr.operations.reductions import * # noqa: F403 diff --git a/third_party/bigframes_vendored/ibis/expr/operations/analytic.py b/third_party/bigframes_vendored/ibis/expr/operations/analytic.py new file mode 100644 index 00000000000..038987cac93 --- /dev/null +++ b/third_party/bigframes_vendored/ibis/expr/operations/analytic.py @@ -0,0 +1,26 @@ +# Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/analytic.py + +from __future__ import annotations + +from ibis.expr.operations.analytic import Analytic +import ibis.expr.rules as rlz + + +class FirstNonNullValue(Analytic): + """Retrieve the first element.""" + + arg = rlz.column(rlz.any) + output_dtype = rlz.dtype_like("arg") + + +class LastNonNullValue(Analytic): + """Retrieve the last element.""" + + arg = rlz.column(rlz.any) + output_dtype = rlz.dtype_like("arg") + + +__all__ = [ + "FirstNonNullValue", + "LastNonNullValue", +] diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 8c81b23b6ca..113c6547a0e 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -13,7 +13,7 @@ from typing import Iterable, Literal, Mapping, Optional, Sequence, Union -import numpy +import numpy as np from bigframes import constants from third_party.bigframes_vendored.pandas.core.generic import NDFrame @@ -56,7 +56,7 @@ def axes(self) -> list: return [self.index, self.columns] @property - def values(self) -> numpy.ndarray: + def values(self) -> np.ndarray: """Return the values of DataFrame in the form of a NumPy array. Args: @@ -72,9 +72,7 @@ def values(self) -> numpy.ndarray: # ---------------------------------------------------------------------- # IO methods (to / from other formats) - def to_numpy( - self, dtype=None, copy=False, na_value=None, **kwargs - ) -> numpy.ndarray: + def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarray: """ Convert the DataFrame to a NumPy array. @@ -154,6 +152,250 @@ def to_parquet( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def to_dict( + self, + orient: Literal[ + "dict", "list", "series", "split", "tight", "records", "index" + ] = "dict", + into: type[dict] = dict, + **kwargs, + ) -> dict | list[dict]: + """ + Convert the DataFrame to a dictionary. + + The type of the key-value pairs can be customized with the parameters + (see below). + + Args: + orient (str {'dict', 'list', 'series', 'split', 'tight', 'records', 'index'}): + Determines the type of the values of the dictionary. + 'dict' (default) : dict like {column -> {index -> value}}. + 'list' : dict like {column -> [values]}. + 'series' : dict like {column -> Series(values)}. + split' : dict like {'index' -> [index], 'columns' -> [columns], 'data' -> [values]}. + 'tight' : dict like {'index' -> [index], 'columns' -> [columns], 'data' -> [values], + 'index_names' -> [index.names], 'column_names' -> [column.names]}. + 'records' : list like [{column -> value}, ... , {column -> value}]. + 'index' : dict like {index -> {column -> value}}. + into (class, default dict): + The collections.abc.Mapping subclass used for all Mappings + in the return value. Can be the actual class or an empty + instance of the mapping type you want. If you want a + collections.defaultdict, you must pass it initialized. + + index (bool, default True): + Whether to include the index item (and index_names item if `orient` + is 'tight') in the returned dictionary. Can only be ``False`` + when `orient` is 'split' or 'tight'. + + Returns: + dict or list of dict: Return a collections.abc.Mapping object representing the DataFrame. + The resulting transformation depends on the `orient` parameter. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_excel(self, excel_writer, sheet_name: str = "Sheet1", **kwargs) -> None: + """ + Write DataFrame to an Excel sheet. + + To write a single DataFrame to an Excel .xlsx file it is only necessary to + specify a target file name. To write to multiple sheets it is necessary to + create an `ExcelWriter` object with a target file name, and specify a sheet + in the file to write to. + + Multiple sheets may be written to by specifying unique `sheet_name`. + With all data written to the file it is necessary to save the changes. + Note that creating an `ExcelWriter` object with a file name that already + exists will result in the contents of the existing file being erased. + + Args: + excel_writer (path-like, file-like, or ExcelWriter object): + File path or existing ExcelWriter. + sheet_name (str, default 'Sheet1'): + Name of sheet which will contain DataFrame. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_latex( + self, buf=None, columns=None, header=True, index=True, **kwargs + ) -> str | None: + r""" + Render object to a LaTeX tabular, longtable, or nested table. + + Requires ``\usepackage{{booktabs}}``. The output can be copy/pasted + into a main LaTeX document or read from an external file + with ``\input{{table.tex}}``. + + Args: + buf (str, Path or StringIO-like, optional, default None): + Buffer to write to. If None, the output is returned as a string. + columns (list of label, optional): + The subset of columns to write. Writes all columns by default. + header (bool or list of str, default True): + Write out the column names. If a list of strings is given, + it is assumed to be aliases for the column names. + index (bool, default True): + Write row names (index). + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_records( + self, index: bool = True, column_dtypes=None, index_dtypes=None + ) -> np.recarray: + """ + Convert DataFrame to a NumPy record array. + + Index will be included as the first field of the record array if + requested. + + Args: + index (bool, default True): + Include index in resulting record array, stored in 'index' + field or using the index label, if set. + column_dtypes (str, type, dict, default None): + If a string or type, the data type to store all columns. If + a dictionary, a mapping of column names and indices (zero-indexed) + to specific data types. + index_dtypes (str, type, dict, default None): + If a string or type, the data type to store all index levels. If + a dictionary, a mapping of index level names and indices + (zero-indexed) to specific data types. + + This mapping is applied only if `index=True`. + + Returns: + np.recarray: NumPy ndarray with the DataFrame labels as fields and each row + of the DataFrame as entries. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_string( + self, + buf=None, + columns: Sequence[str] | None = None, + col_space=None, + header: bool | Sequence[str] = True, + index: bool = True, + na_rep: str = "NaN", + formatters=None, + float_format=None, + sparsify: bool | None = None, + index_names: bool = True, + justify: str | None = None, + max_rows: int | None = None, + max_cols: int | None = None, + show_dimensions: bool = False, + decimal: str = ".", + line_width: int | None = None, + min_rows: int | None = None, + max_colwidth: int | None = None, + encoding: str | None = None, + ): + """Render a DataFrame to a console-friendly tabular output. + + Args: + buf (str, Path or StringIO-like, optional, default None): + Buffer to write to. If None, the output is returned as a string. + columns (sequence, optional, default None): + The subset of columns to write. Writes all columns by default. + col_space (int, list or dict of int, optional): + The minimum width of each column. + header (bool or sequence, optional): + Write out the column names. If a list of strings is given, it is assumed to be aliases for the column names. + index (bool, optional, default True): + Whether to print index (row) labels. + na_rep (str, optional, default 'NaN'): + String representation of NAN to use. + formatters (list, tuple or dict of one-param. functions, optional): + Formatter functions to apply to columns' elements by position or + name. + The result of each function must be a unicode string. + List/tuple must be of length equal to the number of columns. + float_format (one-parameter function, optional, default None): + Formatter function to apply to columns' elements if they are + floats. The result of this function must be a unicode string. + sparsify (bool, optional, default True): + Set to False for a DataFrame with a hierarchical index to print + every multiindex key at each row. + index_names (bool, optional, default True): + Prints the names of the indexes. + justify (str, default None): + How to justify the column labels. If None uses the option from + the print configuration (controlled by set_option), 'right' out + of the box. Valid values are, 'left', 'right', 'center', 'justify', + 'justify-all', 'start', 'end', 'inherit', 'match-parent', 'initial', + 'unset'. + max_rows (int, optional): + Maximum number of rows to display in the console. + min_rows (int, optional): + The number of rows to display in the console in a truncated repr + (when number of rows is above `max_rows`). + max_cols (int, optional): + Maximum number of columns to display in the console. + show_dimensions (bool, default False): + Display DataFrame dimensions (number of rows by number of columns). + decimal (str, default '.'): + Character recognized as decimal separator, e.g. ',' in Europe. + line_width (int, optional): + Width to wrap a line in characters. + max_colwidth (int, optional): + Max width to truncate each column in characters. By default, no limit. + encoding (str, default "utf-8"): + Set character encoding. + + Returns: + str or None: If buf is None, returns the result as a string. Otherwise returns + None. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_markdown( + self, + buf=None, + mode: str = "wt", + index: bool = True, + **kwargs, + ): + """Print DataFrame in Markdown-friendly format. + + Args: + buf (str, Path or StringIO-like, optional, default None): + Buffer to write to. If None, the output is returned as a string. + mode (str, optional): + Mode in which file is opened. + index (bool, optional, default True): + Add index (row) labels. + **kwargs + These parameters will be passed to `tabulate `_. + + Returns: + DataFrame in Markdown-friendly format. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_pickle(self, path, **kwargs) -> None: + """Pickle (serialize) object to file. + + Args: + path (str): + File path where the pickled object will be stored. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def to_orc(self, path=None, **kwargs) -> bytes | None: + """ + Write a DataFrame to the ORC format. + + Args: + path (str, file-like object or None, default None): + If a string, it will be used as Root Directory path + when writing a partitioned dataset. By file-like object, + we refer to objects with a write() method, such as a file handle + (e.g. via builtin open function). If path is None, + a bytes object is returned. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Unsorted @@ -185,6 +427,53 @@ def assign(self, **kwargs) -> DataFrame: # ---------------------------------------------------------------------- # Reindexing and alignment + def reindex( + self, + labels=None, + *, + index=None, + columns=None, + axis=None, + ): + """Conform DataFrame to new index with optional filling logic. + + Places NA in locations having no value in the previous index. A new object + is produced. + + Args: + labels (array-like, optional): + New labels / index to conform the axis specified by 'axis' to. + index (array-like, optional): + New labels for the index. Preferably an Index object to avoid + duplicating data. + columns (array-like, optional): + New labels for the columns. Preferably an Index object to avoid + duplicating data. + axis (int or str, optional): + Axis to target. Can be either the axis name ('index', 'columns') + or number (0, 1). + Returns: + DataFrame: DataFrame with changed index. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def reindex_like(self, other): + """Return an object with matching indices as other object. + + Conform the object to the same index on all axes. Optional + filling logic, placing Null in locations having no value + in the previous index. + + Args: + other (Object of the same data type): + Its row and column indices are used to define the new indices + of this object. + + Returns: + Series or DataFrame: Same type as caller, but with changed indices on each axis. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def drop( self, labels=None, *, axis=0, index=None, columns=None, level=None ) -> DataFrame | None: @@ -276,7 +565,9 @@ def set_index( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: + def reorder_levels( + self, order: Sequence[int | str], axis: str | int = 0 + ) -> DataFrame: """ Rearrange index levels using input order. May not drop or duplicate levels. @@ -284,13 +575,33 @@ def reorder_levels(self, order: Sequence[int | str]) -> DataFrame: order (list of int or list of str): List representing new level order. Reference level by number (position) or by key (label). + axis ({0 or 'index', 1 or 'columns'}, default 0): + Where to reorder levels. Returns: DataFrame: DataFrame of rearranged index. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def droplevel(self, level): + def swaplevel(self, i, j, axis: str | int = 0) -> DataFrame: + """ + Swap levels i and j in a :class:`MultiIndex`. + + Default is to swap the two innermost levels of the index. + + Args: + i, j (int or str): + Levels of the indices to be swapped. Can pass level name as string. + axis ({0 or 'index', 1 or 'columns'}, default 0): + The axis to swap levels on. 0 or 'index' for row-wise, 1 or + 'columns' for column-wise. + + Returns: + DataFrame: DataFrame with levels swapped in MultiIndex. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def droplevel(self, level, axis: str | int = 0): """ Return DataFrame with requested index / column level(s) removed. @@ -299,6 +610,11 @@ def droplevel(self, level): If a string is given, must be the name of a level If list-like, elements must be names or positional indexes of levels. + axis ({0 or 'index', 1 or 'columns'}, default 0): + Axis along which the level(s) is removed: + + * 0 or 'index': remove level(s) in column. + * 1 or 'columns': remove level(s) in row. Returns: DataFrame: DataFrame with requested index / column level(s) removed. """ @@ -889,6 +1205,54 @@ def rpow(self, other, axis: str | int = "columns") -> DataFrame: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def combine( + self, other, func, fill_value=None, overwrite: bool = True + ) -> DataFrame: + """Perform column-wise combine with another DataFrame. + + Combines a DataFrame with `other` DataFrame using `func` + to element-wise combine columns. The row and column indexes of the + resulting DataFrame will be the union of the two. + + Args: + other (DataFrame): + The DataFrame to merge column-wise. + func (function): + Function that takes two series as inputs and return a Series or a + scalar. Used to merge the two dataframes column by columns. + fill_value (scalar value, default None): + The value to fill NaNs with prior to passing any column to the + merge func. + overwrite (bool, default True): + If True, columns in `self` that do not exist in `other` will be + overwritten with NaNs. + + Returns: + DataFrame: Combination of the provided DataFrames. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def combine_first(self, other) -> DataFrame: + """ + Update null elements with value in the same location in `other`. + + Combine two DataFrame objects by filling null values in one DataFrame + with non-null values from other DataFrame. The row and column indexes + of the resulting DataFrame will be the union of the two. The resulting + dataframe contains the 'first' dataframe values and overrides the + second one values where both first.loc[index, col] and + second.loc[index, col] are not missing values, upon calling + first.combine_first(second). + + Args: + other (DataFrame): + Provided DataFrame to use to fill null values. + + Returns: + DataFrame: The result of combining the provided DataFrame with the other object. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Data reshaping @@ -1191,6 +1555,20 @@ def var(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def skew(self, *, numeric_only: bool = False): + """Return unbiased skew over requested axis. + + Normalized by N-1. + + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def std(self, *, numeric_only: bool = False): """Return sample standard deviation over requested axis. @@ -1222,6 +1600,76 @@ def count(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def nlargest(self, n: int, columns, keep: str = "first"): + """ + Return the first `n` rows ordered by `columns` in descending order. + + Return the first `n` rows with the largest values in `columns`, in + descending order. The columns that are not specified are returned as + well, but not used for ordering. + + This method is equivalent to + ``df.sort_values(columns, ascending=False).head(n)``, but more + performant. + + Args: + n (int): + Number of rows to return. + columns (label or list of labels): + Column label(s) to order by. + keep ({'first', 'last', 'all'}, default 'first'): + Where there are duplicate values: + + - ``first`` : prioritize the first occurrence(s) + - ``last`` : prioritize the last occurrence(s) + - ``all`` : do not drop any duplicates, even it means + selecting more than `n` items. + + Returns: + DataFrame: The first `n` rows ordered by the given columns in descending order. + + .. note:: + This function cannot be used with all column types. For example, when + specifying columns with `object` or `category` dtypes, ``TypeError`` is + raised. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def nsmallest(self, n: int, columns, keep: str = "first"): + """ + Return the first `n` rows ordered by `columns` in ascending order. + + Return the first `n` rows with the smallest values in `columns`, in + ascending order. The columns that are not specified are returned as + well, but not used for ordering. + + This method is equivalent to + ``df.sort_values(columns, ascending=True).head(n)``, but more + performant. + + Args: + n (int): + Number of rows to return. + columns (label or list of labels): + Column label(s) to order by. + keep ({'first', 'last', 'all'}, default 'first'): + Where there are duplicate values: + + - ``first`` : prioritize the first occurrence(s) + - ``last`` : prioritize the last occurrence(s) + - ``all`` : do not drop any duplicates, even it means + selecting more than `n` items. + + Returns: + DataFrame: The first `n` rows ordered by the given columns in ascending order. + + .. note:: + This function cannot be used with all column types. For example, when + specifying columns with `object` or `category` dtypes, ``TypeError`` is + raised. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def nunique(self): """ Count number of distinct elements in specified axis. @@ -1271,6 +1719,25 @@ def cumprod(self) -> DataFrame: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def diff( + self, + periods: int = 1, + ) -> NDFrame: + """First discrete difference of element. + + Calculates the difference of a DataFrame element compared with another + element in the DataFrame (default is element in previous row). + + Args: + periods (int, default 1): + Periods to shift for calculating difference, accepts negative + values. + + Returns: + bigframes.dataframe.DataFrame: First differences of the Series. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def agg(self, func): """ Aggregate using one or more operations over the specified axis. diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 56d3b2434f0..7d496891b0b 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -335,6 +335,41 @@ def copy(self): # ---------------------------------------------------------------------- # Action Methods + def ffill(self, *, limit: Optional[int] = None): + """Fill NA/NaN values by propagating the last valid observation to next valid. + + Args: + limit : int, default None + If method is specified, this is the maximum number of consecutive + NaN values to forward/backward fill. In other words, if there is + a gap with more than this number of consecutive NaNs, it will only + be partially filled. If method is not specified, this is the + maximum number of entries along the entire axis where NaNs will be + filled. Must be greater than 0 if not None. + + + Returns: + Series/DataFrame or None: Object with missing values filled. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def bfill(self, *, limit: Optional[int] = None): + """Fill NA/NaN values by using the next valid observation to fill the gap. + + Args: + limit : int, default None + If method is specified, this is the maximum number of consecutive + NaN values to forward/backward fill. In other words, if there is + a gap with more than this number of consecutive NaNs, it will only + be partially filled. If method is not specified, this is the + maximum number of entries along the entire axis where NaNs will be + filled. Must be greater than 0 if not None. + + Returns: + Series/DataFrame or None: Object with missing values filled. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def isna(self) -> NDFrame: """Detect missing values. @@ -367,6 +402,36 @@ def notna(self) -> NDFrame: notnull = notna + def filter( + self, + items=None, + like: str | None = None, + regex: str | None = None, + axis=None, + ) -> NDFrame: + """ + Subset the dataframe rows or columns according to the specified index labels. + + Note that this routine does not filter a dataframe on its + contents. The filter is applied to the labels of the index. + + Args: + items (list-like): + Keep labels from axis which are in items. + like (str): + Keep labels from axis for which "like in label == True". + regex (str (regular expression)): + Keep labels from axis for which re.search(regex, label) == True. + axis ({0 or 'index', 1 or 'columns', None}, default None): + The axis to filter on, expressed either as an index (int) + or axis name (str). By default this is the info axis, 'columns' for + DataFrame. For `Series` this parameter is unused and defaults to `None`. + + Returns: + same type as input object + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def shift( self, periods: int = 1, @@ -384,6 +449,30 @@ def shift( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def pct_change(self, periods: int = 1): + """ + Fractional change between the current and a prior element. + + Computes the fractional change from the immediately previous row by + default. This is useful in comparing the fraction of change in a time + series of elements. + + .. note:: + + Despite the name of this method, it calculates fractional change + (also known as per unit change or relative change) and not + percentage change. If you need the percentage change, multiply + these values by 100. + + Args: + periods (int, default 1): + Periods to shift for forming percent change. + + Returns: + Series or DataFrame: The same type as the calling object. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def rank( self, axis=0, diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index 9271da8a5ea..7849a3afd52 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -124,6 +124,26 @@ def var( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def skew( + self, + *, + numeric_only: bool = False, + ): + """ + Return unbiased skew within groups. + + Normalized by N-1. + + Args: + numeric_only (bool, default False): + Include only `float`, `int` or `boolean` data. + + Returns: + Series or DataFrame + Variance of values within each group. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def sum( self, numeric_only: bool = False, diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py index 404a99809c0..864007b7749 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/base.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -21,6 +21,16 @@ def shape(self): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property + def nlevels(self) -> int: + """Number of levels.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def is_unique(self) -> bool: + """Return if the index has unique values.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def to_numpy(self, dtype): """ A NumPy ndarray representing the values in this Series or Index. diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 76fb46a7003..c6dd973372c 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -758,6 +758,41 @@ def groupby( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def reindex(self, index=None): + """ + Conform Series to new index with optional filling logic. + + Places NA/NaN in locations having no value in the previous index. A new object + is produced unless the new index is equivalent to the current one and + ``copy=False``. + + Args: + index (array-like, optional): + New labels for the index. Preferably an Index object to avoid + duplicating data. + + Returns: + Series: Series with changed index. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def reindex_like(self, other): + """Return an object with matching indices as other object. + + Conform the object to the same index on all axes. Optional + filling logic, placing Null in locations having no value + in the previous index. + + Args: + other (Object of the same data type): + Its row and column indices are used to define the new indices + of this object. + + Returns: + Series or DataFrame: Same type as caller, but with changed indices on each axis. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def drop( self, labels=None, *, axis=0, index=None, columns=None, level=None ) -> Series | None: @@ -790,7 +825,7 @@ def drop( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def reorder_levels(self, order: Sequence) -> Series: + def reorder_levels(self, order: Sequence, axis) -> Series: """ Rearrange index levels using input order. @@ -800,12 +835,31 @@ def reorder_levels(self, order: Sequence) -> Series: order (list of int representing new level order): Reference level by number or key. + axis ({0 or 'index', 1 or 'columns'}, default 0): + For `Series` this parameter is unused and defaults to 0. + + Returns: type of caller (new object) """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def droplevel(self, level): + def swaplevel(self, i, j): + """ + Swap levels i and j in a `MultiIndex`. + + Default is to swap the two innermost levels of the index. + + Args: + i, j (int or str): + Levels of the indices to be swapped. Can pass level name as string. + + Returns: + Series: Series with levels swapped in MultiIndex + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def droplevel(self, level, axis): """ Return Series with requested index / column level(s) removed. @@ -815,6 +869,9 @@ def droplevel(self, level): If list-like, elements must be names or positional indexes of levels. + axis ({0 or 'index', 1 or 'columns'}, default 0): + For `Series` this parameter is unused and defaults to 0. + Returns: Series with requested index / column level(s) removed. """ @@ -836,6 +893,69 @@ def fillna( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def replace( + self, + to_replace, + value=None, + ) -> Series | None: + """ + Replace values given in `to_replace` with `value`. + + Values of the Series/DataFrame are replaced with other values dynamically. + This differs from updating with ``.loc`` or ``.iloc``, which require + you to specify a location to update with some value. + + Args: + to_replace (str, regex, list, int, float or None): + How to find the values that will be replaced. + + * numeric, str or regex: + + - numeric: numeric values equal to `to_replace` will be + replaced with `value` + - str: string exactly matching `to_replace` will be replaced + with `value` + - regex: regexs matching `to_replace` will be replaced with + `value` + + * list of str, regex, or numeric: + + - First, if `to_replace` and `value` are both lists, they + **must** be the same length. + - Second, if ``regex=True`` then all of the strings in **both** + lists will be interpreted as regexs otherwise they will match + directly. This doesn't matter much for `value` since there + are only a few possible substitution regexes you can use. + - str, regex and numeric rules apply as above. + + value (scalar, default None): + Value to replace any values matching `to_replace` with. + For a DataFrame a dict of values can be used to specify which + value to use for each column (columns not in the dict will not be + filled). Regular expressions, strings and lists or dicts of such + objects are also allowed. + regex (bool, default False): + Whether to interpret `to_replace` and/or `value` as regular + expressions. If this is ``True`` then `to_replace` *must* be a + string. + + Returns: + Series/DataFrame: Object after replacement. + + Raises: + TypeError: + * If `to_replace` is not a scalar, array-like, ``dict``, or ``None`` + * If `to_replace` is a ``dict`` and `value` is not a ``list``, + ``dict``, ``ndarray``, or ``Series`` + * If `to_replace` is ``None`` and `regex` is not compilable + into a regular expression or is a list, dict, ndarray, or + Series. + * When replacing multiple ``bool`` or ``datetime64`` objects and + the arguments to `to_replace` does not match the type of the + value being replaced + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def dropna(self, *, axis=0, inplace: bool = False, how=None) -> Series: """ Return a new Series with missing values removed. diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py index 847ad06f750..42868ce51f2 100644 --- a/third_party/bigframes_vendored/sklearn/base.py +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -144,6 +144,7 @@ def fit_transform(self, X, y=None): bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_features_new) Transformed DataFrame. """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) class MetaEstimatorMixin: diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index ff1c04edbe3..ece62dc1470 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -49,7 +49,6 @@ def fit( self, X, y=None, - transforms: Optional[List[str]] = None, ): """Compute k-means clustering. @@ -58,10 +57,6 @@ def fit( DataFrame of shape (n_samples, n_features). Training data. y (default None): Not used, present here for API consistency by convention. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: KMeans: Fitted Estimator. diff --git a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py index 85feab00245..97fee5a501e 100644 --- a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py +++ b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py @@ -49,10 +49,6 @@ def fit(self, X, y=None): y (default None): Ignored. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: PCA: Fitted estimator. """ diff --git a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py index 79224a772d6..6be41bf9aa5 100644 --- a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py +++ b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py @@ -54,10 +54,6 @@ def fit(self, X, y): Series or DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: Fitted Estimator. diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py index 8141da4e3bc..81b4fca1572 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_base.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -74,7 +74,6 @@ def fit( self, X, y, - transforms: Optional[List[str]] = None, ): """Fit linear model. @@ -86,10 +85,6 @@ def fit( Series or DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: LinearRegression: Fitted Estimator. """ diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py index a06035eef66..133dc4498e4 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py @@ -37,7 +37,6 @@ def fit( self, X, y, - transforms: Optional[List[str]] = None, ): """Fit the model according to the given training data. @@ -50,10 +49,6 @@ def fit( y (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples,). Target vector relative to X. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: LogisticRegression: Fitted Estimator. diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index d0130434678..89981e34c03 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -8,10 +8,10 @@ # License: BSD 3 clause from bigframes import constants -from third_party.bigframes_vendored.sklearn.base import BaseEstimator +from third_party.bigframes_vendored.sklearn.base import BaseEstimator, TransformerMixin -class StandardScaler(BaseEstimator): +class StandardScaler(BaseEstimator, TransformerMixin): """Standardize features by removing the mean and scaling to unit variance. The standard score of a sample `x` is calculated as:z = (x - u) / s @@ -28,30 +28,23 @@ class StandardScaler(BaseEstimator): machine learning estimators: they might behave badly if the individual features do not more or less look like standard normally distributed data (e.g. Gaussian with 0 mean and unit variance). - """ - - def fit(self, X): - """Compute the mean and std to be used for later scaling. - Examples: + Examples: .. code-block:: from bigframes.ml.preprocessing import StandardScaler + import bigframes.pandas as bpd - enc = StandardScaler() - X = [['Male', 1], ['Female', 3], ['Female', 2]] - enc.fit(X) - - Examples: - - .. code-block:: - - from bigframes.ml import StandardScaler + scaler = StandardScaler() + data = bpd.DataFrame({"a": [0, 0, 1, 1], "b":[0, 0, 1, 1]}) + scaler.fit(data) + print(scaler.transform(data)) + print(scaler.transform(bpd.DataFrame({"a": [2], "b":[2]}))) + """ - enc = StandardScaler() - X = [['Male', 1], ['Female', 3], ['Female', 2]] - enc.fit(X) + def fit(self, X): + """Compute the mean and std to be used for later scaling. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py index b1cf17e539d..b0f0df8e157 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -37,12 +37,8 @@ class OneHotEncoder(BaseEstimator): when considering infrequent categories. If there are infrequent categories, max_categories includes the category representing the infrequent categories along with the frequent categories. Default None, set limit to 1,000,000. - """ - - def fit(self, X): - """Fit OneHotEncoder to X. - Examples: + Examples: Given a dataset with two features, we let the encoder find the unique values per feature and transform the data to a binary one-hot encoding. @@ -50,10 +46,16 @@ def fit(self, X): .. code-block:: from bigframes.ml.preprocessing import OneHotEncoder + import bigframes.pandas as bpd enc = OneHotEncoder() - X = [['Male', 1], ['Female', 3], ['Female', 2]] + X = bpd.DataFrame({"a": ["Male", "Female", "Female"], "b": ["1", "3", "2"]}) enc.fit(X) + print(enc.transform(bpd.DataFrame({"a": ["Female", "Male"], "b": ["1", "4"]}))) + """ + + def fit(self, X): + """Fit OneHotEncoder to X. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): diff --git a/third_party/bigframes_vendored/xgboost/sklearn.py b/third_party/bigframes_vendored/xgboost/sklearn.py index 620c87fa3db..b7b43b85a3e 100644 --- a/third_party/bigframes_vendored/xgboost/sklearn.py +++ b/third_party/bigframes_vendored/xgboost/sklearn.py @@ -37,10 +37,6 @@ def fit(self, X, y): DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. - transforms (Optional[List[str]], default None): - Do not use. Internal param to be deprecated. - Use bigframes.ml.pipeline instead. - Returns: XGBModel: Fitted Estimator. """ From 94779edca7240a8893498ca537266bc628c9bba3 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Mon, 18 Sep 2023 09:13:43 -0500 Subject: [PATCH 20/80] chore(main): release 0.4.0 (#28) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 43 +++++++++++++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa736877f4c..de2edcf31e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,49 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.4.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.2...v0.4.0) (2023-09-16) + + +### Features + +* Add `axis` parameter to `droplevel` and `reorder_levels` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `bfill` and `ffill` to `DataFrame` and `Series` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `DataFrame.combine` and `DataFrame.combine_first` ([#27](https://github.com/googleapis/python-bigquery-dataframes/issues/27)) ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `DataFrame.nlargest`, `nsmallest` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `DataFrame.pct_change` and `Series.pct_change` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `DataFrame.skew` and `GroupBy.skew` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `DataFrame.to_dict`, `to_excel`, `to_latex`, `to_records`, `to_string`, `to_markdown`, `to_pickle`, `to_orc` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `diff` method to `DataFrame` and `GroupBy` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `filter` and `reindex` to `Series` and `DataFrame` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `reindex_like` to `DataFrame` and `Series` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add `swaplevel` to `DataFrame` and `Series` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add partial support for `Sereies.replace` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Support `DataFrame.loc[bool_series, column] = scalar` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Support a persistent `name` in `remote_function` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) + + +### Bug Fixes + +* `remote_function` uses same credentials as other APIs ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Add type hints to models ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Raise error when ARIMAPlus is used with Pipeline ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Remove `transforms` parameter in `model.fit` (**breaking change**) ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Support column joins with "None indexer" ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Use for literals `Int64Dtype` in `cut` ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Use lowercase strings for parameter literals in `bigframes.ml` (**breaking change**) ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) + + +### Performance Improvements + +* `bigframes-api` label to I/O query jobs ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) + + +### Documentation + +* Document possible parameter values for PaLM2TextGenerator ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Document region logic in README ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) +* Fix OneHotEncoder sample ([7c6b0dd](https://github.com/googleapis/python-bigquery-dataframes/commit/7c6b0dd2f99139c8830e762201a45b28486532ff)) + ## [0.3.2](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.1...v0.3.2) (2023-09-06) diff --git a/bigframes/version.py b/bigframes/version.py index ef4f01f6294..65b984a0d73 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.3.2" +__version__ = "0.4.0" From bbbd21ea0d8c5fa13ba66877ce28d20247884afe Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 18 Sep 2023 16:38:16 +0000 Subject: [PATCH 21/80] chore: use correct SCM name for release trigger (#29) Closes #19 --- .github/release-trigger.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/release-trigger.yml b/.github/release-trigger.yml index b0a6cadefca..4fbd4aa427b 100644 --- a/.github/release-trigger.yml +++ b/.github/release-trigger.yml @@ -1,2 +1,2 @@ enabled: true -multiScmName: bigframes +multiScmName: python-bigquery-dataframes From 69e51a67fb4aba9001fd5ea2b25cfcd6a6fdb80f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 18 Sep 2023 23:12:13 +0000 Subject: [PATCH 22/80] chore: enforce use of conventional commits (#31) This will prevent accidental merging of commits that release-please can't handle. --- .github/sync-repo-settings.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index b7dae76ba3b..cfa62f787c9 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -7,6 +7,7 @@ branchProtectionRules: requiresCodeOwnerReviews: true requiresStrictStatusChecks: true requiredStatusCheckContexts: + - 'conventionalcommits.org' - 'cla/google' - 'OwlBot Post Processor' - 'docs' From c1900c29a44199d5d8d036d6d842b4f00448fa79 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 20 Sep 2023 16:26:56 +0000 Subject: [PATCH 23/80] chore: sync internal changes to GitHub (#34) feat: support `optimize_strategy` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `l2_reg` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `max_iterations` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `learn_rate_strategy` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `early_stop` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `min_rel_progress` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `ls_init_learn_rate` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `calculate_p_values` parameter in `bigframes.ml.linear_model.LinearRegression` feat: support `enable_global_explain` parameter in `bigframes.ml.linear_model.LinearRegression` test: add golden SQL test for logistic model test: extend ml golden sql test linear_reg docs: link to Remote Functions code samples from README and API reference feat: support `df[column_name] = df_only_one_column` feat: add `DataFrame.rolling` and `DataFrame.expanding` methods feat: add `DataFrame.kurtosis` / `DF.kurt` method feat: support `class_weights="balanced"` in `LogisticRegression` model --- README.rst | 4 +- bigframes/core/__init__.py | 21 +-- bigframes/core/block_transforms.py | 101 +++++++++++++- bigframes/core/blocks.py | 12 +- bigframes/core/groupby/__init__.py | 70 +++++++++- bigframes/core/window/__init__.py | 50 ++++--- bigframes/dataframe.py | 37 +++++ bigframes/ml/core.py | 4 +- bigframes/ml/ensemble.py | 4 +- bigframes/ml/linear_model.py | 79 +++++++++-- bigframes/operations/__init__.py | 11 ++ bigframes/series.py | 4 +- bigframes/session.py | 3 +- tests/system/large/ml/test_linear_model.py | 93 ++++++++----- tests/system/small/test_dataframe.py | 48 +++++++ tests/system/small/test_groupby.py | 24 ++++ tests/system/small/test_window.py | 42 +++++- tests/unit/ml/test_golden_sql.py | 130 ++++++++++++++++-- .../bigframes_vendored/pandas/core/frame.py | 15 ++ .../bigframes_vendored/pandas/core/generic.py | 55 ++++++++ .../pandas/core/groupby/__init__.py | 21 +++ .../bigframes_vendored/pandas/core/series.py | 55 -------- .../sklearn/linear_model/_base.py | 22 ++- .../sklearn/linear_model/_logistic.py | 11 +- 24 files changed, 745 insertions(+), 171 deletions(-) diff --git a/README.rst b/README.rst index 23aea446ff7..77c42e43251 100644 --- a/README.rst +++ b/README.rst @@ -241,7 +241,9 @@ Remote functions BigQuery DataFrames gives you the ability to turn your custom scalar functions into `BigQuery remote functions `_ . Creating a remote -function in BigQuery DataFrames creates a BigQuery remote function, a `BigQuery +function in BigQuery DataFrames (See `code samples +`_) +creates a BigQuery remote function, a `BigQuery connection `_ , and a `Cloud Functions (2nd gen) function diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 27fe4a4fe6c..8e7beb73dbb 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -600,7 +600,7 @@ def project_window_op( window_spec: WindowSpec, output_name=None, *, - skip_null_groups=False, + never_skip_nulls=False, skip_reproject_unsafe: bool = False, ) -> ArrayValue: """ @@ -609,7 +609,7 @@ def project_window_op( op: the windowable operator to apply to the input column window_spec: a specification of the window over which to apply the operator output_name: the id to assign to the output of the operator, by default will replace input col if distinct output id not provided - skip_null_groups: will filter out any rows where any of the grouping keys is null + never_skip_nulls: will disable null skipping for operators that would otherwise do so skip_reproject_unsafe: skips the reprojection step, can be used when performing many non-dependent window operations, user responsible for not nesting window expressions, or using outputs as join, filter or aggregation keys before a reprojection """ column = typing.cast(ibis_types.Column, self.get_column(column_name)) @@ -618,20 +618,23 @@ def project_window_op( window_op = op._as_ibis(column, window) clauses = [] - if op.skips_nulls: + if op.skips_nulls and not never_skip_nulls: clauses.append((column.isnull(), ibis.NA)) - if skip_null_groups: - for key in window_spec.grouping_keys: - clauses.append((self.get_column(key).isnull(), ibis.NA)) if window_spec.min_periods: + if op.skips_nulls: + # Most operations do not count NULL values towards min_periods + observation_count = agg_ops.count_op._as_ibis(column, window) + else: + # Operations like count treat even NULLs as valid observations for the sake of min_periods + # notnull is just used to convert null values to non-null (FALSE) values to be counted + denulled_value = typing.cast(ibis_types.BooleanColumn, column.notnull()) + observation_count = agg_ops.count_op._as_ibis(denulled_value, window) clauses.append( ( - agg_ops.count_op._as_ibis(column, window) - < ibis_types.literal(window_spec.min_periods), + observation_count < ibis_types.literal(window_spec.min_periods), ibis.NA, ) ) - if clauses: case_statement = ibis.case() for clause in clauses: diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 5dcd9fe7538..da6ba65b8a9 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -355,6 +355,46 @@ def skew( return block +def kurt( + block: blocks.Block, + skew_column_ids: typing.Sequence[str], + grouping_column_ids: typing.Sequence[str] = (), +) -> blocks.Block: + original_columns = skew_column_ids + column_labels = block.select_columns(original_columns).column_labels + + block, delta4_ids = _mean_delta_to_power( + block, 4, original_columns, grouping_column_ids + ) + # counts, moment4 for each column + aggregations = [] + for i, col in enumerate(original_columns): + count_agg = (col, agg_ops.count_op) + moment4_agg = (delta4_ids[i], agg_ops.mean_op) + variance_agg = (col, agg_ops.PopVarOp()) + aggregations.extend([count_agg, moment4_agg, variance_agg]) + + block, agg_ids = block.aggregate( + by_column_ids=grouping_column_ids, aggregations=aggregations + ) + + kurt_ids = [] + for i, col in enumerate(original_columns): + # Corresponds to order of aggregations in preceding loop + count_id, moment4_id, var_id = agg_ids[i * 3 : (i * 3) + 3] + block, kurt_id = _kurt_from_moments_and_count( + block, count_id, moment4_id, var_id + ) + kurt_ids.append(kurt_id) + + block = block.select_columns(kurt_ids).with_column_labels(column_labels) + if not grouping_column_ids: + # When ungrouped, stack everything into single column so can be returned as series + block = block.stack() + block = block.drop_levels([block.index_columns[0]]) + return block + + def _mean_delta_to_power( block: blocks.Block, n_power, @@ -375,13 +415,13 @@ def _mean_delta_to_power( def _skew_from_moments_and_count( - block: blocks.Block, count_id: str, moment3_id: str, var_id: str + block: blocks.Block, count_id: str, moment3_id: str, moment2_id: str ) -> typing.Tuple[blocks.Block, str]: # Calculate skew using count, third moment and population variance # See G1 estimator: # https://en.wikipedia.org/wiki/Skewness#Sample_skewness block, denominator_id = block.apply_unary_op( - var_id, ops.partial_right(ops.pow_op, 3 / 2) + moment2_id, ops.partial_right(ops.unsafe_pow_op, 3 / 2) ) block, base_id = block.apply_binary_op(moment3_id, denominator_id, ops.div_op) block, countminus1_id = block.apply_unary_op( @@ -392,7 +432,7 @@ def _skew_from_moments_and_count( ) block, adjustment_id = block.apply_binary_op(count_id, countminus1_id, ops.mul_op) block, adjustment_id = block.apply_unary_op( - adjustment_id, ops.partial_right(ops.pow_op, 1 / 2) + adjustment_id, ops.partial_right(ops.unsafe_pow_op, 1 / 2) ) block, adjustment_id = block.apply_binary_op( adjustment_id, countminus2_id, ops.div_op @@ -405,3 +445,58 @@ def _skew_from_moments_and_count( skew_id, na_cond_id, ops.partial_arg3(ops.where_op, None) ) return block, skew_id + + +def _kurt_from_moments_and_count( + block: blocks.Block, count_id: str, moment4_id: str, moment2_id: str +) -> typing.Tuple[blocks.Block, str]: + # Kurtosis is often defined as the second standardize moment: moment(4)/moment(2)**2 + # Pandas however uses Fisher’s estimator, implemented below + # numerator = (count + 1) * (count - 1) * moment4 + # denominator = (count - 2) * (count - 3) * moment2**2 + # adjustment = 3 * (count - 1) ** 2 / ((count - 2) * (count - 3)) + # kurtosis = (numerator / denominator) - adjustment + + # Numerator + block, countminus1_id = block.apply_unary_op( + count_id, ops.partial_right(ops.sub_op, 1) + ) + block, countplus1_id = block.apply_unary_op( + count_id, ops.partial_right(ops.add_op, 1) + ) + block, num_adj = block.apply_binary_op(countplus1_id, countminus1_id, ops.mul_op) + block, numerator_id = block.apply_binary_op(moment4_id, num_adj, ops.mul_op) + + # Denominator + block, countminus2_id = block.apply_unary_op( + count_id, ops.partial_right(ops.sub_op, 2) + ) + block, countminus3_id = block.apply_unary_op( + count_id, ops.partial_right(ops.sub_op, 3) + ) + block, denom_adj = block.apply_binary_op(countminus2_id, countminus3_id, ops.mul_op) + block, popvar_squared = block.apply_unary_op( + moment2_id, ops.partial_right(ops.unsafe_pow_op, 2) + ) + block, denominator_id = block.apply_binary_op(popvar_squared, denom_adj, ops.mul_op) + + # Adjustment + block, countminus1_square = block.apply_unary_op( + countminus1_id, ops.partial_right(ops.unsafe_pow_op, 2) + ) + block, adj_num = block.apply_unary_op( + countminus1_square, ops.partial_right(ops.mul_op, 3) + ) + block, adj_denom = block.apply_binary_op(countminus2_id, countminus3_id, ops.mul_op) + block, adjustment_id = block.apply_binary_op(adj_num, adj_denom, ops.div_op) + + # Combine + block, base_id = block.apply_binary_op(numerator_id, denominator_id, ops.div_op) + block, kurt_id = block.apply_binary_op(base_id, adjustment_id, ops.sub_op) + + # Need to produce NA if have less than 4 data points + block, na_cond_id = block.apply_unary_op(count_id, ops.partial_right(ops.ge_op, 4)) + block, kurt_id = block.apply_binary_op( + kurt_id, na_cond_id, ops.partial_arg3(ops.where_op, None) + ) + return block, kurt_id diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 5b414252ee9..fb9ede9f4c9 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -709,6 +709,7 @@ def multi_apply_window_op( window_spec: core.WindowSpec, *, skip_null_groups: bool = False, + never_skip_nulls: bool = False, ) -> typing.Tuple[Block, typing.Sequence[str]]: block = self result_ids = [] @@ -721,6 +722,7 @@ def multi_apply_window_op( skip_reproject_unsafe=(i + 1) < len(columns), result_label=label, skip_null_groups=skip_null_groups, + never_skip_nulls=never_skip_nulls, ) result_ids.append(result_id) return block, result_ids @@ -751,15 +753,21 @@ def apply_window_op( result_label: Label = None, skip_null_groups: bool = False, skip_reproject_unsafe: bool = False, + never_skip_nulls: bool = False, ) -> typing.Tuple[Block, str]: + block = self + if skip_null_groups: + for key in window_spec.grouping_keys: + block, not_null_id = block.apply_unary_op(key, ops.notnull_op) + block = block.filter(not_null_id).drop_columns([not_null_id]) result_id = guid.generate_guid() - expr = self._expr.project_window_op( + expr = block._expr.project_window_op( column, op, window_spec, result_id, - skip_null_groups=skip_null_groups, skip_reproject_unsafe=skip_reproject_unsafe, + never_skip_nulls=never_skip_nulls, ) block = Block( expr, diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 810e145d33f..9be7f22a714 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -156,6 +156,18 @@ def skew( block = block_ops.skew(self._block, self._selected_cols, self._by_col_ids) return df.DataFrame(block) + def kurt( + self, + *, + numeric_only: bool = False, + ) -> df.DataFrame: + if not numeric_only: + self._raise_on_non_numeric("kurt") + block = block_ops.kurt(self._block, self._selected_cols, self._by_col_ids) + return df.DataFrame(block) + + kurtosis = kurt + def all(self) -> df.DataFrame: return self._aggregate_all(agg_ops.all_op) @@ -195,6 +207,36 @@ def diff(self, periods=1) -> series.Series: ) return self._apply_window_op(agg_ops.DiffOp(periods), window=window) + def rolling(self, window: int, min_periods=None) -> windows.Window: + # To get n size window, need current row and n-1 preceding rows. + window_spec = core.WindowSpec( + grouping_keys=self._by_col_ids, + preceding=window - 1, + following=0, + min_periods=min_periods or window, + ) + block = self._block.order_by( + [order.OrderingColumnReference(col) for col in self._by_col_ids], + stable=True, + ) + return windows.Window( + block, window_spec, self._selected_cols, drop_null_groups=self._dropna + ) + + def expanding(self, min_periods: int = 1) -> windows.Window: + window_spec = core.WindowSpec( + grouping_keys=self._by_col_ids, + following=0, + min_periods=min_periods, + ) + block = self._block.order_by( + [order.OrderingColumnReference(col) for col in self._by_col_ids], + stable=True, + ) + return windows.Window( + block, window_spec, self._selected_cols, drop_null_groups=self._dropna + ) + def agg(self, func=None, **kwargs) -> df.DataFrame: if func: if isinstance(func, str): @@ -351,7 +393,7 @@ def _apply_window_op( ) columns = self._aggregated_columns(numeric_only=numeric_only) block, result_ids = self._block.multi_apply_window_op( - columns, op, window_spec=window_spec, skip_null_groups=self._dropna + columns, op, window_spec=window_spec ) block = block.select_columns(result_ids) return df.DataFrame(block) @@ -422,6 +464,12 @@ def skew(self, *args, **kwargs) -> series.Series: block = block_ops.skew(self._block, [self._value_column], self._by_col_ids) return series.Series(block) + def kurt(self, *args, **kwargs) -> series.Series: + block = block_ops.kurt(self._block, [self._value_column], self._by_col_ids) + return series.Series(block) + + kurtosis = kurt + def prod(self, *args) -> series.Series: return self._aggregate(agg_ops.product_op) @@ -510,7 +558,13 @@ def rolling(self, window: int, min_periods=None) -> windows.Window: [order.OrderingColumnReference(col) for col in self._by_col_ids], stable=True, ) - return windows.Window(block, window_spec, self._value_column) + return windows.Window( + block, + window_spec, + [self._value_column], + drop_null_groups=self._dropna, + is_series=True, + ) def expanding(self, min_periods: int = 1) -> windows.Window: window_spec = core.WindowSpec( @@ -522,10 +576,13 @@ def expanding(self, min_periods: int = 1) -> windows.Window: [order.OrderingColumnReference(col) for col in self._by_col_ids], stable=True, ) - return windows.Window(block, window_spec, self._value_column) - - def _ungroup(self) -> series.Series: - return series.Series(self._block.select_column(self._value_column)) + return windows.Window( + block, + window_spec, + [self._value_column], + drop_null_groups=self._dropna, + is_series=True, + ) def _aggregate(self, aggregate_op: agg_ops.AggregateOp) -> series.Series: result_block, _ = self._block.aggregate( @@ -553,6 +610,5 @@ def _apply_window_op( op, result_label=label, window_spec=window_spec, - skip_null_groups=self._dropna, ) return series.Series(block.select_column(result_id)) diff --git a/bigframes/core/window/__init__.py b/bigframes/core/window/__init__.py index 8994004e0b8..d3d081124e9 100644 --- a/bigframes/core/window/__init__.py +++ b/bigframes/core/window/__init__.py @@ -21,54 +21,57 @@ import bigframes.operations.aggregations as agg_ops import third_party.bigframes_vendored.pandas.core.window.rolling as vendored_pandas_rolling -if typing.TYPE_CHECKING: - from bigframes.series import Series - class Window(vendored_pandas_rolling.Window): __doc__ = vendored_pandas_rolling.Window.__doc__ - # TODO(tbergeron): Windows with groupings should create multi-indexed results - def __init__( self, block: blocks.Block, window_spec: core.WindowSpec, - value_column_id: str, + value_column_ids: typing.Sequence[str], + drop_null_groups: bool = True, + is_series: bool = False, ): self._block = block self._window_spec = window_spec - self._value_column_id = value_column_id + self._value_column_ids = value_column_ids + self._drop_null_groups = drop_null_groups + self._is_series = is_series - def count(self) -> Series: + def count(self): return self._apply_aggregate(agg_ops.count_op) - def sum(self) -> Series: + def sum(self): return self._apply_aggregate(agg_ops.sum_op) - def mean(self) -> Series: + def mean(self): return self._apply_aggregate(agg_ops.mean_op) - def var(self) -> Series: + def var(self): return self._apply_aggregate(agg_ops.var_op) - def std(self) -> Series: + def std(self): return self._apply_aggregate(agg_ops.std_op) - def max(self) -> Series: + def max(self): return self._apply_aggregate(agg_ops.max_op) - def min(self) -> Series: + def min(self): return self._apply_aggregate(agg_ops.min_op) def _apply_aggregate( self, op: agg_ops.AggregateOp, - ) -> Series: + ): block = self._block - label = block.col_id_to_label[self._value_column_id] - block, result_id = block.apply_window_op( - self._value_column_id, op, self._window_spec, result_label=label + labels = [block.col_id_to_label[col] for col in self._value_column_ids] + block, result_ids = block.multi_apply_window_op( + self._value_column_ids, + op, + self._window_spec, + skip_null_groups=self._drop_null_groups, + never_skip_nulls=True, ) if self._window_spec.grouping_keys: @@ -80,6 +83,13 @@ def _apply_aggregate( ) block = block.set_index(col_ids=index_ids) - from bigframes.series import Series + if self._is_series: + from bigframes.series import Series - return Series(block.select_column(result_id)) + return Series(block.select_columns(result_ids).with_column_labels(labels)) + else: + from bigframes.dataframe import DataFrame + + return DataFrame( + block.select_columns(result_ids).with_column_labels(labels) + ) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index d65d4ce344e..de4adb912e6 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -49,6 +49,7 @@ import bigframes.core.joins as joins import bigframes.core.ordering as order import bigframes.core.utils as utils +import bigframes.core.window import bigframes.dtypes import bigframes.formatting_helpers as formatter import bigframes.operations as ops @@ -282,6 +283,10 @@ def empty(self) -> bool: def values(self) -> numpy.ndarray: return self.to_numpy() + @property + def _session(self) -> bigframes.Session: + return self._get_block().expr._session + def __len__(self): rows, _ = self.shape return rows @@ -1056,6 +1061,13 @@ def _assign_single_item( ) -> DataFrame: if isinstance(v, bigframes.series.Series): return self._assign_series_join_on_index(k, v) + elif isinstance(v, bigframes.dataframe.DataFrame): + v_df_col_count = len(v._block.value_columns) + if v_df_col_count != 1: + raise ValueError( + f"Cannot set a DataFrame with {v_df_col_count} columns to the single column {k}" + ) + return self._assign_series_join_on_index(k, v[v.columns[0]]) elif callable(v): copy = self.copy() copy[k] = v(copy) @@ -1627,6 +1639,16 @@ def skew(self, *, numeric_only: bool = False): result_block = block_ops.skew(frame._block, frame._block.value_columns) return bigframes.series.Series(result_block) + def kurt(self, *, numeric_only: bool = False): + if not numeric_only: + frame = self._raise_on_non_numeric("kurt") + else: + frame = self._drop_non_numeric() + result_block = block_ops.kurt(frame._block, frame._block.value_columns) + return bigframes.series.Series(result_block) + + kurtosis = kurt + def pivot( self, *, @@ -1882,6 +1904,21 @@ def _perform_join_by_index(self, other: DataFrame, *, how: str = "left"): ) return DataFrame(combined_index._block) + def rolling(self, window: int, min_periods=None) -> bigframes.core.window.Window: + # To get n size window, need current row and n-1 preceding rows. + window_spec = bigframes.core.WindowSpec( + preceding=window - 1, following=0, min_periods=min_periods or window + ) + return bigframes.core.window.Window( + self._block, window_spec, self._block.value_columns + ) + + def expanding(self, min_periods: int = 1) -> bigframes.core.window.Window: + window_spec = bigframes.core.WindowSpec(following=0, min_periods=min_periods) + return bigframes.core.window.Window( + self._block, window_spec, self._block.value_columns + ) + def groupby( self, by: typing.Union[ diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 57f610c4c4b..667d42f7ee6 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -221,7 +221,7 @@ def create_bqml_model( input_data = X_train.join(y_train, how="outer") options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) - session = X_train._get_block().expr._session + session = X_train._session source_sql = input_data.sql options_sql = ml_sql.options(**options) @@ -255,7 +255,7 @@ def create_bqml_time_series_model( options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]}) options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) - session = X_train._get_block().expr._session + session = X_train._session source_sql = input_data.sql options_sql = ml_sql.options(**options) diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index 142edaa00f6..b0f3e5f0813 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -110,7 +110,7 @@ def _from_bq( dummy_regressor = cls() for bf_param, bf_value in dummy_regressor.__dict__.items(): bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) - if bqml_param is not None: + if bqml_param in last_fitting: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) new_xgb_regressor = cls(**kwargs) @@ -431,7 +431,7 @@ def _from_bq( dummy_model = cls() for bf_param, bf_value in dummy_model.__dict__.items(): bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) - if bqml_param is not None: + if bqml_param in last_fitting: kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) new_random_forest_regressor = cls(**kwargs) diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index 1606a15d730..f27b798eea1 100644 --- a/bigframes/ml/linear_model.py +++ b/bigframes/ml/linear_model.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import cast, Dict, List, Optional, Union +from typing import cast, Dict, List, Literal, Optional, Union from google.cloud import bigquery @@ -28,6 +28,23 @@ import third_party.bigframes_vendored.sklearn.linear_model._base import third_party.bigframes_vendored.sklearn.linear_model._logistic +_BQML_PARAMS_MAPPING = { + "optimize_strategy": "optimizationStrategy", + "fit_intercept": "fitIntercept", + "l1_reg": "l1Regularization", + "l2_reg": "l2Regularization", + "max_iterations": "maxIterations", + "learn_rate_strategy": "learnRateStrategy", + "learn_rate": "learnRate", + "early_stop": "earlyStop", + "min_rel_progress": "minRelativeProgress", + "ls_init_learn_rate": "initialLearnRate", + "warm_start": "warmStart", + "calculate_p_values": "calculatePValues", + "enable_global_explain": "enableGlobalExplain", + "category_encoding_method": "categoryEncodingMethod", +} + class LinearRegression( base.SupervisedTrainablePredictor, @@ -39,9 +56,29 @@ class LinearRegression( def __init__( self, + optimize_strategy: Literal[ + "auto_strategy", "batch_gradient_descent", "normal_equation" + ] = "normal_equation", fit_intercept: bool = True, + l2_reg: float = 0.0, + max_iterations: int = 20, + learn_rate_strategy: Literal["line_search", "constant"] = "line_search", + early_stop: bool = True, + min_rel_progress: float = 0.01, + ls_init_learn_rate: float = 0.1, + calculate_p_values: bool = False, + enable_global_explain: bool = False, ): + self.optimize_strategy = optimize_strategy self.fit_intercept = fit_intercept + self.l2_reg = l2_reg + self.max_iterations = max_iterations + self.learn_rate_strategy = learn_rate_strategy + self.early_stop = early_stop + self.min_rel_progress = min_rel_progress + self.ls_init_learn_rate = ls_init_learn_rate + self.calculate_p_values = calculate_p_values + self.enable_global_explain = enable_global_explain self._bqml_model: Optional[core.BqmlModel] = None @classmethod @@ -55,8 +92,12 @@ def _from_bq( # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun last_fitting = model.training_runs[-1]["trainingOptions"] - if "fitIntercept" in last_fitting: - kwargs["fit_intercept"] = last_fitting["fitIntercept"] + + dummy_linear = cls() + for bf_param, bf_value in dummy_linear.__dict__.items(): + bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) + if bqml_param in last_fitting: + kwargs[bf_param] = type(bf_value)(last_fitting[bqml_param]) new_linear_regression = cls(**kwargs) new_linear_regression._bqml_model = core.BqmlModel(session, model) @@ -65,10 +106,20 @@ def _from_bq( @property def _bqml_options(self) -> Dict[str, str | int | bool | float | List[str]]: """The model options as they will be set for BQML""" + # TODO: Support l1_reg, warm_start, and learn_rate with error catching. return { "model_type": "LINEAR_REG", "data_split_method": "NO_SPLIT", + "optimize_strategy": self.optimize_strategy, "fit_intercept": self.fit_intercept, + "l2_reg": self.l2_reg, + "max_iterations": self.max_iterations, + "learn_rate_strategy": self.learn_rate_strategy, + "early_stop": self.early_stop, + "min_rel_progress": self.min_rel_progress, + "ls_init_learn_rate": self.ls_init_learn_rate, + "calculate_p_values": self.calculate_p_values, + "enable_global_explain": self.enable_global_explain, } def _fit( @@ -147,10 +198,11 @@ class LogisticRegression( def __init__( self, fit_intercept: bool = True, - auto_class_weights: bool = False, + class_weights: Optional[Union[Literal["balanced"], Dict[str, float]]] = None, ): self.fit_intercept = fit_intercept - self.auto_class_weights = auto_class_weights + self.class_weights = class_weights + self._auto_class_weight = class_weights == "balanced" self._bqml_model: Optional[core.BqmlModel] = None @classmethod @@ -165,10 +217,8 @@ def _from_bq( last_fitting = model.training_runs[-1]["trainingOptions"] if "fitIntercept" in last_fitting: kwargs["fit_intercept"] = last_fitting["fitIntercept"] - # TODO(ashleyxu): b/285162045 support auto_class_weights once the API is - # fixed and enable the tests. - if "autoClassWeights" in last_fitting: - kwargs["auto_class_weights"] = last_fitting["autoClassWeights"] + if last_fitting["autoClassWeights"]: + kwargs["class_weights"] = "balanced" # TODO(ashleyxu) support class_weights in the constructor. # if "labelClassWeights" in last_fitting: # kwargs["class_weights"] = last_fitting["labelClassWeights"] @@ -184,8 +234,8 @@ def _bqml_options(self) -> Dict[str, str | int | float | List[str]]: "model_type": "LOGISTIC_REG", "data_split_method": "NO_SPLIT", "fit_intercept": self.fit_intercept, - "auto_class_weights": self.auto_class_weights, - # TODO(ashleyxu): support class_weights (struct array) + "auto_class_weights": self._auto_class_weight, + # TODO(ashleyxu): support class_weights (struct array as dict in our API) # "class_weights": self.class_weights, } @@ -253,11 +303,10 @@ def to_gbq(self, model_name: str, replace: bool = False) -> LogisticRegression: if not self._bqml_model: raise RuntimeError("A model must be fitted before it can be saved") - # TODO(ashleyxu): b/285162045 support auto_class_weights once the API is - # fixed and enable the tests. - if self.auto_class_weights is True: + # TODO(ashleyxu): support class_weights (struct array as dict in our API) + if self.class_weights not in (None, "balanced"): raise NotImplementedError( - f"auto_class_weight is not supported yet. {constants.FEEDBACK_LINK}" + f"class_weights is not supported yet. {constants.FEEDBACK_LINK}" ) new_model = self._bqml_model.copy(model_name, replace) diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index fa43f725f64..c5c55607ae5 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -773,6 +773,17 @@ def pow_op( return _float_pow_op(x, y) +@short_circuit_nulls(ibis_dtypes.float) +def unsafe_pow_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + """For internal use only - where domain and overflow checks are not needed.""" + return typing.cast(ibis_types.NumericValue, x) ** typing.cast( + ibis_types.NumericValue, y + ) + + def _int_pow_op( x: ibis_types.Value, y: ibis_types.Value, diff --git a/bigframes/series.py b/bigframes/series.py index 12e72c58b61..9db64fae9c9 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1016,13 +1016,13 @@ def rolling(self, window: int, min_periods=None) -> bigframes.core.window.Window preceding=window - 1, following=0, min_periods=min_periods or window ) return bigframes.core.window.Window( - self._block, window_spec, self._value_column + self._block, window_spec, self._block.value_columns, is_series=True ) def expanding(self, min_periods: int = 1) -> bigframes.core.window.Window: window_spec = WindowSpec(following=0, min_periods=min_periods) return bigframes.core.window.Window( - self._block, window_spec, self._value_column + self._block, window_spec, self._block.value_columns, is_series=True ) def groupby( diff --git a/bigframes/session.py b/bigframes/session.py index 3ca79a7b531..04ae6ba4541 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -1324,7 +1324,8 @@ def remote_function( reuse: bool = True, name: Optional[str] = None, ): - """Decorator to turn a user defined function into a BigQuery remote function. + """Decorator to turn a user defined function into a BigQuery remote function. Check out + the code samples at: https://cloud.google.com/bigquery/docs/remote-functions#bigquery-dataframes. .. note:: Please make sure following is setup before using this API: diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py index 3b90568450b..a0f4182e6fb 100644 --- a/tests/system/large/ml/test_linear_model.py +++ b/tests/system/large/ml/test_linear_model.py @@ -18,7 +18,7 @@ def test_linear_regression_configure_fit_score(penguins_df_default_index, dataset_id): - model = bigframes.ml.linear_model.LinearRegression(fit_intercept=False) + model = bigframes.ml.linear_model.LinearRegression() df = penguins_df_default_index.dropna() X_train = df[ @@ -55,15 +55,24 @@ def test_linear_regression_configure_fit_score(penguins_df_default_index, datase assert ( f"{dataset_id}.temp_configured_model" in reloaded_model._bqml_model.model_name ) - - # TODO(yunmengxie): enable this once b/277242951 (fit_intercept missing from API) is fixed - # assert reloaded_model.fit_intercept == False + assert reloaded_model.optimize_strategy == "NORMAL_EQUATION" + assert reloaded_model.fit_intercept is True + assert reloaded_model.calculate_p_values is False + assert reloaded_model.early_stop is True + assert reloaded_model.enable_global_explain is False + assert reloaded_model.l2_reg == 0.0 + assert reloaded_model.learn_rate_strategy == "line_search" + assert reloaded_model.ls_init_learn_rate == 0.1 + assert reloaded_model.max_iterations == 20 + assert reloaded_model.min_rel_progress == 0.01 -def test_linear_regression_manual_split_configure_fit_score( +def test_linear_regression_customized_params_fit_score( penguins_df_default_index, dataset_id ): - model = bigframes.ml.linear_model.LinearRegression(fit_intercept=True) + model = bigframes.ml.linear_model.LinearRegression( + fit_intercept=False, l2_reg=0.1, min_rel_progress=0.01 + ) df = penguins_df_default_index.dropna() X_train = df[ @@ -83,12 +92,12 @@ def test_linear_regression_manual_split_configure_fit_score( result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { - "mean_absolute_error": [225.735767], - "mean_squared_error": [80417.461828], - "mean_squared_log_error": [0.004967], - "median_absolute_error": [172.543702], - "r2_score": [0.87548], - "explained_variance": [0.87548], + "mean_absolute_error": [226.108411], + "mean_squared_error": [80459.668456], + "mean_squared_log_error": [0.00497], + "median_absolute_error": [171.618872], + "r2_score": [0.875415], + "explained_variance": [0.875417], }, dtype="Float64", ) @@ -100,13 +109,21 @@ def test_linear_regression_manual_split_configure_fit_score( assert ( f"{dataset_id}.temp_configured_model" in reloaded_model._bqml_model.model_name ) - assert reloaded_model.fit_intercept is True - - -def test_logistic_regression_auto_class_weights_configure_fit_score( - penguins_df_default_index, dataset_id -): + assert reloaded_model.optimize_strategy == "NORMAL_EQUATION" + assert reloaded_model.fit_intercept is False + assert reloaded_model.calculate_p_values is False + assert reloaded_model.early_stop is True + assert reloaded_model.enable_global_explain is False + assert reloaded_model.l2_reg == 0.1 + assert reloaded_model.learn_rate_strategy == "line_search" + assert reloaded_model.ls_init_learn_rate == 0.1 + assert reloaded_model.max_iterations == 20 + assert reloaded_model.min_rel_progress == 0.01 + + +def test_logistic_regression_configure_fit_score(penguins_df_default_index, dataset_id): model = bigframes.ml.linear_model.LogisticRegression() + df = penguins_df_default_index.dropna() X_train = df[ [ @@ -115,6 +132,7 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", + "body_mass_g", ] ] y_train = df[["sex"]] @@ -124,12 +142,12 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { - "precision": [0.58085], - "recall": [0.582576], - "accuracy": [0.871257], - "f1_score": [0.58171], - "log_loss": [1.59285], - "roc_auc": [0.9602], + "precision": [0.616753], + "recall": [0.618615], + "accuracy": [0.92515], + "f1_score": [0.617681], + "log_loss": [1.498832], + "roc_auc": [0.975807], }, dtype="Float64", ) @@ -145,15 +163,15 @@ def test_logistic_regression_auto_class_weights_configure_fit_score( in reloaded_model._bqml_model.model_name ) assert reloaded_model.fit_intercept is True - # TODO(gaotianxiang): enable this once (auto_class_weights missing from API) is fixed - # assert reloaded_model.auto_class_weights is True + assert reloaded_model.class_weights is None -def test_logistic_regression_manual_split_configure_fit_score( +def test_logistic_regression_customized_params_fit_score( penguins_df_default_index, dataset_id ): - model = bigframes.ml.linear_model.LogisticRegression(fit_intercept=True) - + model = bigframes.ml.linear_model.LogisticRegression( + fit_intercept=False, class_weights="balanced" + ) df = penguins_df_default_index.dropna() X_train = df[ [ @@ -162,7 +180,6 @@ def test_logistic_regression_manual_split_configure_fit_score( "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", - "body_mass_g", ] ] y_train = df[["sex"]] @@ -172,12 +189,12 @@ def test_logistic_regression_manual_split_configure_fit_score( result = model.score(X_train, y_train).to_pandas() expected = pd.DataFrame( { - "precision": [0.616753], - "recall": [0.618615], - "accuracy": [0.92515], - "f1_score": [0.617681], - "log_loss": [1.498832], - "roc_auc": [0.975807], + "precision": [0.58483], + "recall": [0.586616], + "accuracy": [0.877246], + "f1_score": [0.58571], + "log_loss": [1.032699], + "roc_auc": [0.924132], }, dtype="Float64", ) @@ -192,5 +209,5 @@ def test_logistic_regression_manual_split_configure_fit_score( f"{dataset_id}.temp_configured_logistic_reg_model" in reloaded_model._bqml_model.model_name ) - assert reloaded_model.fit_intercept is True - assert reloaded_model.auto_class_weights is False + assert reloaded_model.fit_intercept is False + assert reloaded_model.class_weights == "balanced" diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a85777c59d6..ed682c855b7 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -411,6 +411,30 @@ def test_assign_new_column_w_setitem(scalars_dfs): pd.testing.assert_frame_equal(bf_result, pd_result) +def test_assign_new_column_w_setitem_dataframe(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + bf_df["int64_col"] = bf_df["int64_too"].to_frame() + pd_df["int64_col"] = pd_df["int64_too"].to_frame() + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_df["int64_col"] = pd_df["int64_col"].astype("Int64") + + pd.testing.assert_frame_equal(bf_df.to_pandas(), pd_df) + + +def test_assign_new_column_w_setitem_dataframe_error(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_df = scalars_df.copy() + pd_df = scalars_pandas_df.copy() + + with pytest.raises(ValueError): + bf_df["impossible_col"] = bf_df[["int64_too", "string_col"]] + with pytest.raises(ValueError): + pd_df["impossible_col"] = pd_df[["int64_too", "string_col"]] + + def test_assign_new_column_w_setitem_list(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_df = scalars_df.copy() @@ -1996,6 +2020,30 @@ def test_df_skew(scalars_dfs): pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) +def test_df_kurt_too_few_values(scalars_dfs): + columns = ["float64_col", "int64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[columns].head(2).kurt().to_pandas() + pd_result = scalars_pandas_df[columns].head(2).kurt() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + +def test_df_kurt(scalars_dfs): + columns = ["float64_col", "int64_col"] + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[columns].kurt().to_pandas() + pd_result = scalars_pandas_df[columns].kurt() + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + @pytest.mark.parametrize( ("frac", "n", "random_state"), [ diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 18741468c5d..05154f7ab7d 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -238,6 +238,16 @@ def test_series_groupby_skew(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) +def test_series_groupby_kurt(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.groupby("bool_col")["int64_too"].kurt().to_pandas() + # Pandas doesn't have groupby.kurt yet: https://github.com/pandas-dev/pandas/issues/40139 + pd_result = scalars_pandas_df_index.groupby("bool_col")["int64_too"].apply( + pd.Series.kurt + ) + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + def test_dataframe_groupby_skew(scalars_df_index, scalars_pandas_df_index): col_names = ["float64_col", "int64_col", "bool_col"] bf_result = scalars_df_index[col_names].groupby("bool_col").skew().to_pandas() @@ -246,6 +256,20 @@ def test_dataframe_groupby_skew(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) +def test_dataframe_groupby_kurt(scalars_df_index, scalars_pandas_df_index): + col_names = ["float64_col", "int64_col", "bool_col"] + bf_result = scalars_df_index[col_names].groupby("bool_col").kurt().to_pandas() + # Pandas doesn't have groupby.kurt yet: https://github.com/pandas-dev/pandas/issues/40139 + pd_result = ( + scalars_pandas_df_index[col_names] + .groupby("bool_col") + .apply(pd.Series.kurt) + .drop("bool_col", axis=1) + ) + + pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) + + def test_dataframe_groupby_diff(scalars_df_index, scalars_pandas_df_index): col_names = ["float64_col", "int64_col", "string_col"] bf_result = scalars_df_index[col_names].groupby("string_col").diff(-1) diff --git a/tests/system/small/test_window.py b/tests/system/small/test_window.py index e2f0fe999b1..2b9ec1a3c0f 100644 --- a/tests/system/small/test_window.py +++ b/tests/system/small/test_window.py @@ -41,7 +41,9 @@ pytest.param(lambda x: x.var(), id="var"), ], ) -def test_window_agg_ops(scalars_df_index, scalars_pandas_df_index, windowing, agg_op): +def test_series_window_agg_ops( + scalars_df_index, scalars_pandas_df_index, windowing, agg_op +): col_name = "int64_too" bf_series = agg_op(windowing(scalars_df_index[col_name])).to_pandas() pd_series = agg_op(windowing(scalars_pandas_df_index[col_name])) @@ -53,3 +55,41 @@ def test_window_agg_ops(scalars_df_index, scalars_pandas_df_index, windowing, ag pd_series, bf_series, ) + + +@pytest.mark.parametrize( + ("windowing"), + [ + pytest.param(lambda x: x.expanding(), id="expanding"), + pytest.param(lambda x: x.rolling(3, min_periods=3), id="rolling"), + pytest.param( + lambda x: x.groupby(level=0).rolling(3, min_periods=3), id="rollinggroupby" + ), + pytest.param( + lambda x: x.groupby("int64_too").expanding(min_periods=2), + id="expandinggroupby", + ), + ], +) +@pytest.mark.parametrize( + ("agg_op"), + [ + pytest.param(lambda x: x.sum(), id="sum"), + pytest.param(lambda x: x.min(), id="min"), + pytest.param(lambda x: x.max(), id="max"), + pytest.param(lambda x: x.mean(), id="mean"), + pytest.param(lambda x: x.count(), id="count"), + pytest.param(lambda x: x.std(), id="std"), + pytest.param(lambda x: x.var(), id="var"), + ], +) +def test_dataframe_window_agg_ops( + scalars_df_index, scalars_pandas_df_index, windowing, agg_op +): + scalars_df_index = scalars_df_index.set_index("bool_col") + scalars_pandas_df_index = scalars_pandas_df_index.set_index("bool_col") + col_names = ["int64_too", "float64_col"] + bf_result = agg_op(windowing(scalars_df_index[col_names])).to_pandas() + pd_result = agg_op(windowing(scalars_pandas_df_index[col_names])) + + pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) diff --git a/tests/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py index 8d4932a3c3c..584d080d428 100644 --- a/tests/unit/ml/test_golden_sql.py +++ b/tests/unit/ml/test_golden_sql.py @@ -14,34 +14,146 @@ from unittest import mock +from google.cloud import bigquery +import pandas as pd +import pytest import pytest_mock import bigframes -from bigframes.ml import linear_model +from bigframes.ml import core, linear_model import bigframes.pandas as bpd -def test_linear_regression_default_fit(mocker: pytest_mock.MockerFixture): +@pytest.fixture +def mock_session(): mock_session = mock.create_autospec(spec=bigframes.Session) - mock_X = mock.create_autospec(spec=bpd.DataFrame) - mock_X._get_block().expr._session = mock_session + # return values we don't care about, but need to provide to continue the program when calling session._start_query() + mock_session._start_query.return_value = (None, mock.MagicMock()) + + return mock_session + +@pytest.fixture +def mock_y(): mock_y = mock.create_autospec(spec=bpd.DataFrame) - mock_y.columns.tolist.return_value = ["input_label_column"] + mock_y.columns = pd.Index(["input_column_label"]) - mock_X.join(mock_y).sql = "input_dataframe_sql" + return mock_y - # return values we don't care about, but need to provide to continue the program - mock_session._start_query.return_value = (None, mock.MagicMock()) +@pytest.fixture +def mock_X(mock_y, mock_session): + mock_X = mock.create_autospec(spec=bpd.DataFrame) + mock_X._session = mock_session + mock_X._to_sql_query.return_value = ( + "input_X_sql", + ["index_column_id"], + ["index_column_label"], + ) + mock_X.join(mock_y).sql = "input_X_y_sql" + mock_X.join(mock_y)._to_sql_query.return_value = ( + "input_X_y_sql", + ["index_column_id"], + ["index_column_label"], + ) + + return mock_X + + +@pytest.fixture +def bqml_model(mock_session): + bqml_model = core.BqmlModel( + mock_session, bigquery.Model("model_project.model_dataset.model_name") + ) + + return bqml_model + + +@pytest.fixture +def ml_mocker(mocker: pytest_mock.MockerFixture): mocker.patch( "bigframes.ml.core._create_temp_model_name", return_value="temp_model_name" ) + return mocker + + +def test_linear_regression_default_fit(ml_mocker, mock_session, mock_X, mock_y): + model = linear_model.LinearRegression() + model.fit(mock_X, mock_y) + + mock_session._start_query.assert_called_once_with( + 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=True,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + ) + + +def test_linear_regression_params_fit(ml_mocker, mock_session, mock_X, mock_y): + model = linear_model.LinearRegression(fit_intercept=False) + model.fit(mock_X, mock_y) + + mock_session._start_query.assert_called_once_with( + 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=False,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + ) + + +def test_linear_regression_predict(mock_session, bqml_model, mock_X): + model = linear_model.LinearRegression() + model._bqml_model = bqml_model + model.predict(mock_X) + + mock_session.read_gbq.assert_called_once_with( + "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_name`,\n (input_X_sql))", + index_col=["index_column_id"], + ) + + +def test_linear_regression_score(mock_session, bqml_model, mock_X, mock_y): model = linear_model.LinearRegression() + model._bqml_model = bqml_model + model.score(mock_X, mock_y) + + mock_session.read_gbq.assert_called_once_with( + "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_name`,\n (input_X_y_sql))" + ) + + +def test_logistic_regression_default_fit(ml_mocker, mock_session, mock_X, mock_y): + model = linear_model.LogisticRegression() model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=True,\n INPUT_LABEL_COLS=["input_label_column"])\nAS input_dataframe_sql' + 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=True,\n auto_class_weights=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + ) + + +def test_logistic_regression_params_fit(ml_mocker, mock_session, mock_X, mock_y): + model = linear_model.LogisticRegression( + fit_intercept=False, class_weights="balanced" + ) + model.fit(mock_X, mock_y) + + mock_session._start_query.assert_called_once_with( + 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=False,\n auto_class_weights=True,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + ) + + +def test_logistic_regression_predict(mock_session, bqml_model, mock_X): + model = linear_model.LogisticRegression() + model._bqml_model = bqml_model + model.predict(mock_X) + + mock_session.read_gbq.assert_called_once_with( + "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_name`,\n (input_X_sql))", + index_col=["index_column_id"], + ) + + +def test_logistic_regression_score(mock_session, bqml_model, mock_X, mock_y): + model = linear_model.LogisticRegression() + model._bqml_model = bqml_model + model.score(mock_X, mock_y) + + mock_session.read_gbq.assert_called_once_with( + "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_name`,\n (input_X_y_sql))" ) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 113c6547a0e..27cc2144e0f 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -1569,6 +1569,21 @@ def skew(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def kurt(self, *, numeric_only: bool = False): + """Return unbiased kurtosis over requested axis. + + Kurtosis obtained using Fisher's definition of + kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + + Args: + numeric_only (bool, default False): + Include only float, int, boolean columns. + + Returns: + Series + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def std(self, *, numeric_only: bool = False): """Return sample standard deviation over requested axis. diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 7d496891b0b..27d2e845376 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -511,6 +511,61 @@ def rank( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def rolling( + self, + window, + min_periods: int | None = None, + ): + """ + Provide rolling window calculations. + + Args: + window (int, timedelta, str, offset, or BaseIndexer subclass): + Size of the moving window. + + If an integer, the fixed number of observations used for + each window. + + If a timedelta, str, or offset, the time period of each window. Each + window will be a variable sized based on the observations included in + the time-period. This is only valid for datetime-like indexes. + To learn more about the offsets & frequency strings, please see `this link + `__. + + If a BaseIndexer subclass, the window boundaries + based on the defined ``get_window_bounds`` method. Additional rolling + keyword arguments, namely ``min_periods``, ``center``, ``closed`` and + ``step`` will be passed to ``get_window_bounds``. + + min_periods (int, default None): + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. + + For a window that is specified by an offset, ``min_periods`` will default to 1. + + For a window that is specified by an integer, ``min_periods`` will default + to the size of the window. + + Returns: + bigframes.core.window.Window: ``Window`` subclass if a ``win_type`` is passed. + ``Rolling`` subclass if ``win_type`` is not passed. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def expanding(self, min_periods=1): + """ + Provide expanding window calculations. + + Args: + min_periods (int, default 1): + Minimum number of observations in window required to have a value; + otherwise, result is ``np.nan``. + + Returns: + bigframes.core.window.Window: ``Expanding`` subclass. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def __nonzero__(self): raise ValueError( f"The truth value of a {type(self).__name__} is ambiguous. " diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index 7849a3afd52..b05319b4f7d 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -144,6 +144,27 @@ def skew( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def kurt( + self, + *, + numeric_only: bool = False, + ): + """ + Return unbiased kurtosis over requested axis. + + Kurtosis obtained using Fisher's definition of + kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + + Args: + numeric_only (bool, default False): + Include only `float`, `int` or `boolean` data. + + Returns: + Series or DataFrame + Variance of values within each group. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def sum( self, numeric_only: bool = False, diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index c6dd973372c..d58c1ccc3b4 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1673,61 +1673,6 @@ def rename_axis(self, mapper, **kwargs): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def rolling( - self, - window, - min_periods: int | None = None, - ): - """ - Provide rolling window calculations. - - Args: - window (int, timedelta, str, offset, or BaseIndexer subclass): - Size of the moving window. - - If an integer, the fixed number of observations used for - each window. - - If a timedelta, str, or offset, the time period of each window. Each - window will be a variable sized based on the observations included in - the time-period. This is only valid for datetime-like indexes. - To learn more about the offsets & frequency strings, please see `this link - `__. - - If a BaseIndexer subclass, the window boundaries - based on the defined ``get_window_bounds`` method. Additional rolling - keyword arguments, namely ``min_periods``, ``center``, ``closed`` and - ``step`` will be passed to ``get_window_bounds``. - - min_periods (int, default None): - Minimum number of observations in window required to have a value; - otherwise, result is ``np.nan``. - - For a window that is specified by an offset, ``min_periods`` will default to 1. - - For a window that is specified by an integer, ``min_periods`` will default - to the size of the window. - - Returns: - bigframes.core.window.Window: ``Window`` subclass if a ``win_type`` is passed. - ``Rolling`` subclass if ``win_type`` is not passed. - """ - raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - - def expanding(self, min_periods=1): - """ - Provide expanding window calculations. - - Args: - min_periods (int, default 1): - Minimum number of observations in window required to have a value; - otherwise, result is ``np.nan``. - - Returns: - bigframes.core.window.Window: ``Expanding`` subclass. - """ - raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def value_counts( self, normalize: bool = False, diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py index 81b4fca1572..8dc3b6280a4 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_base.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -64,10 +64,30 @@ class LinearRegression(RegressorMixin, LinearModel): the dataset, and the targets predicted by the linear approximation. Args: - fit_intercept (default True): + optimize_strategy (str, default "normal_equation"): + The strategy to train linear regression models. Possible values are + "auto_strategy", "batch_gradient_descent", "normal_equation". Default + to "normal_equation". + fit_intercept (bool, default True): Default ``True``. Whether to calculate the intercept for this model. If set to False, no intercept will be used in calculations (i.e. data is expected to be centered). + l2_reg (float, default 0.0): + The amount of L2 regularization applied. Default to 0. + max_iterations (int, default 20): + The maximum number of training iterations or steps. Default to 20. + learn_rate_strategy (str, default "line_search"): + The strategy for specifying the learning rate during training. Default to "line_search". + early_stop (bool, default True): + Whether training should stop after the first iteration in which the relative loss improvement is less than the value specified for min_rel_progress. Default to True. + min_rel_progress (float, default 0.01): + The minimum relative loss improvement that is necessary to continue training when EARLY_STOP is set to true. For example, a value of 0.01 specifies that each iteration must reduce the loss by 1% for training to continue. Default to 0.01. + ls_init_learn_rate (float, default 0.1): + Sets the initial learning rate that learn_rate_strategy='line_search' uses. This option can only be used if line_search is specified. Default to 0.1. + calculate_p_values (bool, default False): + Specifies whether to compute p-values and standard errors during training. Default to False. + enable_global_explain (bool, default False): + Whether to compute global explanations using explainable AI to evaluate global feature importance to the model. Default to False. """ def fit( diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py index 133dc4498e4..989ca03c827 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_logistic.py @@ -28,9 +28,14 @@ class LogisticRegression(LinearClassifierMixin, BaseEstimator): fit_intercept (default True): Default True. Specifies if a constant (a.k.a. bias or intercept) should be added to the decision function. - auto_class_weights (default False): - Default False. If True, balance class labels using weights for each - class in inverse proportion to the frequency of that class. + class_weights (dict or 'balanced', default None): + Default None. Weights associated with classes in the form + ``{class_label: weight}``.If not given, all classes are supposed + to have weight one. The "balanced" mode uses the values of y to + automatically adjust weights inversely proportional to class + frequencies in the input data as + ``n_samples / (n_classes * np.bincount(y))``. Dict isn't + supported now. """ def fit( From 5056da6b385dbcfc179d2bcbb6549fa539428cda Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 20 Sep 2023 15:16:38 -0700 Subject: [PATCH 24/80] perf: simplify join order to use multiple order keys instead of string. (#36) Change-Id: I8c37e9296b2e4e0ea87f6a7e836d48988d161d37 --- bigframes/core/joins/single_column.py | 118 +++++++++++--------------- 1 file changed, 50 insertions(+), 68 deletions(-) diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 4c865fffdf7..7aba71fd95e 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -63,7 +63,6 @@ def join_by_column( allow_row_identity_join (bool): If True, allow matching by row identity. Set to False to always perform a true JOIN in generated SQL. - Returns: The joined expression and the objects needed to interpret it. @@ -123,13 +122,13 @@ def join_by_column( ), ) else: - # Generate offsets if non-default ordering is applied - # Assumption, both sides are totally ordered, otherwise offsets will be nondeterministic left_table = left.to_ibis_expr( - ordering_mode="string_encoded", order_col_name=core.ORDER_ID_COLUMN + ordering_mode="unordered", + expose_hidden_cols=True, ) right_table = right.to_ibis_expr( - ordering_mode="string_encoded", order_col_name=core.ORDER_ID_COLUMN + ordering_mode="unordered", + expose_hidden_cols=True, ) join_conditions = [ value_to_join_key(left_table[left_index]) @@ -178,41 +177,13 @@ def get_column_right(key: str) -> str: return key - left_ordering_encoding_size = ( - left._ordering.string_encoding.length - if left._ordering.is_string_encoded - else bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH - ) - right_ordering_encoding_size = ( - right._ordering.string_encoding.length - if right._ordering.is_string_encoded - else bigframes.core.ordering.DEFAULT_ORDERING_ID_LENGTH - ) - - # Preserve original ordering accross joins. - left_order_id = get_column_left(core.ORDER_ID_COLUMN) - right_order_id = get_column_right(core.ORDER_ID_COLUMN) - new_order_id_col = _merge_order_ids( - typing.cast(ibis_types.StringColumn, combined_table[left_order_id]), - left_ordering_encoding_size, - typing.cast(ibis_types.StringColumn, combined_table[right_order_id]), - right_ordering_encoding_size, - how, - ) - new_order_id = new_order_id_col.get_name() - if new_order_id is None: - raise ValueError("new_order_id unexpectedly has no name") - - hidden_columns = (new_order_id_col,) - ordering = core.ExpressionOrdering( - # Order id is non-nullable but na_last=False generates simpler sql with current impl - ordering_value_columns=[ - core.OrderingColumnReference(new_order_id, na_last=False) - ], - total_ordering_columns=frozenset([new_order_id]), - string_encoding=core.StringEncoding( - True, left_ordering_encoding_size + right_ordering_encoding_size - ), + # Preserve ordering accross joins. + ordering = join_orderings( + left._ordering, + right._ordering, + get_column_left, + get_column_right, + left_order_dominates=(how != "right"), ) left_join_keys = [ @@ -234,11 +205,21 @@ def get_column_right(key: str) -> str: for col in right.columns ] ) + hidden_ordering_columns = [ + *[ + combined_table[get_column_left(col.get_name())] + for col in left.hidden_ordering_columns + ], + *[ + combined_table[get_column_right(col.get_name())] + for col in right.hidden_ordering_columns + ], + ] combined_expr = core.ArrayValue( left._session, combined_table, columns=columns, - hidden_ordering_columns=hidden_columns, + hidden_ordering_columns=hidden_ordering_columns, ordering=ordering, ) if sort: @@ -313,32 +294,33 @@ def value_to_join_key(value: ibis_types.Value): return value.fillna(ibis_types.literal("$NULL_SENTINEL$")) -def _merge_order_ids( - left_id: ibis_types.StringColumn, - left_encoding_size: int, - right_id: ibis_types.StringColumn, - right_encoding_size: int, - how: str, -) -> ibis_types.StringColumn: - if how == "right": - return _merge_order_ids( - right_id, right_encoding_size, left_id, left_encoding_size, "left" - ) +def join_orderings( + left: core.ExpressionOrdering, + right: core.ExpressionOrdering, + left_id_mapping: Callable[[str], str], + right_id_mapping: Callable[[str], str], + left_order_dominates: bool = True, +) -> core.ExpressionOrdering: + left_ordering_refs = [ + ref.with_name(left_id_mapping(ref.column_id)) + for ref in left.all_ordering_columns + ] + right_ordering_refs = [ + ref.with_name(right_id_mapping(ref.column_id)) + for ref in right.all_ordering_columns + ] + if left_order_dominates: + joined_refs = [*left_ordering_refs, *right_ordering_refs] + else: + joined_refs = [*right_ordering_refs, *left_ordering_refs] - if how == "left": - right_id = typing.cast( - ibis_types.StringColumn, - right_id.fillna(ibis_types.literal(":" * right_encoding_size)), - ) - elif how != "inner": # outer join - left_id = typing.cast( - ibis_types.StringColumn, - left_id.fillna(ibis_types.literal(":" * left_encoding_size)), - ) - right_id = typing.cast( - ibis_types.StringColumn, - right_id.fillna(ibis_types.literal(":" * right_encoding_size)), - ) - return (left_id + right_id).name( - bigframes.core.guid.generate_guid(prefix="bigframes_ordering_id_") + left_total_order_cols = frozenset( + [left_id_mapping(id) for id in left.total_ordering_columns] + ) + right_total_order_cols = frozenset( + [right_id_mapping(id) for id in right.total_ordering_columns] + ) + return core.ExpressionOrdering( + ordering_value_columns=joined_refs, + total_ordering_columns=left_total_order_cols | right_total_order_cols, ) From edabdbb131150707ea9211292cacbb60b8d076dd Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 20 Sep 2023 17:22:07 -0700 Subject: [PATCH 25/80] fix: loosen filter items tests to accomodate shifting pandas impl (#41) --- tests/system/small/test_dataframe.py | 10 ++++++---- tests/system/small/test_series.py | 4 +++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index ed682c855b7..6c96387e97e 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2129,10 +2129,10 @@ def test_df_columns_filter_items(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.filter(items=["string_col", "int64_col"]).to_pandas() pd_result = scalars_pandas_df_index.filter(items=["string_col", "int64_col"]) - + # Ignore column ordering as pandas order differently depending on version pd.testing.assert_frame_equal( - bf_result, - pd_result, + bf_result.sort_index(axis=1), + pd_result.sort_index(axis=1), ) @@ -2167,9 +2167,11 @@ def test_df_rows_filter_items(scalars_df_index, scalars_pandas_df_index): # Pandas uses int64 instead of Int64 (nullable) dtype. pd_result.index = pd_result.index.astype(pd.Int64Dtype()) - pd.testing.assert_frame_equal( + # Ignore ordering as pandas order differently depending on version + assert_pandas_df_equal_ignore_ordering( bf_result, pd_result, + check_names=False, ) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index d3560540cc9..d702049e684 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -1950,9 +1950,11 @@ def test_series_filter_items(scalars_df_index, scalars_pandas_df_index): # Pandas uses int64 instead of Int64 (nullable) dtype. pd_result.index = pd_result.index.astype(pd.Int64Dtype()) - pd.testing.assert_series_equal( + # Ignore ordering as pandas order differently depending on version + assert_series_equal_ignoring_order( bf_result, pd_result, + check_names=False, ) From 109ee24108875389a654674569583e2c7a32d853 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 20 Sep 2023 19:52:36 -0700 Subject: [PATCH 26/80] refactor: remove ibis references outside of arrayvalue code. (#37) Change-Id: I1386355446e90f89a43cee8a9f447f0775639902 --- bigframes/core/__init__.py | 59 ++++++++++++++----- bigframes/core/blocks.py | 84 ++++++++++++++++++--------- bigframes/core/indexers.py | 44 +++++--------- bigframes/core/joins/single_column.py | 4 +- bigframes/core/scalar.py | 50 +--------------- bigframes/dataframe.py | 46 +++------------ bigframes/operations/__init__.py | 34 ++++++++--- bigframes/series.py | 16 +---- tests/unit/test_core.py | 12 ++-- 9 files changed, 160 insertions(+), 189 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 8e7beb73dbb..dd91f80e63b 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -269,7 +269,7 @@ def _get_hidden_ordering_column(self, key: str) -> ibis_types.Column: return typing.cast(ibis_types.Column, self._hidden_ordering_column_names[key]) def apply_limit(self, max_results: int) -> ArrayValue: - table = self.to_ibis_expr( + table = self._to_ibis_expr( ordering_mode="order_by", expose_hidden_cols=True, ).limit(max_results) @@ -285,11 +285,23 @@ def apply_limit(self, max_results: int) -> ArrayValue: ordering=self._ordering, ) - def filter(self, predicate: ibis_types.BooleanValue) -> ArrayValue: + def filter(self, predicate_id: str, keep_null: bool = False) -> ArrayValue: + """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" + condition = typing.cast(ibis_types.BooleanValue, self.get_column(predicate_id)) + if keep_null: + condition = typing.cast( + ibis_types.BooleanValue, + condition.fillna( + typing.cast(ibis_types.BooleanScalar, ibis_types.literal(True)) + ), + ) + return self._filter(condition) + + def _filter(self, predicate_value: ibis_types.BooleanValue) -> ArrayValue: """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" expr = self.builder() expr.ordering = expr.ordering.with_non_sequential() - expr.predicates = [*self._predicates, predicate] + expr.predicates = [*self._predicates, predicate_value] return expr.build() def order_by( @@ -310,7 +322,7 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue: .. warning:: The row numbers of result is non-deterministic, avoid to use. """ - table = self.to_ibis_expr( + table = self._to_ibis_expr( ordering_mode="order_by", expose_hidden_cols=True, fraction=fraction ) columns = [table[column_name] for column_name in self._column_names] @@ -342,7 +354,7 @@ def project_offsets(self) -> ArrayValue: if self._ordering.is_sequential: return self # TODO(tbergeron): Enforce total ordering - table = self.to_ibis_expr( + table = self._to_ibis_expr( ordering_mode="offset_col", order_col_name=ORDER_ID_COLUMN ) columns = [table[column_name] for column_name in self._column_names] @@ -412,7 +424,7 @@ def projection(self, columns: Iterable[ibis_types.Value]) -> ArrayValue: def shape(self) -> typing.Tuple[int, int]: """Returns dimensions as (length, width) tuple.""" width = len(self.columns) - count_expr = self.to_ibis_expr(ordering_mode="unordered").count() + count_expr = self._to_ibis_expr(ordering_mode="unordered").count() sql = self._session.ibis_client.compile(count_expr) row_iterator, _ = self._session._start_query( sql=sql, @@ -435,7 +447,7 @@ def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: ) for i, expr in enumerate([self, *other]): ordering_prefix = str(i).zfill(prefix_size) - table = expr.to_ibis_expr( + table = expr._to_ibis_expr( ordering_mode="string_encoded", order_col_name=ORDER_ID_COLUMN ) # Rename the value columns based on horizontal offset before applying union. @@ -522,7 +534,7 @@ def aggregate( by_column_id: column id of the aggregation key, this is preserved through the transform dropna: whether null keys should be dropped """ - table = self.to_ibis_expr(ordering_mode="unordered") + table = self._to_ibis_expr(ordering_mode="unordered") stats = { col_out: agg_op._as_ibis(table[col_in]) for col_in, agg_op, col_out in aggregations @@ -541,7 +553,7 @@ def aggregate( expr = ArrayValue(self._session, result, columns=columns, ordering=ordering) if dropna: for column_id in by_column_ids: - expr = expr.filter( + expr = expr._filter( ops.notnull_op._as_ibis(expr.get_column(column_id)) ) # Can maybe remove this as Ordering id is redundant as by_column is unique after aggregation @@ -572,7 +584,7 @@ def corr_aggregate( Arguments: corr_aggregations: left_column_id, right_column_id, output_column_id tuples """ - table = self.to_ibis_expr(ordering_mode="unordered") + table = self._to_ibis_expr(ordering_mode="unordered") stats = { col_out: table[col_left].corr(table[col_right], how="pop") for col_left, col_right, col_out in corr_aggregations @@ -646,7 +658,24 @@ def project_window_op( # TODO(tbergeron): Automatically track analytic expression usage and defer reprojection until required for valid query generation. return result._reproject_to_table() if not skip_reproject_unsafe else result - def to_ibis_expr( + def to_sql( + self, + ordering_mode: Literal[ + "order_by", "string_encoded", "offset_col", "unordered" + ] = "order_by", + order_col_name: Optional[str] = ORDER_ID_COLUMN, + col_id_overrides: typing.Mapping[str, str] = {}, + ) -> str: + sql = self._session.ibis_client.compile( + self._to_ibis_expr( + ordering_mode=ordering_mode, + order_col_name=order_col_name, + col_id_overrides=col_id_overrides, + ) + ) + return typing.cast(str, sql) + + def _to_ibis_expr( self, ordering_mode: Literal[ "order_by", "string_encoded", "offset_col", "unordered" @@ -814,7 +843,7 @@ def start_query( # a LocalSession for unit testing. # TODO(swast): Add a timeout here? If the query is taking a long time, # maybe we just print the job metadata that we have so far? - table = self.to_ibis_expr(expose_hidden_cols=expose_extra_columns) + table = self._to_ibis_expr(expose_hidden_cols=expose_extra_columns) sql = self._session.ibis_client.compile(table) # type:ignore return self._session._start_query( sql=sql, @@ -833,7 +862,7 @@ def _reproject_to_table(self) -> ArrayValue: some operations such as window operations that cannot be used recursively in projections. """ - table = self.to_ibis_expr( + table = self._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True, ) @@ -912,7 +941,7 @@ def unpivot( Returns: ArrayValue: The unpivoted ArrayValue """ - table = self.to_ibis_expr(ordering_mode="offset_col") + table = self._to_ibis_expr(ordering_mode="offset_col") sub_expressions = [] # Use ibis memtable to infer type of rowlabels (if possible) @@ -1054,7 +1083,7 @@ def slice( start = start if (start is not None) else last_offset cond_list.append((start - expr_with_offsets.offsets) % (-step) == 0) - sliced_expr = expr_with_offsets.filter( + sliced_expr = expr_with_offsets._filter( functools.reduce(lambda x, y: x & y, cond_list) ) return sliced_expr if step > 0 else sliced_expr.reversed() diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index fb9ede9f4c9..e691a30f9cd 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -30,8 +30,6 @@ import geopandas as gpd # type: ignore import google.cloud.bigquery as bigquery -import ibis.expr.schema as ibis_schema -import ibis.expr.types as ibis_types import numpy import pandas as pd import pyarrow as pa # type: ignore @@ -42,6 +40,7 @@ import bigframes.core.indexes as indexes import bigframes.core.ordering as ordering import bigframes.core.utils +import bigframes.core.utils as utils import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -368,7 +367,10 @@ def reorder_levels(self, ids: typing.Sequence[str]): level_names = [self.col_id_to_index_name[index_id] for index_id in ids] return Block(self.expr, ids, self.column_labels, level_names) - def _to_dataframe(self, result, schema: ibis_schema.Schema) -> pd.DataFrame: + @classmethod + def _to_dataframe( + cls, result, schema: typing.Mapping[str, bigframes.dtypes.Dtype] + ) -> pd.DataFrame: """Convert BigQuery data to pandas DataFrame with specific dtypes.""" df = result.to_dataframe( bool_dtype=pd.BooleanDtype(), @@ -382,8 +384,8 @@ def _to_dataframe(self, result, schema: ibis_schema.Schema) -> pd.DataFrame: ) # Convert Geography column from StringDType to GeometryDtype. - for column_name, ibis_dtype in schema.items(): - if ibis_dtype.is_geospatial(): + for column_name, dtype in schema.items(): + if dtype == gpd.array.GeometryDtype(): df[column_name] = gpd.GeoSeries.from_wkt( # https://github.com/geopandas/geopandas/issues/1879 df[column_name].replace({numpy.nan: None}), @@ -473,7 +475,8 @@ def _compute_and_count( if sampling_method == _HEAD: total_rows = int(results_iterator.total_rows * fraction) results_iterator.max_results = total_rows - df = self._to_dataframe(results_iterator, expr.to_ibis_expr().schema()) + schema = dict(zip(self.value_columns, self.dtypes)) + df = self._to_dataframe(results_iterator, schema) if self.index_columns: df.set_index(list(self.index_columns), inplace=True) @@ -508,7 +511,8 @@ def _compute_and_count( ) else: total_rows = results_iterator.total_rows - df = self._to_dataframe(results_iterator, expr.to_ibis_expr().schema()) + schema = dict(zip(self.value_columns, self.dtypes)) + df = self._to_dataframe(results_iterator, schema) if self.index_columns: df.set_index(list(self.index_columns), inplace=True) @@ -639,13 +643,6 @@ def with_index_labels(self, value: typing.Sequence[Label]) -> Block: index_labels=tuple(value), ) - def get_value_col_exprs( - self, column_names: Optional[Sequence[str]] = None - ) -> List[ibis_types.Value]: - """Retrive value column expressions.""" - column_names = self.value_columns if column_names is None else column_names - return [self._expr.get_column(column_name) for column_name in column_names] - def apply_unary_op( self, column: str, op: ops.UnaryOp, result_label: Label = None ) -> typing.Tuple[Block, str]: @@ -816,20 +813,9 @@ def assign_label(self, column_id: str, new_label: Label) -> Block: ) return self.with_column_labels(new_labels) - def filter(self, column_name: str, keep_null: bool = False): - condition = typing.cast( - ibis_types.BooleanValue, self._expr.get_column(column_name) - ) - if keep_null: - condition = typing.cast( - ibis_types.BooleanValue, - condition.fillna( - typing.cast(ibis_types.BooleanScalar, ibis_types.literal(True)) - ), - ) - filtered_expr = self.expr.filter(condition) + def filter(self, column_id: str, keep_null: bool = False): return Block( - filtered_expr, + self._expr.filter(column_id, keep_null), index_columns=self.index_columns, column_labels=self.column_labels, index_labels=self.index.names, @@ -1436,6 +1422,50 @@ def is_monotonic_decreasing( ) -> bool: return self._is_monotonic(column_id, increasing=False) + def to_sql_query( + self, include_index: bool + ) -> typing.Tuple[str, list[str], list[Label]]: + """ + Compiles this DataFrame's expression tree to SQL, optionally + including index columns. + + Args: + include_index (bool): + whether to include index columns. + + Returns: + a tuple of (sql_string, index_column_id_list, index_column_label_list). + If include_index is set to False, index_column_id_list and index_column_label_list + return empty lists. + """ + array_value = self._expr + col_labels, idx_labels = list(self.column_labels), list(self.index_labels) + old_col_ids, old_idx_ids = list(self.value_columns), list(self.index_columns) + + if not include_index: + idx_labels, old_idx_ids = [], [] + array_value = array_value.drop_columns(self.index_columns) + + old_ids = old_idx_ids + old_col_ids + + new_col_ids, new_idx_ids = utils.get_standardized_ids(col_labels, idx_labels) + new_ids = new_idx_ids + new_col_ids + + substitutions = {} + for old_id, new_id in zip(old_ids, new_ids): + # TODO(swast): Do we need to further escape this, or can we rely on + # the BigQuery unicode column name feature? + substitutions[old_id] = new_id + + sql = array_value.to_sql( + ordering_mode="unordered", col_id_overrides=substitutions + ) + return ( + sql, + new_ids[: len(idx_labels)], + idx_labels, + ) + def _is_monotonic( self, column_ids: typing.Union[str, Sequence[str]], increasing: bool ) -> bool: diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 28bce05338c..a538c80711b 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -21,11 +21,11 @@ import pandas as pd import bigframes.constants as constants -import bigframes.core as core import bigframes.core.guid as guid import bigframes.core.indexes as indexes import bigframes.core.scalar import bigframes.dataframe +import bigframes.operations as ops import bigframes.series if typing.TYPE_CHECKING: @@ -59,35 +59,23 @@ def __setitem__(self, key, value) -> None: # Assume the key is for the index label. block = self._series._block - value_column = self._series._value - index_column = block.expr.get_column(block.index_columns[0]) - new_value = ( - ibis.case() - .when( - index_column == ibis.literal(key, index_column.type()), - ibis.literal(value, value_column.type()), - ) - .else_(value_column) - .end() - .name(value_column.get_name()) + value_column = self._series._value_column + index_column = block.index_columns[0] + + # if index == key return value else value_colum + block, insert_cond = block.apply_unary_op( + index_column, ops.partial_right(ops.eq_op, key) ) - all_columns = [] - for column in block.expr.columns: - if column.get_name() != value_column.get_name(): - all_columns.append(column) - else: - all_columns.append(new_value) - new_expr = block.expr.projection(all_columns) - - # TODO(tbergeron): Use block operators rather than directly building desired ibis expressions. - self._series._set_block( - core.blocks.Block( - new_expr, - self._series._block.index_columns, - self._series._block.column_labels, - self._series._block.index.names, - ) + block, result_id = block.apply_binary_op( + insert_cond, + self._series._value_column, + ops.partial_arg1(ops.where_op, value), ) + block = block.copy_values(result_id, value_column).drop_columns( + [insert_cond, result_id] + ) + + self._series._set_block(block) class IlocSeriesIndexer: diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 7aba71fd95e..434cc2cd794 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -122,11 +122,11 @@ def join_by_column( ), ) else: - left_table = left.to_ibis_expr( + left_table = left._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True, ) - right_table = right.to_ibis_expr( + right_table = right._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True, ) diff --git a/bigframes/core/scalar.py b/bigframes/core/scalar.py index 6dfbd31b77b..5db83b4a627 100644 --- a/bigframes/core/scalar.py +++ b/bigframes/core/scalar.py @@ -14,55 +14,7 @@ from __future__ import annotations -import typing -from typing import Any, Optional - -import google.cloud.bigquery as bigquery -import ibis.expr.types as ibis_types - -import bigframes -import bigframes.formatting_helpers as formatter - -if typing.TYPE_CHECKING: - import bigframes.session - - -class DeferredScalar: - """A deferred scalar object.""" - - def __init__(self, value: ibis_types.Scalar, session: bigframes.session.Session): - self._value = value - self._session = session - self._query_job: Optional[bigquery.QueryJob] = None - - @property - def query_job(self) -> Optional[bigquery.QueryJob]: - """BigQuery job metadata for the most recent query.""" - if self._query_job is None: - self._query_job = self._compute_dry_run() - return self._query_job - - def __repr__(self) -> str: - """Converts a Series to a string.""" - # TODO(swast): Add a timeout here? If the query is taking a long time, - # maybe we just print the job metadata that we have so far? - opts = bigframes.options.display - if opts.repr_mode == "deferred": - return formatter.repr_query_job(self.query_job) - else: - return repr(self.to_pandas()) - - def to_pandas(self) -> Any: - """Executes deferred operations and downloads the resulting scalar.""" - result, query_job = self._session._start_query(self._value.compile()) - self._query_job = query_job - df = self._session._rows_to_dataframe(result) - return df.iloc[0, 0] - - def _compute_dry_run(self): - job_config = bigquery.QueryJobConfig(dry_run=True) - return self._session._start_query(self._value.compile(), job_config=job_config) - +from typing import Any # All public APIs return Any at present # Later implementation may sometimes return a lazy scalar diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index de4adb912e6..46c27eaccbf 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -312,37 +312,7 @@ def _to_sql_query( If include_index is set to False, index_column_id_list and index_column_label_list return empty lists. """ - # Has to be unordered as it is impossible to order the sql without - # including metadata columns in selection with ibis. - ibis_expr = self._block.expr.to_ibis_expr(ordering_mode="unordered") - col_labels, idx_labels = list(self._block.column_labels), list( - self._block.index_labels - ) - old_col_ids, old_idx_ids = list(self._block.value_columns), list( - self._block.index_columns - ) - - if not include_index: - idx_labels, old_idx_ids = [], [] - ibis_expr = ibis_expr.drop(*self._block.index_columns) - - old_ids = old_idx_ids + old_col_ids - - new_col_ids, new_idx_ids = utils.get_standardized_ids(col_labels, idx_labels) - new_ids = new_idx_ids + new_col_ids - - substitutions = {} - for old_id, new_id in zip(old_ids, new_ids): - # TODO(swast): Do we need to further escape this, or can we rely on - # the BigQuery unicode column name feature? - substitutions[old_id] = new_id - - ibis_expr = ibis_expr.relabel(substitutions) - return ( - typing.cast(str, ibis_expr.compile()), - new_ids[: len(idx_labels)], - idx_labels, - ) + return self._block.to_sql_query(include_index) @property def sql(self) -> str: @@ -2340,8 +2310,7 @@ def _apply_unary_op(self, operation: ops.UnaryOp) -> DataFrame: def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: """Create query text representing this dataframe for I/O.""" - expr = self._block.expr - session = expr._session + array_value = self._block.expr columns = list(self._block.value_columns) column_labels = list(self._block.column_labels) # This code drops unnamed indexes to keep consistent with the behavior of @@ -2352,7 +2321,7 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: columns.extend(self._block.index_columns) column_labels.extend(self.index.names) else: - expr = expr.drop_columns(self._block.index_columns) + array_value = array_value.drop_columns(self._block.index_columns) # Make columns in SQL reflect _labels_ not _ids_. Note: This may use # the arbitrary unicode column labels feature in BigQuery, which is @@ -2365,19 +2334,17 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: } if ordering_id is not None: - ibis_expr = expr.to_ibis_expr( + return array_value.to_sql( ordering_mode="offset_col", col_id_overrides=id_overrides, order_col_name=ordering_id, ) else: - ibis_expr = expr.to_ibis_expr( + return array_value.to_sql( ordering_mode="unordered", col_id_overrides=id_overrides, ) - return session.ibis_client.compile(ibis_expr) # type: ignore - def _run_io_query( self, index: bool, @@ -2458,6 +2425,9 @@ def rank( df = self._drop_non_numeric() if numeric_only else self return DataFrame(block_ops.rank(df._block, method, na_option, ascending)) + def first_valid_index(self): + return + applymap = map def _slice( diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index c5c55607ae5..bc08298eb7b 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -206,7 +206,7 @@ def _as_ibis(self, x: ibis_types.Value): class LenOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.StringValue, x).length() + return typing.cast(ibis_types.StringValue, x).length().cast(ibis_dtypes.int64) class NotNullOp(UnaryOp): @@ -443,7 +443,7 @@ def _as_ibis(self, x: ibis_types.Value): ## Datetime Ops class DayOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).day() + return typing.cast(ibis_types.TimestampValue, x).day().cast(ibis_dtypes.int64) class DateOp(UnaryOp): @@ -453,32 +453,42 @@ def _as_ibis(self, x: ibis_types.Value): class DayofweekOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).day_of_week.index() + return ( + typing.cast(ibis_types.TimestampValue, x) + .day_of_week.index() + .cast(ibis_dtypes.int64) + ) class HourOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).hour() + return typing.cast(ibis_types.TimestampValue, x).hour().cast(ibis_dtypes.int64) class MinuteOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).minute() + return ( + typing.cast(ibis_types.TimestampValue, x).minute().cast(ibis_dtypes.int64) + ) class MonthOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).month() + return typing.cast(ibis_types.TimestampValue, x).month().cast(ibis_dtypes.int64) class QuarterOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).quarter() + return ( + typing.cast(ibis_types.TimestampValue, x).quarter().cast(ibis_dtypes.int64) + ) class SecondOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).second() + return ( + typing.cast(ibis_types.TimestampValue, x).second().cast(ibis_dtypes.int64) + ) class TimeOp(UnaryOp): @@ -488,7 +498,7 @@ def _as_ibis(self, x: ibis_types.Value): class YearOp(UnaryOp): def _as_ibis(self, x: ibis_types.Value): - return typing.cast(ibis_types.TimestampValue, x).year() + return typing.cast(ibis_types.TimestampValue, x).year().cast(ibis_dtypes.int64) # Parameterized ops @@ -969,6 +979,12 @@ def fillna_op( return x.fillna(typing.cast(ibis_types.Scalar, y)) +def round_op(x: ibis_types.Value, y: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).round( + digits=typing.cast(ibis_types.IntegerValue, y) + ) + + def clip_lower( value: ibis_types.Value, lower: ibis_types.Value, diff --git a/bigframes/series.py b/bigframes/series.py index 9db64fae9c9..8f3a24698b3 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -22,7 +22,6 @@ from typing import Any, Mapping, Optional, Tuple, Union import google.cloud.bigquery as bigquery -import ibis.expr.types as ibis_types import numpy import pandas import pandas.core.dtypes.common @@ -223,14 +222,6 @@ def __repr__(self) -> str: return repr(pandas_df.iloc[:, 0]) - def _to_ibis_expr(self): - """Creates an Ibis table expression representing the Series.""" - expr = self._block.expr.projection([self._value]) - ibis_expr = expr.to_ibis_expr()[self._value_column] - if self._name: - return ibis_expr.name(self._name) - return ibis_expr - def astype( self, dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], @@ -661,12 +652,7 @@ def abs(self) -> Series: return self._apply_unary_op(ops.abs_op) def round(self, decimals=0) -> "Series": - def round_op(x: ibis_types.Value, y: ibis_types.Value): - return typing.cast(ibis_types.NumericValue, x).round( - digits=typing.cast(ibis_types.IntegerValue, y) - ) - - return self._apply_binary_op(decimals, round_op) + return self._apply_binary_op(decimals, ops.round_op) def corr(self, other: Series, method="pearson", min_periods=None) -> float: """ diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 8f3e0beb0e4..ee0cefb3d28 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -103,7 +103,7 @@ def test_arrayvalue_to_ibis_expr_with_projection(): value.table["col2"].name("string_col"), ] ) - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(actual.columns) == 3 assert actual.columns[0] == "int64_col" assert actual.columns[1] == "literals" @@ -138,7 +138,7 @@ def test_arrayvalues_to_ibis_expr_with_concat(): total_ordering_columns=["col1"], ) expr = value.concat([value]) - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(actual.columns) == 3 # TODO(ashleyxu, b/299631930): test out the union expression assert actual.columns[0] == "column_0" @@ -175,7 +175,7 @@ def test_arrayvalues_to_ibis_expr_with_project_binary_op(): ) expr = value.project_binary_op("col2", "col3", ops.add_op, "col4") assert expr.columns[3].type().is_float64() - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(expr.columns) == 4 assert actual.columns[3] == "col4" @@ -194,7 +194,7 @@ def test_arrayvalues_to_ibis_expr_with_project_ternary_op(): ) expr = value.project_ternary_op("col2", "col3", "col4", ops.where_op, "col5") assert expr.columns[4].type().is_float64() - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(expr.columns) == 5 assert actual.columns[4] == "col5" @@ -215,7 +215,7 @@ def test_arrayvalue_to_ibis_expr_with_aggregate(): by_column_ids=["col1"], dropna=False, ) - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(expr.columns) == 2 assert actual.columns[0] == "col1" assert actual.columns[1] == "col4" @@ -234,7 +234,7 @@ def test_arrayvalue_to_ibis_expr_with_corr_aggregate(): total_ordering_columns=["col1"], ) expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]) - actual = expr.to_ibis_expr() + actual = expr._to_ibis_expr() assert len(expr.columns) == 1 assert actual.columns[0] == "col4" assert expr.columns[0].type().is_float64() From 3adc1b3aa3e2b218d4fa5debdaa4298276bdf801 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 21 Sep 2023 09:56:41 -0700 Subject: [PATCH 27/80] feat: add `items`, `apply` methods to `DataFrame`. (#43) Change-Id: Id3a0e78da3bb9ccce64e190f7797f737b239c33f Co-authored-by: Tim Swast --- bigframes/dataframe.py | 18 +++++++ tests/system/small/test_dataframe.py | 51 +++++++++++++++++++ .../bigframes_vendored/pandas/core/frame.py | 34 +++++++++++++ 3 files changed, 103 insertions(+) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 46c27eaccbf..0b741feff67 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1418,6 +1418,12 @@ def isin(self, values) -> DataFrame: f"isin(), you passed a [{type(values).__name__}]" ) + def items(self): + column_ids = self._block.value_columns + column_labels = self._block.column_labels + for col_id, col_label in zip(column_ids, column_labels): + yield col_label, bigframes.series.Series(self._block.select_column(col_id)) + def dropna( self, *, @@ -2382,6 +2388,18 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: ops.RemoteFunctionOp(func, apply_on_null=(na_action is None)) ) + def apply(self, func, *, args: typing.Tuple = (), **kwargs): + results = {name: func(col, *args, **kwargs) for name, col in self.items()} + if all( + [ + isinstance(val, bigframes.series.Series) or utils.is_list_like(val) + for val in results.values() + ] + ): + return DataFrame(data=results) + else: + return pandas.Series(data=results) + def drop_duplicates( self, subset: typing.Union[blocks.Label, typing.Sequence[blocks.Label]] = None, diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 6c96387e97e..3eeb368ad23 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -663,6 +663,57 @@ def test_df_bfill(scalars_dfs): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_apply_series_series_callable( + scalars_df_index, + scalars_pandas_df_index, +): + columns = ["int64_too", "int64_col"] + + def foo(series, arg1, arg2, *, kwarg1=0, kwarg2=0): + return series**2 + (arg1 * arg2 % 4) + (kwarg1 * kwarg2 % 7) + + bf_result = ( + scalars_df_index[columns] + .apply(foo, args=(33, 61), kwarg1=52, kwarg2=21) + .to_pandas() + ) + + pd_result = scalars_pandas_df_index[columns].apply( + foo, args=(33, 61), kwarg1=52, kwarg2=21 + ) + + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_apply_series_listlike_callable( + scalars_df_index, + scalars_pandas_df_index, +): + columns = ["int64_too", "int64_col"] + bf_result = ( + scalars_df_index[columns].apply(lambda x: [len(x), x.min(), 24]).to_pandas() + ) + + pd_result = scalars_pandas_df_index[columns].apply(lambda x: [len(x), x.min(), 24]) + + # Convert default pandas dtypes `int64` to match BigQuery DataFrames dtypes. + pd_result.index = pd_result.index.astype("Int64") + pd_result = pd_result.astype("Int64") + pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_apply_series_scalar_callable( + scalars_df_index, + scalars_pandas_df_index, +): + columns = ["int64_too", "int64_col"] + bf_result = scalars_df_index[columns].apply(lambda x: x.sum()) + + pd_result = scalars_pandas_df_index[columns].apply(lambda x: x.sum()) + + pandas.testing.assert_series_equal(bf_result, pd_result) + + def test_df_isin_list(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs values = ["Hello, World!", 55555, 2.51, pd.NA, True] diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 27cc2144e0f..9d26938e080 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -734,6 +734,18 @@ def isin(self, values): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def items(self): + """ + Iterate over (column name, Series) pairs. + + Iterates over the DataFrame columns, returning a tuple with + the column name and the content as a Series. + + Returns: + Iterator: Iterator of label, Series for each column. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Sorting @@ -1420,6 +1432,28 @@ def merge( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def apply(self, func, *, args=(), **kwargs): + """Apply a function along an axis of the DataFrame. + + Objects passed to the function are Series objects whose index is + the DataFrame's index (``axis=0``) the final return type + is inferred from the return type of the applied function. + + Args: + func (function): + Function to apply to each column or row. + args (tuple): + Positional arguments to pass to `func` in addition to the + array/series. + **kwargs: + Additional keyword arguments to pass as keywords arguments to + `func`. + + Returns: + pandas.Series or bigframes.DataFrame: Result of applying ``func`` along the given axis of the DataFrame. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # ndarray-like stats methods From 1a254a496633957b9506dd8392dcc6fd10762201 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 21 Sep 2023 10:42:59 -0700 Subject: [PATCH 28/80] feat: add index `dtype`, `astype`, `drop`, `fillna`, aggregate attributes. (#38) Change-Id: I4af249d10b2fcd779ad05d1f1d95049893e40135 --- bigframes/core/indexes/index.py | 155 ++++++++++++++++-- bigframes/series.py | 4 +- tests/system/small/test_index.py | 110 +++++++++++++ .../pandas/core/indexes/base.py | 120 ++++++++++++++ 4 files changed, 373 insertions(+), 16 deletions(-) diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 748a68c9447..c08c851c91d 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -17,7 +17,7 @@ from __future__ import annotations import typing -from typing import Callable, Tuple +from typing import Callable, Sequence, Tuple, Union import numpy as np import pandas @@ -26,7 +26,11 @@ import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.joins as joins +import bigframes.core.utils as utils +import bigframes.dtypes import bigframes.dtypes as bf_dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops import third_party.bigframes_vendored.pandas.core.indexes.base as vendored_pandas_index @@ -51,16 +55,34 @@ def names(self) -> typing.Sequence[blocks.Label]: @names.setter def names(self, values: typing.Sequence[blocks.Label]): - return self._data._set_block(self._data._get_block().with_index_labels(values)) + return self._data._set_block(self._block.with_index_labels(values)) @property def nlevels(self) -> int: return len(self._data._get_block().index_columns) + @property + def values(self) -> np.ndarray: + return self.to_numpy() + + @property + def ndim(self) -> int: + return 1 + @property def shape(self) -> typing.Tuple[int]: return (self._data._get_block().shape[0],) + @property + def dtype(self): + return self._block.index_dtypes[0] if self.nlevels == 1 else np.dtype("O") + + @property + def dtypes(self) -> pandas.Series: + return pandas.Series( + data=self._block.index_dtypes, index=self._block.index_labels # type:ignore + ) + @property def size(self) -> int: """Returns the size of the Index.""" @@ -103,23 +125,120 @@ def is_monotonic_decreasing(self) -> bool: @property def is_unique(self) -> bool: + # TODO: Cache this at block level + # Avoid circular imports + return not self.has_duplicates + + @property + def has_duplicates(self) -> bool: # TODO: Cache this at block level # Avoid circular imports import bigframes.core.block_transforms as block_ops import bigframes.dataframe as df - duplicates_block, _ = block_ops.indicate_duplicates( - self._data._get_block(), self._data._get_block().index_columns - ) - duplicates_block = duplicates_block.with_column_labels( - ["values", "is_duplicate"] + duplicates_block, indicator = block_ops.indicate_duplicates( + self._block, self._block.index_columns ) + duplicates_block = duplicates_block.select_columns( + [indicator] + ).with_column_labels(["is_duplicate"]) duplicates_df = df.DataFrame(duplicates_block) - return not duplicates_df["is_duplicate"].any() + return duplicates_df["is_duplicate"].any() + + @property + def _block(self) -> blocks.Block: + return self._data._get_block() + + def astype( + self, + dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], + ) -> Index: + if self.nlevels > 1: + raise TypeError("Multiindex does not support 'astype'") + return self._apply_unary_op(ops.AsTypeOp(dtype)) + + def all(self) -> bool: + if self.nlevels > 1: + raise TypeError("Multiindex does not support 'all'") + return typing.cast(bool, self._apply_aggregation(agg_ops.all_op)) + + def any(self) -> bool: + if self.nlevels > 1: + raise TypeError("Multiindex does not support 'any'") + return typing.cast(bool, self._apply_aggregation(agg_ops.any_op)) + + def nunique(self) -> int: + return typing.cast(int, self._apply_aggregation(agg_ops.nunique_op)) + + def max(self) -> typing.Any: + return self._apply_aggregation(agg_ops.max_op) + + def min(self) -> typing.Any: + return self._apply_aggregation(agg_ops.min_op) + + def fillna(self, value=None) -> Index: + if self.nlevels > 1: + raise TypeError("Multiindex does not support 'fillna'") + return self._apply_unary_op(ops.partial_right(ops.fillna_op, value)) + + def rename(self, name: Union[str, Sequence[str]]) -> Index: + names = [name] if isinstance(name, str) else list(name) + if len(names) != self.nlevels: + raise ValueError("'name' must be same length as levels") + + import bigframes.dataframe as df + + return Index(df.DataFrame(self._block.with_index_labels(names))) + + def drop( + self, + labels: typing.Any, + ) -> Index: + # ignore axis, columns params + block = self._block + level_id = self._block.index_columns[0] + if utils.is_list_like(labels): + block, inverse_condition_id = block.apply_unary_op( + level_id, ops.IsInOp(labels, match_nulls=True) + ) + block, condition_id = block.apply_unary_op( + inverse_condition_id, ops.invert_op + ) + else: + block, condition_id = block.apply_unary_op( + level_id, ops.partial_right(ops.ne_op, labels) + ) + block = block.filter(condition_id, keep_null=True) + block = block.drop_columns([condition_id]) + import bigframes.dataframe as df + + return Index(df.DataFrame(block.select_columns([]))) + + def _apply_unary_op( + self, + op: ops.UnaryOp, + ) -> Index: + """Applies a unary operator to the index.""" + block = self._block + result_ids = [] + for col in self._block.index_columns: + block, result_id = block.apply_unary_op(col, op) + result_ids.append(result_id) + + block = block.set_index(result_ids, index_labels=self._block.index_labels) + import bigframes.dataframe as df + + return Index(df.DataFrame(block)) + + def _apply_aggregation(self, op: agg_ops.AggregateOp) -> typing.Any: + if self.nlevels > 1: + raise NotImplementedError(f"Multiindex does not yet support {op.name}") + column_id = self._block.index_columns[0] + return self._block.get_stat(column_id, op) def __getitem__(self, key: int) -> typing.Any: if isinstance(key, int): - result_pd_df, _ = self._data._get_block().slice(key, key + 1, 1).to_pandas() + result_pd_df, _ = self._block.slice(key, key + 1, 1).to_pandas() if result_pd_df.empty: raise IndexError("single positional indexer is out-of-bounds") return result_pd_df.index[0] @@ -133,7 +252,7 @@ def to_pandas(self) -> pandas.Index: pandas.Index: A pandas Index with all of the labels from this Index. """ - return IndexValue(self._data._get_block()).to_pandas() + return IndexValue(self._block).to_pandas() def to_numpy(self, dtype=None, **kwargs) -> np.ndarray: return self.to_pandas().to_numpy(dtype, **kwargs) @@ -184,13 +303,15 @@ def __repr__(self) -> str: def to_pandas(self) -> pandas.Index: """Executes deferred operations and downloads the results.""" # Project down to only the index column. So the query can be cached to visualize other data. - index_column = self._block.index_columns[0] - expr = self._expr.projection([self._expr.get_any_column(index_column)]) + index_columns = list(self._block.index_columns) + expr = self._expr.projection( + [self._expr.get_any_column(col) for col in index_columns] + ) results, _ = expr.start_query() df = expr._session._rows_to_dataframe(results) - df.set_index(index_column) + df = df.set_index(index_columns) index = df.index - index.name = self._block._index_labels[0] + index.names = list(self._block._index_labels) return index def join( @@ -235,6 +356,12 @@ def resolve_level_name(self: IndexValue, label: blocks.Label) -> str: def is_uniquely_named(self: IndexValue): return len(set(self.names)) == len(self.names) + def _set_block(self, block: blocks.Block): + self._block = block + + def _get_block(self) -> blocks.Block: + return self._block + def join_mono_indexed( left: IndexValue, diff --git a/bigframes/series.py b/bigframes/series.py index 8f3a24698b3..8e47088c146 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -287,7 +287,7 @@ def drop( # ignore axis, columns params block = self._block level_id = self._resolve_levels(level or 0)[0] - if _is_list_like(labels): + if _is_list_like(index): block, inverse_condition_id = block.apply_unary_op( level_id, ops.IsInOp(index, match_nulls=True) ) @@ -296,7 +296,7 @@ def drop( ) else: block, condition_id = block.apply_unary_op( - level_id, ops.partial_right(ops.ne_op, labels) + level_id, ops.partial_right(ops.ne_op, index) ) block = block.filter(condition_id, keep_null=True) block = block.drop_columns([condition_id]) diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py index 558dd12e697..7f09e3a9d51 100644 --- a/tests/system/small/test_index.py +++ b/tests/system/small/test_index.py @@ -13,6 +13,7 @@ # limitations under the License. import numpy +import pandas as pd from tests.system.utils import assert_pandas_index_equal_ignore_index_type @@ -25,6 +26,44 @@ def test_get_index(scalars_df_index, scalars_pandas_df_index): assert_pandas_index_equal_ignore_index_type(bf_result, pd_result) +def test_index_has_duplicates(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.has_duplicates + pd_result = scalars_pandas_df_index.set_index("int64_col").index.has_duplicates + assert bf_result == pd_result + + +def test_index_values(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.index.values + pd_result = scalars_pandas_df_index.index.values + + # Numpy isn't equipped to compare non-numeric objects, so convert back to dataframe + pd.testing.assert_series_equal( + pd.Series(bf_result), pd.Series(pd_result), check_dtype=False + ) + + +def test_index_ndim(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.index.ndim + pd_result = scalars_pandas_df_index.index.ndim + + assert pd_result == bf_result + + +def test_index_dtype(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.index.dtype + pd_result = scalars_pandas_df_index.index.dtype + + assert pd_result == bf_result + + +def test_index_dtypes(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index(["string_col", "int64_too"]).index.dtypes + pd_result = scalars_pandas_df_index.set_index( + ["string_col", "int64_too"] + ).index.dtypes + pd.testing.assert_series_equal(bf_result, pd_result) + + def test_index_shape(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.index.shape pd_result = scalars_pandas_df_index.index.shape @@ -32,6 +71,77 @@ def test_index_shape(scalars_df_index, scalars_pandas_df_index): assert bf_result == pd_result +def test_index_astype(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.set_index("int64_col").index.astype("Float64").to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index("int64_col").index.astype("Float64") + pd.testing.assert_index_equal(bf_result, pd_result) + + +def test_index_any(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.any() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.any() + assert bf_result == pd_result + + +def test_index_all(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.all() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.all() + assert bf_result == pd_result + + +def test_index_max(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.max() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.max() + assert bf_result == pd_result + + +def test_index_min(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.min() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.min() + assert bf_result == pd_result + + +def test_index_nunique(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.nunique() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.nunique() + assert bf_result == pd_result + + +def test_index_fillna(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.fillna(42).to_pandas() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.fillna(42) + + pd.testing.assert_index_equal(bf_result, pd_result) + + +def test_index_drop(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.set_index("int64_col").index.drop([2, 314159]).to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index("int64_col").index.drop([2, 314159]) + pd.testing.assert_index_equal(bf_result, pd_result) + + +def test_index_rename(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index("int64_col").index.rename("name").to_pandas() + pd_result = scalars_pandas_df_index.set_index("int64_col").index.rename("name") + pd.testing.assert_index_equal(bf_result, pd_result) + + +def test_index_multi_rename(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index.set_index(["int64_col", "int64_too"]) + .index.rename(["new", "names"]) + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index( + ["int64_col", "int64_too"] + ).index.rename(["new", "names"]) + pd.testing.assert_index_equal(bf_result, pd_result) + + def test_index_len(scalars_df_index, scalars_pandas_df_index): bf_result = len(scalars_df_index.index) pd_result = len(scalars_pandas_df_index.index) diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py index 864007b7749..f89964e2208 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/base.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -1,4 +1,5 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexes/base.py +from __future__ import annotations from bigframes import constants @@ -14,6 +15,11 @@ def name(self): """Returns Index name.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property + def values(self): + """Return an array representing the data in the Index.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property def shape(self): """ @@ -31,6 +37,120 @@ def is_unique(self) -> bool: """Return if the index has unique values.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property + def has_duplicates(self) -> bool: + """Check if the Index has duplicate values.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def dtype(self): + """Return the dtype object of the underlying data.""" + + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def dtypes(self): + """Return the dtypes as a Series for the underlying MultiIndex.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def astype(self, dtype): + """Create an Index with values cast to dtypes. + + The class of a new Index is determined by dtype. When conversion is + impossible, a TypeError exception is raised. + + Args: + dtype (numpy dtype or pandas type): + + Returns: + Index: Index with values cast to specified dtype. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def all(self) -> bool: + """Return whether all elements are Truthy. + + Returns: + bool: A single element array-like may be converted to bool. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def any(self) -> bool: + """Return whether any element is Truthy. + + Returns: + bool: A single element array-like may be converted to bool. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def min(self): + """Return the minimum value of the Index. + + Returns: + scalar: Minimum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def max(self): + """Return the maximum value of the Index. + + Returns: + scalar: Maximum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def nunique(self) -> int: + """Return number of unique elements in the object. + + Excludes NA values by default. + + Returns: + int + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def fillna(self, value) -> Index: + """ + Fill NA/NaN values with the specified value. + + Args: + value (scalar): + Scalar value to use to fill holes (e.g. 0). + This value cannot be a list-likes. + + Returns: + Index + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def rename(self, name) -> Index: + """ + Alter Index or MultiIndex name. + + Able to set new names without level. Defaults to returning new index. + Length of names must match number of levels in MultiIndex. + + Args: + name (label or list of labels): + Name(s) to set. + + Returns: + Index: The same type as the caller. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def drop(self, labels) -> Index: + """ + Make new Index with passed list of labels deleted. + + Args: + labels (array-like or scalar): + + Returns: + Index: Will be same type as self + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def to_numpy(self, dtype): """ A NumPy ndarray representing the values in this Series or Index. From 5e199ecf1ecf13a68a2ed0dd4464afd9db977ab1 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 21 Sep 2023 14:28:13 -0500 Subject: [PATCH 29/80] perf: inline small `Series` and `DataFrames` in query text (#45) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prevents unnecessary load and query jobs. Towards internal issue 296474170 🦕 --- bigframes/core/__init__.py | 62 ++++++++++++++++++++----- bigframes/core/blocks.py | 55 +++++++++++----------- bigframes/dataframe.py | 4 +- bigframes/dtypes.py | 21 +++++---- bigframes/operations/base.py | 4 +- setup.py | 4 +- testing/constraints-3.9.txt | 2 +- tests/unit/core/__init__.py | 13 ++++++ tests/unit/core/test_blocks.py | 85 ++++++++++++++++++++++++++++++++++ 9 files changed, 194 insertions(+), 56 deletions(-) create mode 100644 tests/unit/core/__init__.py create mode 100644 tests/unit/core/test_blocks.py diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index dd91f80e63b..3b3754642e7 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -144,21 +144,56 @@ def mem_expr_from_pandas( """ Builds an in-memory only (SQL only) expr from a pandas dataframe. - Caution: If session is None, only a subset of expr functionality will be available (null Session is usually not supported). + Caution: If session is None, only a subset of expr functionality will + be available (null Session is usually not supported). """ - # must set non-null column labels. these are not the user-facing labels - pd_df = pd_df.set_axis( - [column or bigframes.core.guid.generate_guid() for column in pd_df.columns], - axis="columns", - ) + # We can't include any hidden columns in the ArrayValue constructor, so + # grab the column names before we add the hidden ordering column. + column_names = [str(column) for column in pd_df.columns] + # Make sure column names are all strings. + pd_df = pd_df.set_axis(column_names, axis="columns") pd_df = pd_df.assign(**{ORDER_ID_COLUMN: range(len(pd_df))}) + # ibis memtable cannot handle NA, must convert to None pd_df = pd_df.astype("object") # type: ignore pd_df = pd_df.where(pandas.notnull(pd_df), None) + + # NULL type isn't valid in BigQuery, so retry with an explicit schema in these cases. keys_memtable = ibis.memtable(pd_df) + schema = keys_memtable.schema() + new_schema = [] + for column_index, column in enumerate(schema): + if column == ORDER_ID_COLUMN: + new_type: ibis_dtypes.DataType = ibis_dtypes.int64 + else: + column_type = schema[column] + # The autodetected type might not be one we can support, such + # as NULL type for empty rows, so convert to a type we do + # support. + new_type = bigframes.dtypes.bigframes_dtype_to_ibis_dtype( + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(column_type) + ) + # TODO(swast): Ibis memtable doesn't use backticks in struct + # field names, so spaces and other characters aren't allowed in + # the memtable context. Blocked by + # https://github.com/ibis-project/ibis/issues/7187 + column = f"col_{column_index}" + new_schema.append((column, new_type)) + + # must set non-null column labels. these are not the user-facing labels + pd_df = pd_df.set_axis( + [column for column, _ in new_schema], + axis="columns", + ) + keys_memtable = ibis.memtable(pd_df, schema=ibis.schema(new_schema)) + return cls( session, # type: ignore # Session cannot normally be none, see "caution" above keys_memtable, + columns=[ + keys_memtable[f"col_{column_index}"].name(column) + for column_index, column in enumerate(column_names) + ], ordering=ExpressionOrdering( ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], total_ordering_columns=frozenset([ORDER_ID_COLUMN]), @@ -426,11 +461,16 @@ def shape(self) -> typing.Tuple[int, int]: width = len(self.columns) count_expr = self._to_ibis_expr(ordering_mode="unordered").count() sql = self._session.ibis_client.compile(count_expr) - row_iterator, _ = self._session._start_query( - sql=sql, - max_results=1, - ) - length = next(row_iterator)[0] + + # Support in-memory engines for hermetic unit tests. + if not isinstance(sql, str): + length = self._session.ibis_client.execute(count_expr) + else: + row_iterator, _ = self._session._start_query( + sql=sql, + max_results=1, + ) + length = next(row_iterator)[0] return (length, width) def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue: diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index e691a30f9cd..ad4f72070f5 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -44,6 +44,7 @@ import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops +import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common # Type constraint for wherever column labels are used Label = typing.Hashable @@ -1522,37 +1523,35 @@ def _is_monotonic( return result -def block_from_local(data, session=None, use_index=True) -> Block: - # TODO(tbergeron): Handle duplicate column labels +def block_from_local(data, session=None) -> Block: pd_data = pd.DataFrame(data) + columns = pd_data.columns - column_labels = list(pd_data.columns) - if not all((label is None) or isinstance(label, str) for label in column_labels): - raise NotImplementedError( - f"Only string column labels supported. {constants.FEEDBACK_LINK}" - ) + # Make a flattened version to treat as a table. + if len(pd_data.columns.names) > 1: + pd_data.columns = columns.to_flat_index() - if use_index: - if pd_data.index.nlevels > 1: - raise NotImplementedError( - f"multi-indices not supported. {constants.FEEDBACK_LINK}" - ) - index_label = pd_data.index.name - - index_id = guid.generate_guid() - pd_data = pd_data.reset_index(names=index_id) - keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) - return Block( - keys_expr, - column_labels=column_labels, - index_columns=[index_id], - index_labels=[index_label], - ) - else: - keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) - keys_expr, offsets_id = keys_expr.promote_offsets() - # Constructor will create default range index - return Block(keys_expr, index_columns=[offsets_id], column_labels=column_labels) + index_labels = list(pd_data.index.names) + # The ArrayValue layer doesn't know about indexes, so make sure indexes + # are real columns with unique IDs. + pd_data = pd_data.reset_index( + names=[f"level_{level}" for level in range(len(index_labels))] + ) + pd_data = pd_data.set_axis( + vendored_pandas_io_common.dedup_names( + list(pd_data.columns), is_potential_multiindex=False + ), + axis="columns", + ) + index_ids = pd_data.columns[: len(index_labels)] + + keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) + return Block( + keys_expr, + column_labels=columns, + index_columns=index_ids, + index_labels=index_labels, + ) def _align_block_to_schema( diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 0b741feff67..76377cd477f 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -65,9 +65,7 @@ # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. # TODO(tbergeron): Convert to bytes-based limit -# TODO(swast): Address issues with string escaping and empty tables before -# re-enabling inline data (ibis.memtable) feature. -MAX_INLINE_DF_SIZE = -1 +MAX_INLINE_DF_SIZE = 5000 LevelType = typing.Union[str, int] LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]] diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index af3209b0e1c..271b8aa2f24 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -96,13 +96,13 @@ ), ) -BIGFRAMES_TO_IBIS: Dict[Dtype, IbisDtype] = { +BIGFRAMES_TO_IBIS: Dict[Dtype, ibis_dtypes.DataType] = { pandas: ibis for ibis, pandas in BIDIRECTIONAL_MAPPINGS } -IBIS_TO_BIGFRAMES: Dict[ - Union[IbisDtype, ReadOnlyIbisDtype], Union[Dtype, np.dtype[Any]] -] = {ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS} +IBIS_TO_BIGFRAMES: Dict[ibis_dtypes.DataType, Union[Dtype, np.dtype[Any]]] = { + ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS +} # Allow REQUIRED fields to map correctly. IBIS_TO_BIGFRAMES.update( {ibis.copy(nullable=False): pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS} @@ -130,7 +130,7 @@ def ibis_dtype_to_bigframes_dtype( - ibis_dtype: Union[IbisDtype, ReadOnlyIbisDtype] + ibis_dtype: ibis_dtypes.DataType, ) -> Union[Dtype, np.dtype[Any]]: """Converts an Ibis dtype to a BigQuery DataFrames dtype @@ -155,6 +155,9 @@ def ibis_dtype_to_bigframes_dtype( if ibis_dtype in IBIS_TO_BIGFRAMES: return IBIS_TO_BIGFRAMES[ibis_dtype] + elif isinstance(ibis_dtype, ibis_dtypes.Null): + # Fallback to STRING for NULL values for most flexibility in SQL. + return IBIS_TO_BIGFRAMES[ibis_dtypes.string] else: raise ValueError( f"Unexpected Ibis data type {ibis_dtype}. {constants.FEEDBACK_LINK}" @@ -185,8 +188,8 @@ def ibis_table_to_canonical_types(table: ibis_types.Table) -> ibis_types.Table: def bigframes_dtype_to_ibis_dtype( - bigframes_dtype: Union[DtypeString, Dtype] -) -> IbisDtype: + bigframes_dtype: Union[DtypeString, Dtype, np.dtype[Any]] +) -> ibis_dtypes.DataType: """Converts a BigQuery DataFrames supported dtype to an Ibis dtype. Args: @@ -281,7 +284,9 @@ def literal_to_ibis_scalar( return scalar_expr -def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.Value: +def cast_ibis_value( + value: ibis_types.Value, to_type: ibis_dtypes.DataType +) -> ibis_types.Value: """Perform compatible type casts of ibis values Args: diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index 81a5bc4c412..add6af57f41 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -30,9 +30,7 @@ # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type. # TODO(tbergeron): Convert to bytes-based limit -# TODO(swast): Address issues with string escaping and empty tables before -# re-enabling inline data (ibis.memtable) feature. -MAX_INLINE_SERIES_SIZE = -1 +MAX_INLINE_SERIES_SIZE = 5000 class SeriesMethods: diff --git a/setup.py b/setup.py index 69b71c88f13..29eacb74a9a 100644 --- a/setup.py +++ b/setup.py @@ -44,12 +44,12 @@ "google-cloud-resource-manager >=1.10.3", "google-cloud-storage >=2.0.0", # TODO: Relax upper bound once we have fixed `system_prerelease` tests. - "ibis-framework[bigquery] >=6.0.0,<=6.1.0", + "ibis-framework[bigquery] >=6.2.0,<7.0.0dev", "pandas >=1.5.0", "pydata-google-auth >=1.8.2", "requests >=2.27.1", "scikit-learn >=1.2.2", - "sqlalchemy >=1.4,<3.0", + "sqlalchemy >=1.4,<3.0dev", "ipywidgets >=7.7.1", "humanize >= 4.6.0", ] diff --git a/testing/constraints-3.9.txt b/testing/constraints-3.9.txt index cd69d45dc92..f43d3b4ca01 100644 --- a/testing/constraints-3.9.txt +++ b/testing/constraints-3.9.txt @@ -45,7 +45,7 @@ greenlet==2.0.2 grpc-google-iam-v1==0.12.6 grpcio==1.53.0 grpcio-status==1.48.2 -ibis-framework==6.0.0 +ibis-framework==6.2.0 humanize==4.6.0 identify==2.5.22 idna==3.4 diff --git a/tests/unit/core/__init__.py b/tests/unit/core/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/unit/core/__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/tests/unit/core/test_blocks.py b/tests/unit/core/test_blocks.py new file mode 100644 index 00000000000..a7e9b5a84b7 --- /dev/null +++ b/tests/unit/core/test_blocks.py @@ -0,0 +1,85 @@ +# 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 pandas +import pandas.testing +import pytest + +import bigframes.core.blocks as blocks + +from .. import resources + + +@pytest.mark.parametrize( + ("data",), + ( + pytest.param( + {"test 1": [1, 2, 3], "test 2": [0.25, 0.5, 0.75]}, + id="dict_spaces_in_column_names", + ), + pytest.param( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], + id="nested_list", + ), + pytest.param( + pandas.concat( + [ + pandas.Series([1, 2, 3], name="some col"), + pandas.Series([2, 3, 4], name="some col"), + ], + axis="columns", + ), + id="duplicate_column_names", + ), + pytest.param( + pandas.DataFrame( + {"test": [1, 2, 3]}, + index=pandas.Index(["a", "b", "c"], name="string index"), + ), + id="string_index", + ), + pytest.param( + pandas.DataFrame( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], + columns=pandas.MultiIndex.from_tuples( + [(1, 1), (1, 2), (0, 0), (0, 1)], + names=["some level", "another level"], + ), + ), + marks=[ + pytest.mark.skipif( + tuple(pandas.__version__.split()) < ("2", "0", "0"), + reason="pandas 1.5.3 treats column MultiIndex as Index of tuples", + ), + ], + id="multiindex_columns", + ), + pytest.param( + pandas.DataFrame( + {"test": [1, 2, 3]}, + index=pandas.MultiIndex.from_tuples([(1, 1), (1, 2), (0, 0)]), + ), + id="multiindex_rows", + ), + ), +) +def test_block_from_local(data): + expected = pandas.DataFrame(data) + session = resources.create_pandas_session({}) + + block = blocks.block_from_local(data, session=session) + + pandas.testing.assert_index_equal(block.column_labels, expected.columns) + assert tuple(block.index_labels) == tuple(expected.index.names) + assert block.shape == expected.shape From 33274c2fc5035541fe9eb035fb5a1d92b35af4e5 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Thu, 21 Sep 2023 15:14:44 -0700 Subject: [PATCH 30/80] refactor: ml.sql to Object (#44) * refactor: ml.sql to Object Change-Id: Ibf795b81619778eaf28572fccd95a09b65f8ad58 --- bigframes/ml/cluster.py | 5 +- bigframes/ml/compose.py | 5 +- bigframes/ml/core.py | 301 +++++++++++++------------- bigframes/ml/decomposition.py | 5 +- bigframes/ml/ensemble.py | 14 +- bigframes/ml/forecasting.py | 12 +- bigframes/ml/globals.py | 30 +++ bigframes/ml/imported.py | 12 +- bigframes/ml/linear_model.py | 8 +- bigframes/ml/llm.py | 8 +- bigframes/ml/preprocessing.py | 16 +- bigframes/ml/sql.py | 326 +++++++++++++++-------------- tests/system/large/ml/test_core.py | 16 +- tests/system/small/ml/conftest.py | 3 +- tests/unit/ml/test_golden_sql.py | 52 +++-- tests/unit/ml/test_sql.py | 156 +++++++++----- 16 files changed, 559 insertions(+), 410 deletions(-) create mode 100644 bigframes/ml/globals.py diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 14cce2879e9..772b90f666e 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -22,7 +22,7 @@ from google.cloud import bigquery import bigframes -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.cluster._kmeans @@ -37,6 +37,7 @@ class KMeans( def __init__(self, n_clusters: int = 8): self.n_clusters = n_clusters self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> KMeans: @@ -66,7 +67,7 @@ def _fit( ) -> KMeans: (X,) = utils.convert_to_dataframe(X) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X_train=X, transforms=transforms, options=self._bqml_options, diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index a3d3503ad0a..02365f261c8 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -22,7 +22,7 @@ from typing import List, Optional, Tuple, Union from bigframes import constants -from bigframes.ml import base, core, preprocessing, utils +from bigframes.ml import base, core, globals, preprocessing, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.compose._column_transformer @@ -53,6 +53,7 @@ def __init__( # TODO: if any(transformers) has fitted raise warning self.transformers = transformers self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() # call self.transformers_ to check chained transformers self.transformers_ @@ -114,7 +115,7 @@ def fit( compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 667d42f7ee6..70be0d35ee1 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -36,6 +36,9 @@ class BqmlModel: def __init__(self, session: bigframes.Session, model: bigquery.Model): self._session = session self._model = model + self._model_manipulation_sql_generator = ml_sql.ModelManipulationSqlGenerator( + self.model_name + ) @property def session(self) -> bigframes.Session: @@ -85,18 +88,14 @@ def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: ml_sql.ml_predict( - model_name=self.model_name, source_sql=source_sql - ), + self._model_manipulation_sql_generator.ml_predict, ) def transform(self, input_data: bpd.DataFrame) -> bpd.DataFrame: # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: ml_sql.ml_transform( - model_name=self.model_name, source_sql=source_sql - ), + self._model_manipulation_sql_generator.ml_transform, ) def generate_text( @@ -107,10 +106,11 @@ def generate_text( # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: ml_sql.ml_generate_text( - model_name=self.model_name, + lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text( source_sql=source_sql, - struct_options=ml_sql.struct_options(**options), + struct_options=self._model_manipulation_sql_generator.struct_options( + **options + ), ), ) @@ -122,15 +122,16 @@ def generate_text_embedding( # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: ml_sql.ml_generate_text_embedding( - model_name=self.model_name, + lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text_embedding( source_sql=source_sql, - struct_options=ml_sql.struct_options(**options), + struct_options=self._model_manipulation_sql_generator.struct_options( + **options + ), ), ) def forecast(self) -> bpd.DataFrame: - sql = ml_sql.ml_forecast(self.model_name) + sql = self._model_manipulation_sql_generator.ml_forecast() return self._session.read_gbq(sql) def evaluate(self, input_data: Optional[bpd.DataFrame] = None): @@ -141,28 +142,28 @@ def evaluate(self, input_data: Optional[bpd.DataFrame] = None): if (input_data is not None) else (None, None, None) ) - sql = ml_sql.ml_evaluate(self.model_name, source_sql) + sql = self._model_manipulation_sql_generator.ml_evaluate(source_sql) return self._session.read_gbq(sql) - def centroids(self): + def centroids(self) -> bpd.DataFrame: assert self._model.model_type == "KMEANS" - sql = ml_sql.ml_centroids(self.model_name) + sql = self._model_manipulation_sql_generator.ml_centroids() return self._session.read_gbq(sql) - def principal_components(self): + def principal_components(self) -> bpd.DataFrame: assert self._model.model_type == "PCA" - sql = ml_sql.ml_principal_components(self.model_name) + sql = self._model_manipulation_sql_generator.ml_principal_components() return self._session.read_gbq(sql) - def principal_component_info(self): + def principal_component_info(self) -> bpd.DataFrame: assert self._model.model_type == "PCA" - sql = ml_sql.ml_principal_component_info(self.model_name) + sql = self._model_manipulation_sql_generator.ml_principal_component_info() return self._session.read_gbq(sql) @@ -187,8 +188,12 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: # truncate as Vertex ID only accepts 63 characters, easily exceeding the limit for temp models. # The possibility of conflicts should be low. vertex_ai_model_id = vertex_ai_model_id[:63] - options_sql = ml_sql.options(**{"vertex_ai_model_id": vertex_ai_model_id}) - sql = ml_sql.alter_model(self.model_name, options_sql=options_sql) + options_sql = self._model_manipulation_sql_generator.options( + **{"vertex_ai_model_id": vertex_ai_model_id} + ) + sql = self._model_manipulation_sql_generator.alter_model( + options_sql=options_sql + ) # Register the model and wait it to finish self._session._start_query(sql) @@ -196,137 +201,149 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: return self -def create_bqml_model( - X_train: bpd.DataFrame, - y_train: Optional[bpd.DataFrame] = None, - transforms: Optional[Iterable[str]] = None, - options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, -) -> BqmlModel: - """Create a session-temporary BQML model with the CREATE MODEL statement +class BqmlModelFactory: + def __init__(self): + model_id = self._create_temp_model_id() + self._model_creation_sql_generator = ml_sql.ModelCreationSqlGenerator(model_id) - Args: - X_train: features columns for training - y_train: labels columns for training, if applicable - transforms: an optional list of SQL expressions that implement preprocessing - on top of the input data. Generates a BQML TRANSFORM clause - options: a dict of options to configure the model. Generates a BQML OPTIONS - clause + def _create_temp_model_id(self) -> str: + return uuid.uuid4().hex - Returns: a BqmlModel, wrapping a trained model in BigQuery - """ - options = dict(options) - if y_train is None: - input_data = X_train - else: + def _reset_model_id(self): + self._model_creation_sql_generator._model_id = self._create_temp_model_id() + + def _create_model_with_sql(self, session: bigframes.Session, sql: str) -> BqmlModel: + # fit the model, synchronously + _, job = session._start_query(sql) + + # real model path in the session specific hidden dataset and table prefix + model_name_full = f"{job.destination.dataset_id}.{job.destination.table_id}" + model = session.bqclient.get_model(model_name_full) + + self._reset_model_id() + return BqmlModel(session, model) + + def create_model( + self, + X_train: bpd.DataFrame, + y_train: Optional[bpd.DataFrame] = None, + transforms: Optional[Iterable[str]] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, + ) -> BqmlModel: + """Create a session-temporary BQML model with the CREATE MODEL statement + + Args: + X_train: features columns for training + y_train: labels columns for training, if applicable + transforms: an optional list of SQL expressions that implement preprocessing + on top of the input data. Generates a BQML TRANSFORM clause + options: a dict of options to configure the model. Generates a BQML OPTIONS + clause + + Returns: a BqmlModel, wrapping a trained model in BigQuery + """ + options = dict(options) + if y_train is None: + input_data = X_train + else: + input_data = X_train.join(y_train, how="outer") + options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) + + session = X_train._session + + source_sql = input_data.sql + options_sql = self._model_creation_sql_generator.options(**options) + transform_sql = ( + self._model_creation_sql_generator.transform(*transforms) + if transforms is not None + else None + ) + sql = self._model_creation_sql_generator.create_model( + source_sql=source_sql, + transform_sql=transform_sql, + options_sql=options_sql, + ) + + return self._create_model_with_sql(session=session, sql=sql) + + def create_time_series_model( + self, + X_train: bpd.DataFrame, + y_train: bpd.DataFrame, + transforms: Optional[Iterable[str]] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, + ) -> BqmlModel: + assert ( + X_train.columns.size == 1 + ), "Time series timestamp input must only contain 1 column." + assert ( + y_train.columns.size == 1 + ), "Time stamp data input must only contain 1 column." + + options = dict(options) input_data = X_train.join(y_train, how="outer") - options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) - - session = X_train._session - - source_sql = input_data.sql - options_sql = ml_sql.options(**options) - transform_sql = ml_sql.transform(*transforms) if transforms is not None else None - sql = ml_sql.create_model( - model_name=_create_temp_model_name(), - source_sql=source_sql, - transform_sql=transform_sql, - options_sql=options_sql, - ) - - return _create_bqml_model_with_sql(session=session, sql=sql) - - -def create_bqml_time_series_model( - X_train: bpd.DataFrame, - y_train: bpd.DataFrame, - transforms: Optional[Iterable[str]] = None, - options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, -) -> BqmlModel: - - assert ( - X_train.columns.size == 1 - ), "Time series timestamp input must only contain 1 column." - assert ( - y_train.columns.size == 1 - ), "Time stamp data input must only contain 1 column." - - options = dict(options) - input_data = X_train.join(y_train, how="outer") - options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]}) - options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) - - session = X_train._session - - source_sql = input_data.sql - options_sql = ml_sql.options(**options) - - transform_sql = ml_sql.transform(*transforms) if transforms is not None else None - sql = ml_sql.create_model( - model_name=_create_temp_model_name(), - source_sql=source_sql, - transform_sql=transform_sql, - options_sql=options_sql, - ) - - return _create_bqml_model_with_sql(session=session, sql=sql) - - -def create_bqml_remote_model( - session: bigframes.Session, - connection_name: str, - options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, -) -> BqmlModel: - """Create a session-temporary BQML remote model with the CREATE MODEL statement - - Args: - connection_name: - a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection - options: - a dict of options to configure the model. Generates a BQML OPTIONS clause - - Returns: - BqmlModel: a BqmlModel wrapping a trained model in BigQuery - """ - options_sql = ml_sql.options(**options) - sql = ml_sql.create_remote_model( - model_name=_create_temp_model_name(), - connection_name=connection_name, - options_sql=options_sql, - ) + options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]}) + options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) - return _create_bqml_model_with_sql(session=session, sql=sql) + session = X_train._session + source_sql = input_data.sql + options_sql = self._model_creation_sql_generator.options(**options) -def create_bqml_imported_model( - session: bigframes.Session, - options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, -) -> BqmlModel: - """Create a session-temporary BQML imported model with the CREATE MODEL statement + transform_sql = ( + self._model_creation_sql_generator.transform(*transforms) + if transforms is not None + else None + ) + sql = self._model_creation_sql_generator.create_model( + source_sql=source_sql, + transform_sql=transform_sql, + options_sql=options_sql, + ) - Args: - options: a dict of options to configure the model. Generates a BQML OPTIONS - clause + return self._create_model_with_sql(session=session, sql=sql) - Returns: a BqmlModel, wrapping a trained model in BigQuery - """ - options_sql = ml_sql.options(**options) - sql = ml_sql.create_imported_model( - model_name=_create_temp_model_name(), - options_sql=options_sql, - ) + def create_remote_model( + self, + session: bigframes.Session, + connection_name: str, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, + ) -> BqmlModel: + """Create a session-temporary BQML remote model with the CREATE MODEL statement - return _create_bqml_model_with_sql(session=session, sql=sql) + Args: + connection_name: + a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection + options: + a dict of options to configure the model. Generates a BQML OPTIONS clause + + Returns: + BqmlModel: a BqmlModel wrapping a trained model in BigQuery + """ + options_sql = self._model_creation_sql_generator.options(**options) + sql = self._model_creation_sql_generator.create_remote_model( + connection_name=connection_name, + options_sql=options_sql, + ) + return self._create_model_with_sql(session=session, sql=sql) -def _create_temp_model_name() -> str: - return uuid.uuid4().hex + def create_imported_model( + self, + session: bigframes.Session, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, + ) -> BqmlModel: + """Create a session-temporary BQML imported model with the CREATE MODEL statement + Args: + options: a dict of options to configure the model. Generates a BQML OPTIONS + clause -def _create_bqml_model_with_sql(session: bigframes.Session, sql: str) -> BqmlModel: - # fit the model, synchronously - _, job = session._start_query(sql) + Returns: a BqmlModel, wrapping a trained model in BigQuery + """ + options_sql = self._model_creation_sql_generator.options(**options) + sql = self._model_creation_sql_generator.create_imported_model( + options_sql=options_sql, + ) - # real model path in the session specific hidden dataset and table prefix - model_name_full = f"{job.destination.dataset_id}.{job.destination.table_id}" - model = session.bqclient.get_model(model_name_full) - return BqmlModel(session, model) + return self._create_model_with_sql(session=session, sql=sql) diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 0cfe3b3ddf9..8e6be6d28c3 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -22,7 +22,7 @@ from google.cloud import bigquery import bigframes -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.decomposition._pca @@ -36,6 +36,7 @@ class PCA( def __init__(self, n_components: int = 3): self.n_components = n_components self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> PCA: @@ -60,7 +61,7 @@ def _fit( ) -> PCA: (X,) = utils.convert_to_dataframe(X) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X_train=X, transforms=transforms, options={ diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index b0f3e5f0813..113ad872b5c 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -22,7 +22,7 @@ from google.cloud import bigquery import bigframes -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.ensemble._forest import third_party.bigframes_vendored.xgboost.sklearn @@ -95,6 +95,7 @@ def __init__( self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -151,7 +152,7 @@ def _fit( ) -> XGBRegressor: X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, @@ -259,6 +260,7 @@ def __init__( self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -315,7 +317,7 @@ def _fit( ) -> XGBClassifier: X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, @@ -416,6 +418,7 @@ def __init__( self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -469,7 +472,7 @@ def _fit( ) -> RandomForestRegressor: X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, @@ -584,6 +587,7 @@ def __init__( self.enable_global_explain = enable_global_explain self.xgboost_version = xgboost_version self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -637,7 +641,7 @@ def _fit( ) -> RandomForestClassifier: X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index 22d81294fc3..8a6de1dd81f 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -21,7 +21,7 @@ from google.cloud import bigquery import bigframes -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd _PREDICT_OUTPUT_COLUMNS = ["forecast_timestamp", "forecast_value"] @@ -32,6 +32,7 @@ class ARIMAPlus(base.SupervisedTrainablePredictor): def __init__(self): self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> ARIMAPlus: @@ -69,9 +70,16 @@ def _fit( Returns: ARIMAPlus: Fitted estimator. """ + if X.columns.size != 1: + raise ValueError( + "Time series timestamp input X must only contain 1 column." + ) + if y.columns.size != 1: + raise ValueError("Time series data input y must only contain 1 column.") + X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_time_series_model( + self._bqml_model = self._bqml_model_factory.create_time_series_model( X, y, transforms=transforms, diff --git a/bigframes/ml/globals.py b/bigframes/ml/globals.py new file mode 100644 index 00000000000..c139476daaa --- /dev/null +++ b/bigframes/ml/globals.py @@ -0,0 +1,30 @@ +# 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. + +"""Global Singletons for BigQuery DataFrame ML.""" + +from bigframes.ml import core, sql + +_BASE_SQL_GENERATOR = sql.BaseSqlGenerator() +_BQML_MODEL_FACTORY = core.BqmlModelFactory() + + +def base_sql_generator() -> sql.BaseSqlGenerator: + """Base SQL Generator.""" + return _BASE_SQL_GENERATOR + + +def bqml_model_factory() -> core.BqmlModelFactory: + """BQML Model Factory""" + return _BQML_MODEL_FACTORY diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py index d4571eb3e5b..fb8aa98befd 100644 --- a/bigframes/ml/imported.py +++ b/bigframes/ml/imported.py @@ -21,7 +21,7 @@ from google.cloud import bigquery import bigframes -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd @@ -42,10 +42,13 @@ def __init__( self.session = session or bpd.get_global_session() self.model_path = model_path self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() def _create_bqml_model(self): options = {"model_type": "TENSORFLOW", "model_path": self.model_path} - return core.create_bqml_imported_model(session=self.session, options=options) + return self._bqml_model_factory.create_imported_model( + session=self.session, options=options + ) @classmethod def _from_bq( @@ -124,10 +127,13 @@ def __init__( self.session = session or bpd.get_global_session() self.model_path = model_path self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() def _create_bqml_model(self): options = {"model_type": "ONNX", "model_path": self.model_path} - return core.create_bqml_imported_model(session=self.session, options=options) + return self._bqml_model_factory.create_imported_model( + session=self.session, options=options + ) @classmethod def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> ONNXModel: diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index f27b798eea1..f11879500b9 100644 --- a/bigframes/ml/linear_model.py +++ b/bigframes/ml/linear_model.py @@ -23,7 +23,7 @@ import bigframes import bigframes.constants as constants -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.linear_model._base import third_party.bigframes_vendored.sklearn.linear_model._logistic @@ -80,6 +80,7 @@ def __init__( self.calculate_p_values = calculate_p_values self.enable_global_explain = enable_global_explain self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -130,7 +131,7 @@ def _fit( ) -> LinearRegression: X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, @@ -204,6 +205,7 @@ def __init__( self.class_weights = class_weights self._auto_class_weight = class_weights == "balanced" self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() @classmethod def _from_bq( @@ -248,7 +250,7 @@ def _fit( """Fit model with transforms.""" X, y = utils.convert_to_dataframe(X, y) - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, y, transforms=transforms, diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 973fbf2ad96..c86e5fb3b67 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -21,7 +21,7 @@ import bigframes from bigframes import clients, constants from bigframes.core import blocks -from bigframes.ml import base, core, utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" @@ -52,6 +52,7 @@ def __init__( self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient ) + self._bqml_model_factory = globals.bqml_model_factory() self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): @@ -76,7 +77,7 @@ def _create_bqml_model(self): "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, } - return core.create_bqml_remote_model( + return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -183,6 +184,7 @@ def __init__( self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient ) + self._bqml_model_factory = globals.bqml_model_factory() self._bqml_model: core.BqmlModel = self._create_bqml_model() def _create_bqml_model(self): @@ -207,7 +209,7 @@ def _create_bqml_model(self): "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, } - return core.create_bqml_remote_model( + return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 8add7bdd767..cd4ae27b8c8 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -20,9 +20,7 @@ import typing from typing import Any, cast, List, Literal, Optional, Tuple, Union -from bigframes.ml import base, core -from bigframes.ml import sql as ml_sql -from bigframes.ml import utils +from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.preprocessing._data import third_party.bigframes_vendored.sklearn.preprocessing._encoder @@ -38,6 +36,8 @@ class StandardScaler( def __init__(self): self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() # TODO(garrettwu): implement __hash__ def __eq__(self, other: Any) -> bool: @@ -53,7 +53,7 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: Returns: a list of tuples of (sql_expression, output_name)""" return [ ( - ml_sql.ml_standard_scaler(column, f"scaled_{column}"), + self._base_sql_generator.ml_standard_scaler(column, f"scaled_{column}"), f"scaled_{column}", ) for column in columns @@ -81,7 +81,7 @@ def fit( compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, @@ -132,6 +132,8 @@ def __init__( self.min_frequency = min_frequency self.max_categories = max_categories self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() # TODO(garrettwu): implement __hash__ def __eq__(self, other: Any) -> bool: @@ -167,7 +169,7 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: ) return [ ( - ml_sql.ml_one_hot_encoder( + self._base_sql_generator.ml_one_hot_encoder( column, drop, top_k, frequency_threshold, f"onehotencoded_{column}" ), f"onehotencoded_{column}", @@ -206,7 +208,7 @@ def fit( compiled_transforms = self._compile_to_sql(X.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] - self._bqml_model = core.create_bqml_model( + self._bqml_model = self._bqml_model_factory.create_model( X, options={"model_type": "transform_only"}, transforms=transform_sqls, diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index feb7ff7835b..7cf030485bd 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -21,174 +21,180 @@ import bigframes.constants as constants -def _encode_value(v: Union[str, int, float, Iterable[str]]) -> str: - """Encode a parameter value for SQL""" - if isinstance(v, str): - return f'"{v}"' - elif isinstance(v, int) or isinstance(v, float): - return f"{v}" - elif isinstance(v, Iterable): - inner = ", ".join([_encode_value(x) for x in v]) - return f"[{inner}]" - else: - raise ValueError(f"Unexpected value type. {constants.FEEDBACK_LINK}") - - -def _build_parameters(**kwargs: Union[str, int, float, Iterable[str]]) -> str: - """Encode a dict of values into a formatted Iterable of key-value pairs for SQL""" - indent_str = " " - param_strs = [f"{k}={_encode_value(v)}" for k, v in kwargs.items()] - return "\n" + indent_str + f",\n{indent_str}".join(param_strs) - - -def options(**kwargs: Union[str, int, float, Iterable[str]]) -> str: - """Encode the OPTIONS clause for BQML""" - return f"OPTIONS({_build_parameters(**kwargs)})" - - -def _build_structs(**kwargs: Union[int, float]) -> str: - """Encode a dict of values into a formatted STRUCT items for SQL""" - indent_str = " " - param_strs = [f"{v} AS {k}" for k, v in kwargs.items()] - return "\n" + indent_str + f",\n{indent_str}".join(param_strs) - - -def struct_options(**kwargs: Union[int, float]) -> str: - """Encode a BQ STRUCT as options.""" - return f"STRUCT({_build_structs(**kwargs)})" - - -def _build_expressions(*expr_sqls: str) -> str: - """Encode a Iterable of SQL expressions into a formatted Iterable for SQL""" - indent_str = " " - return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls) - - -def transform(*expr_sqls: str) -> str: - """Encode the TRANSFORM clause for BQML""" - return f"TRANSFORM({_build_expressions(*expr_sqls)})" - - -def connection(conn_name: str) -> str: - """Encode the REMOTE WITH CONNECTION clause for BQML. conn_name is of the format ...""" - return f"REMOTE WITH CONNECTION `{conn_name}`" - - -def ml_standard_scaler(numeric_expr_sql: str, name: str) -> str: - """Encode ML.STANDARD_SCALER for BQML""" - return f"""ML.STANDARD_SCALER({numeric_expr_sql}) OVER() AS {name}""" - - -def ml_one_hot_encoder( - numeric_expr_sql: str, drop: str, top_k: int, frequency_threshold: int, name: str -) -> str: - """Encode ML.ONE_HOT_ENCODER for BQML. - https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder for params.""" - return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}, '{drop}', {top_k}, {frequency_threshold}) OVER() AS {name}""" - - -def create_model( - model_name: str, - source_sql: str, - transform_sql: Optional[str] = None, - options_sql: Optional[str] = None, -) -> str: - """Encode the CREATE TEMP MODEL statement for BQML""" - parts = [f"CREATE TEMP MODEL `{model_name}`"] - if transform_sql: - parts.append(transform_sql) - if options_sql: - parts.append(options_sql) - parts.append(f"AS {source_sql}") - return "\n".join(parts) - - -def create_remote_model( - model_name: str, - connection_name: str, - options_sql: Optional[str] = None, -) -> str: - """Encode the CREATE TEMP MODEL statement for BQML remote model.""" - parts = [f"CREATE TEMP MODEL `{model_name}`"] - parts.append(connection(connection_name)) - if options_sql: - parts.append(options_sql) - return "\n".join(parts) - - -def create_imported_model( - model_name: str, - options_sql: Optional[str] = None, -) -> str: - """Encode the CREATE TEMP MODEL statement for BQML remote model.""" - parts = [f"CREATE TEMP MODEL `{model_name}`"] - if options_sql: - parts.append(options_sql) - return "\n".join(parts) - - -def alter_model( - model_name: str, - options_sql: str, -) -> str: - """Encode the ALTER MODEL statement for BQML""" - parts = [f"ALTER MODEL `{model_name}`"] - parts.append(f"SET {options_sql}") - return "\n".join(parts) - - -def ml_evaluate(model_name: str, source_sql: Optional[str] = None) -> str: - """Encode ML.EVALUATE for BQML""" - if source_sql is None: - return f"""SELECT * FROM ML.EVALUATE(MODEL `{model_name}`)""" - else: - return f"""SELECT * FROM ML.EVALUATE(MODEL `{model_name}`, +class BaseSqlGenerator: + """Generate base SQL strings for ML. Model name isn't needed in this class.""" + + # General methods + def encode_value(self, v: Union[str, int, float, Iterable[str]]) -> str: + """Encode a parameter value for SQL""" + if isinstance(v, str): + return f'"{v}"' + elif isinstance(v, int) or isinstance(v, float): + return f"{v}" + elif isinstance(v, Iterable): + inner = ", ".join([self.encode_value(x) for x in v]) + return f"[{inner}]" + else: + raise ValueError(f"Unexpected value type. {constants.FEEDBACK_LINK}") + + def build_parameters(self, **kwargs: Union[str, int, float, Iterable[str]]) -> str: + """Encode a dict of values into a formatted Iterable of key-value pairs for SQL""" + indent_str = " " + param_strs = [f"{k}={self.encode_value(v)}" for k, v in kwargs.items()] + return "\n" + indent_str + f",\n{indent_str}".join(param_strs) + + def build_structs(self, **kwargs: Union[int, float]) -> str: + """Encode a dict of values into a formatted STRUCT items for SQL""" + indent_str = " " + param_strs = [f"{v} AS {k}" for k, v in kwargs.items()] + return "\n" + indent_str + f",\n{indent_str}".join(param_strs) + + def build_expressions(self, *expr_sqls: str) -> str: + """Encode a Iterable of SQL expressions into a formatted Iterable for SQL""" + indent_str = " " + return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls) + + def options(self, **kwargs: Union[str, int, float, Iterable[str]]) -> str: + """Encode the OPTIONS clause for BQML""" + return f"OPTIONS({self.build_parameters(**kwargs)})" + + def struct_options(self, **kwargs: Union[int, float]) -> str: + """Encode a BQ STRUCT as options.""" + return f"STRUCT({self.build_structs(**kwargs)})" + + # Connection + def connection(self, conn_name: str) -> str: + """Encode the REMOTE WITH CONNECTION clause for BQML. conn_name is of the format ...""" + return f"REMOTE WITH CONNECTION `{conn_name}`" + + # Transformers + def transform(self, *expr_sqls: str) -> str: + """Encode the TRANSFORM clause for BQML""" + return f"TRANSFORM({self.build_expressions(*expr_sqls)})" + + def ml_standard_scaler(self, numeric_expr_sql: str, name: str) -> str: + """Encode ML.STANDARD_SCALER for BQML""" + return f"""ML.STANDARD_SCALER({numeric_expr_sql}) OVER() AS {name}""" + + def ml_one_hot_encoder( + self, + numeric_expr_sql: str, + drop: str, + top_k: int, + frequency_threshold: int, + name: str, + ) -> str: + """Encode ML.ONE_HOT_ENCODER for BQML. + https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder for params.""" + return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}, '{drop}', {top_k}, {frequency_threshold}) OVER() AS {name}""" + + +class ModelCreationSqlGenerator(BaseSqlGenerator): + """Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id.""" + + def __init__(self, model_id: str): + self._model_id = model_id + + # Model create and alter + def create_model( + self, + source_sql: str, + transform_sql: Optional[str] = None, + options_sql: Optional[str] = None, + ) -> str: + """Encode the CREATE TEMP MODEL statement for BQML""" + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] + if transform_sql: + parts.append(transform_sql) + if options_sql: + parts.append(options_sql) + parts.append(f"AS {source_sql}") + return "\n".join(parts) + + def create_remote_model( + self, + connection_name: str, + options_sql: Optional[str] = None, + ) -> str: + """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] + parts.append(self.connection(connection_name)) + if options_sql: + parts.append(options_sql) + return "\n".join(parts) + + def create_imported_model( + self, + options_sql: Optional[str] = None, + ) -> str: + """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] + if options_sql: + parts.append(options_sql) + return "\n".join(parts) + + +class ModelManipulationSqlGenerator(BaseSqlGenerator): + """Sql generator for manipulating a model entity. Model name is the fully model path of project_id.dataset_id.model_id.""" + + def __init__(self, model_name: str): + self._model_name = model_name + + # Alter model + def alter_model( + self, + options_sql: str, + ) -> str: + """Encode the ALTER MODEL statement for BQML""" + parts = [f"ALTER MODEL `{self._model_name}`"] + parts.append(f"SET {options_sql}") + return "\n".join(parts) + + # ML prediction TVFs + def ml_predict(self, source_sql: str) -> str: + """Encode ML.PREDICT for BQML""" + return f"""SELECT * FROM ML.PREDICT(MODEL `{self._model_name}`, ({source_sql}))""" + def ml_forecast(self) -> str: + """Encode ML.FORECAST for BQML""" + return f"""SELECT * FROM ML.FORECAST(MODEL `{self._model_name}`)""" -def ml_centroids( - model_name: str, -) -> str: - """Encode ML.CENTROIDS for BQML""" - return f"""SELECT * FROM ML.CENTROIDS(MODEL `{model_name}`)""" - - -def ml_predict(model_name: str, source_sql: str) -> str: - """Encode ML.PREDICT for BQML""" - return f"""SELECT * FROM ML.PREDICT(MODEL `{model_name}`, - ({source_sql}))""" - - -def ml_transform(model_name: str, source_sql: str) -> str: - """Encode ML.TRANSFORM for BQML""" - return f"""SELECT * FROM ML.TRANSFORM(MODEL `{model_name}`, - ({source_sql}))""" - - -def ml_generate_text(model_name: str, source_sql: str, struct_options: str) -> str: - """Encode ML.GENERATE_TEXT for BQML""" - return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL `{model_name}`, + def ml_generate_text(self, source_sql: str, struct_options: str) -> str: + """Encode ML.GENERATE_TEXT for BQML""" + return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL `{self._model_name}`, ({source_sql}), {struct_options})""" - -def ml_generate_text_embedding( - model_name: str, source_sql: str, struct_options: str -) -> str: - """Encode ML.GENERATE_TEXT_EMBEDDING for BQML""" - return f"""SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `{model_name}`, + def ml_generate_text_embedding(self, source_sql: str, struct_options: str) -> str: + """Encode ML.GENERATE_TEXT_EMBEDDING for BQML""" + return f"""SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `{self._model_name}`, ({source_sql}), {struct_options})""" + # ML evaluation TVFs + def ml_evaluate(self, source_sql: Optional[str] = None) -> str: + """Encode ML.EVALUATE for BQML""" + if source_sql is None: + return f"""SELECT * FROM ML.EVALUATE(MODEL `{self._model_name}`)""" + else: + return f"""SELECT * FROM ML.EVALUATE(MODEL `{self._model_name}`, + ({source_sql}))""" -def ml_forecast(model_name: str) -> str: - """Encode ML.FORECAST for BQML""" - return f"""SELECT * FROM ML.FORECAST(MODEL `{model_name}`)""" - + def ml_centroids(self) -> str: + """Encode ML.CENTROIDS for BQML""" + return f"""SELECT * FROM ML.CENTROIDS(MODEL `{self._model_name}`)""" -def ml_principal_components(model_name: str) -> str: - """Encode ML.PRINCIPAL_COMPONENTS for BQML""" - return f"""SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `{model_name}`)""" + def ml_principal_components(self) -> str: + """Encode ML.PRINCIPAL_COMPONENTS for BQML""" + return f"""SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `{self._model_name}`)""" + def ml_principal_component_info(self) -> str: + """Encode ML.PRINCIPAL_COMPONENT_INFO for BQML""" + return ( + f"""SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `{self._model_name}`)""" + ) -def ml_principal_component_info(model_name: str) -> str: - """Encode ML.PRINCIPAL_COMPONENT_INFO for BQML""" - return f"""SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `{model_name}`)""" + # ML transform TVF, that require a transform_only type model + def ml_transform(self, source_sql: str) -> str: + """Encode ML.TRANSFORM for BQML""" + return f"""SELECT * FROM ML.TRANSFORM(MODEL `{self._model_name}`, + ({source_sql}))""" diff --git a/tests/system/large/ml/test_core.py b/tests/system/large/ml/test_core.py index 133af2dae4a..3b30d7eb1d9 100644 --- a/tests/system/large/ml/test_core.py +++ b/tests/system/large/ml/test_core.py @@ -14,8 +14,7 @@ import pandas -import bigframes.ml.core -import bigframes.ml.sql +from bigframes.ml import globals def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_df): @@ -32,7 +31,7 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d ] y_train = df[["body_mass_g"]] - model = bigframes.ml.core.create_bqml_model( + model = globals.bqml_model_factory().create_model( X_train, y_train, options={"model_type": "linear_reg"} ) @@ -84,6 +83,9 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d def test_bqml_manual_preprocessing_e2e( session, dataset_id, penguins_df_default_index, new_penguins_df ): + base_sql_generator = globals.base_sql_generator() + bqml_model_factory = globals.bqml_model_factory() + df = penguins_df_default_index.dropna() X_train = df[ [ @@ -94,12 +96,12 @@ def test_bqml_manual_preprocessing_e2e( ] y_train = df[["body_mass_g"]] transforms = [ - bigframes.ml.sql.ml_standard_scaler(column, column) + base_sql_generator.ml_standard_scaler(column, column) for column in X_train.columns.tolist() ] transforms.extend(y_train.columns.tolist()) options = {"model_type": "linear_reg"} - model = bigframes.ml.core.create_bqml_model( + model = bqml_model_factory.create_model( X_train, y_train, transforms=transforms, options=options ) @@ -150,8 +152,10 @@ def test_bqml_manual_preprocessing_e2e( def test_bqml_standalone_transform(penguins_df_default_index, new_penguins_df): + bqml_model_factory = globals.bqml_model_factory() + X = penguins_df_default_index[["culmen_length_mm", "species"]] - model = bigframes.ml.core.create_bqml_model( + model = bqml_model_factory.create_model( X, options={"model_type": "transform_only"}, transforms=[ diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 9ca5a2fd0ef..1dd1c813b81 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -25,6 +25,7 @@ decomposition, ensemble, forecasting, + globals, imported, linear_model, llm, @@ -202,7 +203,7 @@ def bqml_palm2_text_generator_model(session, bq_connection) -> core.BqmlModel: options = { "remote_service_type": "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1", } - return core.create_bqml_remote_model( + return globals.bqml_model_factory().create_remote_model( session=session, connection_name=bq_connection, options=options ) diff --git a/tests/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py index 584d080d428..3ca7e144a53 100644 --- a/tests/unit/ml/test_golden_sql.py +++ b/tests/unit/ml/test_golden_sql.py @@ -62,38 +62,44 @@ def mock_X(mock_y, mock_session): @pytest.fixture -def bqml_model(mock_session): - bqml_model = core.BqmlModel( - mock_session, bigquery.Model("model_project.model_dataset.model_name") +def bqml_model_factory(mocker: pytest_mock.MockerFixture): + mocker.patch( + "bigframes.ml.core.BqmlModelFactory._create_temp_model_id", + return_value="temp_model_id", ) + bqml_model_factory = core.BqmlModelFactory() - return bqml_model + return bqml_model_factory @pytest.fixture -def ml_mocker(mocker: pytest_mock.MockerFixture): - mocker.patch( - "bigframes.ml.core._create_temp_model_name", return_value="temp_model_name" +def bqml_model(mock_session): + bqml_model = core.BqmlModel( + mock_session, bigquery.Model("model_project.model_dataset.model_id") ) - return mocker + return bqml_model -def test_linear_regression_default_fit(ml_mocker, mock_session, mock_X, mock_y): +def test_linear_regression_default_fit( + bqml_model_factory, mock_session, mock_X, mock_y +): model = linear_model.LinearRegression() + model._bqml_model_factory = bqml_model_factory model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=True,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + 'CREATE TEMP MODEL `temp_model_id`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=True,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' ) -def test_linear_regression_params_fit(ml_mocker, mock_session, mock_X, mock_y): +def test_linear_regression_params_fit(bqml_model_factory, mock_session, mock_X, mock_y): model = linear_model.LinearRegression(fit_intercept=False) + model._bqml_model_factory = bqml_model_factory model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=False,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + 'CREATE TEMP MODEL `temp_model_id`\nOPTIONS(\n model_type="LINEAR_REG",\n data_split_method="NO_SPLIT",\n optimize_strategy="normal_equation",\n fit_intercept=False,\n l2_reg=0.0,\n max_iterations=20,\n learn_rate_strategy="line_search",\n early_stop=True,\n min_rel_progress=0.01,\n ls_init_learn_rate=0.1,\n calculate_p_values=False,\n enable_global_explain=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' ) @@ -103,7 +109,7 @@ def test_linear_regression_predict(mock_session, bqml_model, mock_X): model.predict(mock_X) mock_session.read_gbq.assert_called_once_with( - "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_name`,\n (input_X_sql))", + "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_id`,\n (input_X_sql))", index_col=["index_column_id"], ) @@ -114,27 +120,33 @@ def test_linear_regression_score(mock_session, bqml_model, mock_X, mock_y): model.score(mock_X, mock_y) mock_session.read_gbq.assert_called_once_with( - "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_name`,\n (input_X_y_sql))" + "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_id`,\n (input_X_y_sql))" ) -def test_logistic_regression_default_fit(ml_mocker, mock_session, mock_X, mock_y): +def test_logistic_regression_default_fit( + bqml_model_factory, mock_session, mock_X, mock_y +): model = linear_model.LogisticRegression() + model._bqml_model_factory = bqml_model_factory model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=True,\n auto_class_weights=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + 'CREATE TEMP MODEL `temp_model_id`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=True,\n auto_class_weights=False,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' ) -def test_logistic_regression_params_fit(ml_mocker, mock_session, mock_X, mock_y): +def test_logistic_regression_params_fit( + bqml_model_factory, mock_session, mock_X, mock_y +): model = linear_model.LogisticRegression( fit_intercept=False, class_weights="balanced" ) + model._bqml_model_factory = bqml_model_factory model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - 'CREATE TEMP MODEL `temp_model_name`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=False,\n auto_class_weights=True,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' + 'CREATE TEMP MODEL `temp_model_id`\nOPTIONS(\n model_type="LOGISTIC_REG",\n data_split_method="NO_SPLIT",\n fit_intercept=False,\n auto_class_weights=True,\n INPUT_LABEL_COLS=["input_column_label"])\nAS input_X_y_sql' ) @@ -144,7 +156,7 @@ def test_logistic_regression_predict(mock_session, bqml_model, mock_X): model.predict(mock_X) mock_session.read_gbq.assert_called_once_with( - "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_name`,\n (input_X_sql))", + "SELECT * FROM ML.PREDICT(MODEL `model_project.model_dataset.model_id`,\n (input_X_sql))", index_col=["index_column_id"], ) @@ -155,5 +167,5 @@ def test_logistic_regression_score(mock_session, bqml_model, mock_X, mock_y): model.score(mock_X, mock_y) mock_session.read_gbq.assert_called_once_with( - "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_name`,\n (input_X_y_sql))" + "SELECT * FROM ML.EVALUATE(MODEL `model_project.model_dataset.model_id`,\n (input_X_y_sql))" ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index c20a17f7d6e..495e8759e8a 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -12,11 +12,32 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest + import bigframes.ml.sql as ml_sql -def test_options_produces_correct_sql(): - sql = ml_sql.options(model_type="lin_reg", input_label_cols=["col_a"], l1_reg=0.6) +@pytest.fixture(scope="session") +def base_sql_generator() -> ml_sql.BaseSqlGenerator: + return ml_sql.BaseSqlGenerator() + + +@pytest.fixture(scope="session") +def model_creation_sql_generator() -> ml_sql.ModelCreationSqlGenerator: + return ml_sql.ModelCreationSqlGenerator(model_id="my_model_id") + + +@pytest.fixture(scope="session") +def model_manipulation_sql_generator() -> ml_sql.ModelManipulationSqlGenerator: + return ml_sql.ModelManipulationSqlGenerator( + model_name="my_project_id.my_dataset_id.my_model_id" + ) + + +def test_options_produces_correct_sql(base_sql_generator: ml_sql.BaseSqlGenerator): + sql = base_sql_generator.options( + model_type="lin_reg", input_label_cols=["col_a"], l1_reg=0.6 + ) assert ( sql == """OPTIONS( @@ -26,8 +47,8 @@ def test_options_produces_correct_sql(): ) -def test_transform_produces_correct_sql(): - sql = ml_sql.transform( +def test_transform_produces_correct_sql(base_sql_generator: ml_sql.BaseSqlGenerator): + sql = base_sql_generator.transform( "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", "ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b", ) @@ -39,141 +60,172 @@ def test_transform_produces_correct_sql(): ) -def test_standard_scaler_produces_correct_sql(): - sql = ml_sql.ml_standard_scaler("col_a", "scaled_col_a") +def test_standard_scaler_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_standard_scaler("col_a", "scaled_col_a") assert sql == "ML.STANDARD_SCALER(col_a) OVER() AS scaled_col_a" -def test_one_hot_encoder_produces_correct_sql(): - sql = ml_sql.ml_one_hot_encoder("col_a", "none", 1000000, 0, "encoded_col_a") +def test_one_hot_encoder_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_one_hot_encoder( + "col_a", "none", 1000000, 0, "encoded_col_a" + ) assert ( sql == "ML.ONE_HOT_ENCODER(col_a, 'none', 1000000, 0) OVER() AS encoded_col_a" ) -def test_create_model_produces_correct_sql(): - sql = ml_sql.create_model( - model_name="my_dataset.my_model", +def test_create_model_produces_correct_sql( + model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, +): + sql = model_creation_sql_generator.create_model( source_sql="my_source_sql", options_sql="my_options_sql", ) assert ( sql - == """CREATE TEMP MODEL `my_dataset.my_model` + == """CREATE TEMP MODEL `my_model_id` my_options_sql AS my_source_sql""" ) -def test_create_model_transform_produces_correct_sql(): - sql = ml_sql.create_model( - model_name="my_model", +def test_create_model_transform_produces_correct_sql( + model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, +): + sql = model_creation_sql_generator.create_model( source_sql="my_source_sql", options_sql="my_options_sql", transform_sql="my_transform_sql", ) assert ( sql - == """CREATE TEMP MODEL `my_model` + == """CREATE TEMP MODEL `my_model_id` my_transform_sql my_options_sql AS my_source_sql""" ) -def test_create_remote_model_produces_correct_sql(): - sql = ml_sql.create_remote_model( - model_name="my_model", +def test_create_remote_model_produces_correct_sql( + model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, +): + sql = model_creation_sql_generator.create_remote_model( connection_name="my_project.us.my_connection", options_sql="my_options_sql", ) assert ( sql - == """CREATE TEMP MODEL `my_model` + == """CREATE TEMP MODEL `my_model_id` REMOTE WITH CONNECTION `my_project.us.my_connection` my_options_sql""" ) -def test_create_imported_model_produces_correct_sql(): - sql = ml_sql.create_imported_model( - model_name="my_model", +def test_create_imported_model_produces_correct_sql( + model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, +): + sql = model_creation_sql_generator.create_imported_model( options_sql="my_options_sql", ) assert ( sql - == """CREATE TEMP MODEL `my_model` + == """CREATE TEMP MODEL `my_model_id` my_options_sql""" ) -def test_alter_model_correct_sql(): - sql = ml_sql.alter_model( - model_name="my_dataset.my_model", +def test_alter_model_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.alter_model( options_sql="my_options_sql", ) assert ( sql - == """ALTER MODEL `my_dataset.my_model` + == """ALTER MODEL `my_project_id.my_dataset_id.my_model_id` SET my_options_sql""" ) -def test_ml_predict_produces_correct_sql(): - sql = ml_sql.ml_predict( - model_name="my_dataset.my_model", source_sql="SELECT * FROM my_table" +def test_ml_predict_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_predict( + source_sql="SELECT * FROM my_table" ) assert ( sql - == """SELECT * FROM ML.PREDICT(MODEL `my_dataset.my_model`, + == """SELECT * FROM ML.PREDICT(MODEL `my_project_id.my_dataset_id.my_model_id`, (SELECT * FROM my_table))""" ) -def test_ml_evaluate_produces_correct_sql(): - sql = ml_sql.ml_evaluate( - model_name="my_dataset.my_model", source_sql="SELECT * FROM my_table" +def test_ml_evaluate_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_evaluate( + source_sql="SELECT * FROM my_table" ) assert ( sql - == """SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`, + == """SELECT * FROM ML.EVALUATE(MODEL `my_project_id.my_dataset_id.my_model_id`, (SELECT * FROM my_table))""" ) -def test_ml_evaluate_no_source_produces_correct_sql(): - sql = ml_sql.ml_evaluate(model_name="my_dataset.my_model") - assert sql == """SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`)""" +def test_ml_evaluate_no_source_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_evaluate() + assert ( + sql + == """SELECT * FROM ML.EVALUATE(MODEL `my_project_id.my_dataset_id.my_model_id`)""" + ) -def test_ml_centroids_produces_correct_sql(): - sql = ml_sql.ml_centroids(model_name="my_dataset.my_model") - assert sql == """SELECT * FROM ML.CENTROIDS(MODEL `my_dataset.my_model`)""" +def test_ml_centroids_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_centroids() + assert ( + sql + == """SELECT * FROM ML.CENTROIDS(MODEL `my_project_id.my_dataset_id.my_model_id`)""" + ) -def test_ml_generate_text_produces_correct_sql(): - sql = ml_sql.ml_generate_text( - model_name="my_dataset.my_model", +def test_ml_generate_text_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_generate_text( source_sql="SELECT * FROM my_table", struct_options="STRUCT(value AS item)", ) assert ( sql - == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_dataset.my_model`, + == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project_id.my_dataset_id.my_model_id`, (SELECT * FROM my_table), STRUCT(value AS item))""" ) -def test_ml_principal_components_produces_correct_sql(): - sql = ml_sql.ml_principal_components(model_name="my_dataset.my_model") +def test_ml_principal_components_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_principal_components() assert ( - sql == """SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `my_dataset.my_model`)""" + sql + == """SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL `my_project_id.my_dataset_id.my_model_id`)""" ) -def test_ml_principal_component_info_produces_correct_sql(): - sql = ml_sql.ml_principal_component_info(model_name="my_dataset.my_model") +def test_ml_principal_component_info_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, +): + sql = model_manipulation_sql_generator.ml_principal_component_info() assert ( sql - == """SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `my_dataset.my_model`)""" + == """SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL `my_project_id.my_dataset_id.my_model_id`)""" ) From 25104610e5ffe526315923946533a66713c1d155 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 22 Sep 2023 11:50:47 -0700 Subject: [PATCH 31/80] feat: add ml.preprocessing.LabelEncoder (#50) --- bigframes/ml/compose.py | 1 + bigframes/ml/pipeline.py | 15 +- bigframes/ml/preprocessing.py | 119 ++++++++++++++ bigframes/ml/sql.py | 11 ++ tests/system/large/ml/test_pipeline.py | 61 ++++++- tests/system/small/ml/test_preprocessing.py | 149 ++++++++++++++++++ tests/unit/ml/test_compose.py | 3 + tests/unit/ml/test_sql.py | 11 +- .../sklearn/preprocessing/_label.py | 52 ++++++ 9 files changed, 415 insertions(+), 7 deletions(-) create mode 100644 third_party/bigframes_vendored/sklearn/preprocessing/_label.py diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index 02365f261c8..db5d8cf260d 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -29,6 +29,7 @@ CompilablePreprocessorType = Union[ preprocessing.OneHotEncoder, preprocessing.StandardScaler, + preprocessing.LabelEncoder, ] diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index bff0bf36ad4..71c21d565a1 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -50,6 +50,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): compose.ColumnTransformer, preprocessing.StandardScaler, preprocessing.OneHotEncoder, + preprocessing.LabelEncoder, ), ): self._transform = transform @@ -143,7 +144,11 @@ def _extract_as_column_transformer( transformers: List[ Tuple[ str, - Union[preprocessing.OneHotEncoder, preprocessing.StandardScaler], + Union[ + preprocessing.OneHotEncoder, + preprocessing.StandardScaler, + preprocessing.LabelEncoder, + ], Union[str, List[str]], ] ] = [] @@ -167,6 +172,13 @@ def _extract_as_column_transformer( *preprocessing.OneHotEncoder._parse_from_sql(transform_sql), ) ) + elif transform_sql.startswith("ML.LABEL_ENCODER"): + transformers.append( + ( + "label_encoder", + *preprocessing.LabelEncoder._parse_from_sql(transform_sql), + ) + ) else: raise NotImplementedError( f"Unsupported transformer type. {constants.FEEDBACK_LINK}" @@ -181,6 +193,7 @@ def _merge_column_transformer( compose.ColumnTransformer, preprocessing.StandardScaler, preprocessing.OneHotEncoder, + preprocessing.LabelEncoder, ]: """Try to merge the column transformer to a simple transformer.""" transformers = column_transformer.transformers_ diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index cd4ae27b8c8..6ee17751df0 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -24,6 +24,7 @@ import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.preprocessing._data import third_party.bigframes_vendored.sklearn.preprocessing._encoder +import third_party.bigframes_vendored.sklearn.preprocessing._label class StandardScaler( @@ -229,3 +230,121 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: bpd.DataFrame, df[self._output_names], ) + + +class LabelEncoder( + base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._label.LabelEncoder, +): + # BQML max value https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder#syntax + TOP_K_DEFAULT = 1000000 + FREQUENCY_THRESHOLD_DEFAULT = 0 + + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._label.LabelEncoder.__doc__ + ) + + # All estimators must implement __init__ to document their parameters, even + # if they don't have any + def __init__( + self, + min_frequency: Optional[int] = None, + max_categories: Optional[int] = None, + ): + if max_categories is not None and max_categories < 2: + raise ValueError( + f"max_categories has to be larger than or equal to 2, input is {max_categories}." + ) + self.min_frequency = min_frequency + self.max_categories = max_categories + self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return ( + type(other) is LabelEncoder + and self._bqml_model == other._bqml_model + and self.min_frequency == other.min_frequency + and self.max_categories == other.max_categories + ) + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: + a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + + # minus one here since BQML's inplimentation always includes index 0, and top_k is on top of that. + top_k = ( + (self.max_categories - 1) + if self.max_categories is not None + else LabelEncoder.TOP_K_DEFAULT + ) + frequency_threshold = ( + self.min_frequency + if self.min_frequency is not None + else LabelEncoder.FREQUENCY_THRESHOLD_DEFAULT + ) + return [ + ( + self._base_sql_generator.ml_label_encoder( + column, top_k, frequency_threshold, f"labelencoded_{column}" + ), + f"labelencoded_{column}", + ) + for column in columns + ] + + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[LabelEncoder, str]: + """Parse SQL to tuple(LabelEncoder, column_label). + + Args: + sql: SQL string of format "ML.LabelEncoder({col_label}, {top_k}, {frequency_threshold}) OVER() " + + Returns: + tuple(LabelEncoder, column_label)""" + s = sql[sql.find("(") + 1 : sql.find(")")] + col_label, top_k, frequency_threshold = s.split(", ") + max_categories = int(top_k) + 1 + min_frequency = int(frequency_threshold) + + return cls(min_frequency, max_categories), col_label + + def fit( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> LabelEncoder: + (X,) = utils.convert_to_dataframe(X) + + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = self._bqml_model_factory.create_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + return self + + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.transform(X) + return typing.cast( + bpd.DataFrame, + df[self._output_names], + ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 7cf030485bd..5d35a10b966 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -88,6 +88,17 @@ def ml_one_hot_encoder( https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder for params.""" return f"""ML.ONE_HOT_ENCODER({numeric_expr_sql}, '{drop}', {top_k}, {frequency_threshold}) OVER() AS {name}""" + def ml_label_encoder( + self, + numeric_expr_sql: str, + top_k: int, + frequency_threshold: int, + name: str, + ) -> str: + """Encode ML.LABEL_ENCODER for BQML. + https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-label-encoder for params.""" + return f"""ML.LABEL_ENCODER({numeric_expr_sql}, {top_k}, {frequency_threshold}) OVER() AS {name}""" + class ModelCreationSqlGenerator(BaseSqlGenerator): """Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id.""" diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index bec1a51a991..87664b4c3d4 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -570,6 +570,11 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "label", + preprocessing.LabelEncoder(), + "species", + ), ] ), ), @@ -632,6 +637,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "label", + preprocessing.LabelEncoder(), + "species", + ), ] ), ), @@ -650,7 +660,7 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id assert isinstance(pl_loaded._transform, compose.ColumnTransformer) transformers = pl_loaded._transform.transformers_ - assert len(transformers) == 3 + assert len(transformers) == 4 assert transformers[0][0] == "ont_hot_encoder" assert isinstance(transformers[0][1], preprocessing.OneHotEncoder) @@ -660,13 +670,20 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id assert one_hot_encoder.max_categories == 100 assert transformers[0][2] == "species" - assert transformers[1][0] == "standard_scaler" - assert isinstance(transformers[1][1], preprocessing.StandardScaler) - assert transformers[1][2] == "culmen_length_mm" + assert transformers[1][0] == "label_encoder" + assert isinstance(transformers[1][1], preprocessing.LabelEncoder) + one_hot_encoder = transformers[1][1] + assert one_hot_encoder.min_frequency == 0 + assert one_hot_encoder.max_categories == 1000001 + assert transformers[1][2] == "species" assert transformers[2][0] == "standard_scaler" assert isinstance(transformers[2][1], preprocessing.StandardScaler) - assert transformers[2][2] == "flipper_length_mm" + assert transformers[2][2] == "culmen_length_mm" + + assert transformers[3][0] == "standard_scaler" + assert isinstance(transformers[2][1], preprocessing.StandardScaler) + assert transformers[3][2] == "flipper_length_mm" assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) assert pl_loaded._estimator.fit_intercept is False @@ -735,3 +752,37 @@ def test_pipeline_one_hot_encoder_to_gbq(penguins_df_default_index, dataset_id): assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) assert pl_loaded._estimator.fit_intercept is False + + +def test_pipeline_label_encoder_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ( + "transform", + preprocessing.LabelEncoder(min_frequency=5, max_categories=100), + ), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "sex", + "species", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_label_encoder", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.LabelEncoder) + + label_encoder = pl_loaded._transform + assert label_encoder.min_frequency == 5 + assert label_encoder.max_categories == 100 + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index 57b9900c489..1f08ef2c2c9 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -264,4 +264,153 @@ def test_one_hot_encoder_different_data(penguins_df_default_index, new_penguins_ pd.testing.assert_frame_equal(result, expected) +def test_label_encoder_default_params(new_penguins_df): + encoder = bigframes.ml.preprocessing.LabelEncoder() + encoder.fit(new_penguins_df[["species", "sex"]]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "labelencoded_sex": [ + 2, + 1, + 1, + ], + "labelencoded_species": [ + 1, + 1, + 2, + ], + }, + dtype="Int64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_label_encoder_default_params_fit_transform(new_penguins_df): + encoder = bigframes.ml.preprocessing.LabelEncoder() + + result = encoder.fit_transform(new_penguins_df[["species", "sex"]]).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "labelencoded_sex": [ + 2, + 1, + 1, + ], + "labelencoded_species": [ + 1, + 1, + 2, + ], + }, + dtype="Int64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_label_encoder_series_default_params(new_penguins_df): + encoder = bigframes.ml.preprocessing.LabelEncoder() + encoder.fit(new_penguins_df["species"]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "labelencoded_species": [ + 1, + 1, + 2, + ], + }, + dtype="Int64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_label_encoder_params(new_penguins_df): + encoder = bigframes.ml.preprocessing.LabelEncoder(100, 2) + encoder.fit(new_penguins_df[["species", "sex"]]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "labelencoded_sex": [ + 0, + 0, + 0, + ], + "labelencoded_species": [ + 0, + 0, + 0, + ], + }, + dtype="Int64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_label_encoder_different_data(penguins_df_default_index, new_penguins_df): + encoder = bigframes.ml.preprocessing.LabelEncoder() + encoder.fit(penguins_df_default_index[["species", "sex"]]) + + result = encoder.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "labelencoded_sex": [ + 3, + 2, + 2, + ], + "labelencoded_species": [ + 1, + 1, + 2, + ], + }, + dtype="Int64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + # TODO(garrettwu): add OneHotEncoder tests to compare with sklearn. diff --git a/tests/unit/ml/test_compose.py b/tests/unit/ml/test_compose.py index 772a148c95f..c5b3b508762 100644 --- a/tests/unit/ml/test_compose.py +++ b/tests/unit/ml/test_compose.py @@ -23,10 +23,12 @@ def test_columntransformer_init_expectedtransforms(): onehot_transformer = bigframes.ml.preprocessing.OneHotEncoder() scaler_transformer = bigframes.ml.preprocessing.StandardScaler() + label_transformer = bigframes.ml.preprocessing.LabelEncoder() column_transformer = bigframes.ml.compose.ColumnTransformer( [ ("onehot", onehot_transformer, "species"), ("scale", scaler_transformer, ["culmen_length_mm", "flipper_length_mm"]), + ("onehot", label_transformer, "species"), ] ) @@ -34,6 +36,7 @@ def test_columntransformer_init_expectedtransforms(): ("onehot", onehot_transformer, "species"), ("scale", scaler_transformer, "culmen_length_mm"), ("scale", scaler_transformer, "flipper_length_mm"), + ("onehot", label_transformer, "species"), ] diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 495e8759e8a..23b68aa150a 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -51,12 +51,14 @@ def test_transform_produces_correct_sql(base_sql_generator: ml_sql.BaseSqlGenera sql = base_sql_generator.transform( "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", "ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b", + "ML.LABEL_ENCODER(col_c) OVER(col_c) AS encoded_col_c", ) assert ( sql == """TRANSFORM( ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a, - ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b)""" + ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b, + ML.LABEL_ENCODER(col_c) OVER(col_c) AS encoded_col_c)""" ) @@ -78,6 +80,13 @@ def test_one_hot_encoder_produces_correct_sql( ) +def test_label_encoder_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_label_encoder("col_a", 1000000, 0, "encoded_col_a") + assert sql == "ML.LABEL_ENCODER(col_a, 1000000, 0) OVER() AS encoded_col_a" + + def test_create_model_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, ): diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_label.py b/third_party/bigframes_vendored/sklearn/preprocessing/_label.py new file mode 100644 index 00000000000..7e60c846d45 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_label.py @@ -0,0 +1,52 @@ +# Authors: Alexandre Gramfort +# Mathieu Blondel +# Olivier Grisel +# Andreas Mueller +# Joel Nothman +# Hamzeh Alsalhi +# License: BSD 3 clause + +from bigframes import constants +from third_party.bigframes_vendored.sklearn.base import BaseEstimator + + +class LabelEncoder(BaseEstimator): + """Encode target labels with value between 0 and n_classes-1. + + This transformer should be used to encode target values, *i.e.* `y`, and + not the input `X`. + + Args: + min_frequency (Optional[int], default None): + Specifies the minimum frequency below which a category will be considered infrequent. + Default None. + int: categories with a smaller cardinality will be considered infrequent as ßindex 0. + max_categories (Optional[int], default None): + Specifies an upper limit to the number of output features for each input feature + when considering infrequent categories. If there are infrequent categories, + max_categories includes the category representing the infrequent categories along with the frequent categories. + Default None, set limit to 1,000,000. + """ + + def fit(self, X): + """Fit LabelEncoder to X. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series with training data. + + Returns: + LabelEncoder: Fitted encoder. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transform(self, X): + """Transform X using label encoding. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. + + Returns: + bigframes.dataframe.DataFrame: The result is an array-like of values.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From f9a93ce71d053aa17b1e3a2946c90e0227076184 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 22 Sep 2023 14:00:16 -0700 Subject: [PATCH 32/80] perf: reimplement unpivot to use cross join rather than union (#47) --- bigframes/core/__init__.py | 97 ++++++++++++++++++++++---------------- 1 file changed, 57 insertions(+), 40 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 3b3754642e7..ce3b1aa630b 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -982,61 +982,78 @@ def unpivot( ArrayValue: The unpivoted ArrayValue """ table = self._to_ibis_expr(ordering_mode="offset_col") - sub_expressions = [] - - # Use ibis memtable to infer type of rowlabels (if possible) - # TODO: Allow caller to specify dtype - labels_ibis_type = ibis.memtable({"col": row_labels})["col"].type() - labels_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(labels_ibis_type) - row_n = len(row_labels) if not all( len(source_columns) == row_n for _, source_columns in unpivot_columns ): raise ValueError("Columns and row labels must all be same length.") - for i in range(row_n): - values = [] - for j in range(len(unpivot_columns)): - result_col, source_cols = unpivot_columns[j] - col_dtype = dtype[j] if utils.is_list_like(dtype) else dtype - if source_cols[i] is not None: - values.append( - ops.AsTypeOp(col_dtype) - ._as_ibis(table[source_cols[i]]) - .name(result_col) - ) - else: - values.append( - bigframes.dtypes.literal_to_ibis_scalar( - None, force_dtype=col_dtype - ).name(result_col) - ) - offsets_value = ( - ((table[ORDER_ID_COLUMN] * row_n) + i) - .cast(ibis_dtypes.int64) - .name(ORDER_ID_COLUMN), + unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") + unpivot_table = table.cross_join( + ibis.memtable({unpivot_offset_id: range(row_n)}) + ) + unpivot_offsets_value = ( + ( + (unpivot_table[ORDER_ID_COLUMN] * row_n) + + unpivot_table[unpivot_offset_id] ) - sub_expr = table.select( - passthrough_columns, + .cast(ibis_dtypes.int64) + .name(ORDER_ID_COLUMN), + ) + + # Use ibis memtable to infer type of rowlabels (if possible) + # TODO: Allow caller to specify dtype + labels_ibis_type = ibis.memtable({"col": row_labels})["col"].type() + labels_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(labels_ibis_type) + cases = [ + ( + i, bigframes.dtypes.literal_to_ibis_scalar( row_labels[i], force_dtype=labels_dtype # type:ignore - ).name(index_col_id), - *values, - offsets_value, + ), + ) + for i in range(len(row_labels)) + ] + labels_value = ( + typing.cast(ibis_types.IntegerColumn, unpivot_table[unpivot_offset_id]) + .cases(cases, default=None) # type:ignore + .name(index_col_id) + ) + + unpivot_values = [] + for j in range(len(unpivot_columns)): + col_dtype = dtype[j] if utils.is_list_like(dtype) else dtype + result_col, source_cols = unpivot_columns[j] + null_value = bigframes.dtypes.literal_to_ibis_scalar( + None, force_dtype=col_dtype + ) + ibis_values = [ + ops.AsTypeOp(col_dtype)._as_ibis(unpivot_table[col]) + if col is not None + else null_value + for col in source_cols + ] + cases = [(i, ibis_values[i]) for i in range(len(ibis_values))] + unpivot_value = typing.cast( + ibis_types.IntegerColumn, unpivot_table[unpivot_offset_id] + ).cases( + cases, default=null_value # type:ignore ) - sub_expressions.append(sub_expr) - rotated_table = ibis.union(*sub_expressions) + unpivot_values.append(unpivot_value.name(result_col)) + + unpivot_table = unpivot_table.select( + passthrough_columns, labels_value, *unpivot_values, unpivot_offsets_value + ) value_columns = [ - rotated_table[value_col_id] for value_col_id, _ in unpivot_columns + unpivot_table[value_col_id] for value_col_id, _ in unpivot_columns ] - passthrough_values = [rotated_table[col] for col in passthrough_columns] + passthrough_values = [unpivot_table[col] for col in passthrough_columns] return ArrayValue( session=self._session, - table=rotated_table, - columns=[rotated_table[index_col_id], *value_columns, *passthrough_values], - hidden_ordering_columns=[rotated_table[ORDER_ID_COLUMN]], + table=unpivot_table, + columns=[unpivot_table[index_col_id], *value_columns, *passthrough_values], + hidden_ordering_columns=[unpivot_table[ORDER_ID_COLUMN]], ordering=ExpressionOrdering( ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), From 416d7cb9b560d7e33dcc0227f03a00d43f55ba0d Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 22 Sep 2023 17:48:02 -0700 Subject: [PATCH 33/80] perf: add ability to cache dataframe and series to session table (#51) --- bigframes/core/__init__.py | 23 +++++++++++++++++++++++ bigframes/core/blocks.py | 9 +++++++++ bigframes/dataframe.py | 3 +++ bigframes/series.py | 3 +++ tests/system/small/test_dataframe.py | 10 ++++++++++ 5 files changed, 48 insertions(+) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index ce3b1aa630b..c529f833510 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -1145,6 +1145,29 @@ def slice( ) return sliced_expr if step > 0 else sliced_expr.reversed() + def cached(self, cluster_cols: typing.Sequence[str]) -> ArrayValue: + """Write the ArrayValue to a session table and create a new block object that references it.""" + ibis_expr = self._to_ibis_expr( + ordering_mode="unordered", expose_hidden_cols=True + ) + destination = self._session._ibis_to_session_table( + ibis_expr, cluster_cols=cluster_cols, api_name="cache" + ) + table_expression = self._session.ibis_client.sql( + f"SELECT * FROM `_SESSION`.`{destination.table_id}`" + ) + new_columns = [table_expression[column] for column in self.column_names] + new_hidden_columns = [ + table_expression[column] for column in self._hidden_ordering_column_names + ] + return ArrayValue( + self._session, + table_expression, + columns=new_columns, + hidden_ordering_columns=new_hidden_columns, + ordering=self._ordering, + ) + class ArrayValueBuilder: """Mutable expression class. diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index ad4f72070f5..c4127c5fd52 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -1467,6 +1467,15 @@ def to_sql_query( idx_labels, ) + def cached(self) -> Block: + """Write the block to a session table and create a new block object that references it.""" + return Block( + self.expr.cached(cluster_cols=self.index_columns), + index_columns=self.index_columns, + column_labels=self.column_labels, + index_labels=self.index_labels, + ) + def _is_monotonic( self, column_ids: typing.Union[str, Sequence[str]], increasing: bool ) -> bool: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 76377cd477f..0d357e7c3d3 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2480,3 +2480,6 @@ def _set_block(self, block: blocks.Block): def _get_block(self) -> blocks.Block: return self._block + + def _cached(self) -> DataFrame: + return DataFrame(self._block.cached()) diff --git a/bigframes/series.py b/bigframes/series.py index 8e47088c146..c1c0cb05372 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1389,6 +1389,9 @@ def _slice( ), ) + def _cached(self) -> Series: + return Series(self._block.cached()) + def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: return pandas.api.types.is_list_like(obj) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 3eeb368ad23..b6ca958c038 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2717,3 +2717,13 @@ def test_query_job_setters(scalars_df_default_index: dataframe.DataFrame): job_ids.add(scalars_df_default_index.query_job.job_id) assert len(job_ids) == 2 + + +def test_df_cached(scalars_df_index): + df = scalars_df_index.set_index(["int64_too", "int64_col"]).sort_values( + "string_col" + ) + df = df[df["rowindex_2"] % 2 == 0] + + df_cached_copy = df._cached() + pandas.testing.assert_frame_equal(df.to_pandas(), df_cached_copy.to_pandas()) From 14b262bde2bb86093bf4df63862e369c5a84b0ad Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Mon, 25 Sep 2023 14:33:40 -0700 Subject: [PATCH 34/80] feat: add ml.preprocessing.MaxAbsScaler (#56) --- bigframes/clients.py | 3 +- bigframes/ml/compose.py | 1 + bigframes/ml/pipeline.py | 10 ++ bigframes/ml/preprocessing.py | 86 +++++++++++++- bigframes/ml/sql.py | 4 + tests/system/large/ml/test_compose.py | 24 ++-- tests/system/large/ml/test_pipeline.py | 82 ++++++++++---- tests/system/small/ml/test_preprocessing.py | 105 ++++++++++++++++-- tests/unit/ml/test_compose.py | 72 ++++++++---- tests/unit/ml/test_sql.py | 7 ++ .../sklearn/preprocessing/_data.py | 42 ++++++- .../sklearn/preprocessing/_encoder.py | 5 +- 12 files changed, 370 insertions(+), 71 deletions(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index 5c019e0fc89..b60fcba04af 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -18,7 +18,7 @@ import logging import time -from typing import Optional +from typing import cast, Optional import google.api_core.exceptions from google.cloud import bigquery_connection_v1, resourcemanager_v3 @@ -80,6 +80,7 @@ def create_bq_connection( logger.info( f"Created BQ connection {connection_name} with service account id: {service_account_id}" ) + service_account_id = cast(str, service_account_id) # Ensure IAM role on the BQ connection # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function self._ensure_iam_binding(project_id, service_account_id, iam_role) diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index db5d8cf260d..a1075c23981 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -29,6 +29,7 @@ CompilablePreprocessorType = Union[ preprocessing.OneHotEncoder, preprocessing.StandardScaler, + preprocessing.MaxAbsScaler, preprocessing.LabelEncoder, ] diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index 71c21d565a1..86b20996194 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -50,6 +50,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): compose.ColumnTransformer, preprocessing.StandardScaler, preprocessing.OneHotEncoder, + preprocessing.MaxAbsScaler, preprocessing.LabelEncoder, ), ): @@ -147,6 +148,7 @@ def _extract_as_column_transformer( Union[ preprocessing.OneHotEncoder, preprocessing.StandardScaler, + preprocessing.MaxAbsScaler, preprocessing.LabelEncoder, ], Union[str, List[str]], @@ -172,6 +174,13 @@ def _extract_as_column_transformer( *preprocessing.OneHotEncoder._parse_from_sql(transform_sql), ) ) + elif transform_sql.startswith("ML.MAX_ABS_SCALER"): + transformers.append( + ( + "max_abs_encoder", + *preprocessing.MaxAbsScaler._parse_from_sql(transform_sql), + ) + ) elif transform_sql.startswith("ML.LABEL_ENCODER"): transformers.append( ( @@ -193,6 +202,7 @@ def _merge_column_transformer( compose.ColumnTransformer, preprocessing.StandardScaler, preprocessing.OneHotEncoder, + preprocessing.MaxAbsScaler, preprocessing.LabelEncoder, ]: """Try to merge the column transformer to a simple transformer.""" diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 6ee17751df0..f4f5446651c 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -54,8 +54,10 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: Returns: a list of tuples of (sql_expression, output_name)""" return [ ( - self._base_sql_generator.ml_standard_scaler(column, f"scaled_{column}"), - f"scaled_{column}", + self._base_sql_generator.ml_standard_scaler( + column, f"standard_scaled_{column}" + ), + f"standard_scaled_{column}", ) for column in columns ] @@ -105,6 +107,86 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +class MaxAbsScaler( + base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._data.MaxAbsScaler, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._data.MaxAbsScaler.__doc__ + ) + + def __init__(self): + self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return type(other) is MaxAbsScaler and self._bqml_model == other._bqml_model + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + ( + self._base_sql_generator.ml_max_abs_scaler( + column, f"max_abs_scaled_{column}" + ), + f"max_abs_scaled_{column}", + ) + for column in columns + ] + + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[MaxAbsScaler, str]: + """Parse SQL to tuple(StandardScaler, column_label). + + Args: + sql: SQL string of format "ML.MAX_ABS_SCALER({col_label}) OVER()" + + Returns: + tuple(StandardScaler, column_label)""" + col_label = sql[sql.find("(") + 1 : sql.find(")")] + return cls(), col_label + + def fit( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> MaxAbsScaler: + (X,) = utils.convert_to_dataframe(X) + + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = self._bqml_model_factory.create_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + return self + + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.transform(X) + return typing.cast( + bpd.DataFrame, + df[self._output_names], + ) + + class OneHotEncoder( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 5d35a10b966..a54d39e6b2f 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -76,6 +76,10 @@ def ml_standard_scaler(self, numeric_expr_sql: str, name: str) -> str: """Encode ML.STANDARD_SCALER for BQML""" return f"""ML.STANDARD_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_max_abs_scaler(self, numeric_expr_sql: str, name: str) -> str: + """Encode ML.MAX_ABS_SCALER for BQML""" + return f"""ML.MAX_ABS_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_one_hot_encoder( self, numeric_expr_sql: str, diff --git a/tests/system/large/ml/test_compose.py b/tests/system/large/ml/test_compose.py index b65baa63ebd..0c280e5d020 100644 --- a/tests/system/large/ml/test_compose.py +++ b/tests/system/large/ml/test_compose.py @@ -56,20 +56,20 @@ def test_columntransformer_standalone_fit_and_transform( [{"index": 1, "value": 1.0}], [{"index": 2, "value": 1.0}], ], - "scaled_culmen_length_mm": [ + "standard_scaled_culmen_length_mm": [ -0.811119671289163, -0.9945520581113803, -1.104611490204711, ], - "scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], + "standard_scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], }, index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"), ) - expected.scaled_culmen_length_mm = expected.scaled_culmen_length_mm.astype( - "Float64" + expected.standard_scaled_culmen_length_mm = ( + expected.standard_scaled_culmen_length_mm.astype("Float64") ) - expected.scaled_flipper_length_mm = expected.scaled_flipper_length_mm.astype( - "Float64" + expected.standard_scaled_flipper_length_mm = ( + expected.standard_scaled_flipper_length_mm.astype("Float64") ) pandas.testing.assert_frame_equal(result, expected, rtol=1e-3) @@ -107,20 +107,20 @@ def test_columntransformer_standalone_fit_transform(new_penguins_df): [{"index": 1, "value": 1.0}], [{"index": 2, "value": 1.0}], ], - "scaled_culmen_length_mm": [ + "standard_scaled_culmen_length_mm": [ 1.313249, -0.20198, -1.111118, ], - "scaled_flipper_length_mm": [1.251098, -1.196588, -0.054338], + "standard_scaled_flipper_length_mm": [1.251098, -1.196588, -0.054338], }, index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"), ) - expected.scaled_culmen_length_mm = expected.scaled_culmen_length_mm.astype( - "Float64" + expected.standard_scaled_culmen_length_mm = ( + expected.standard_scaled_culmen_length_mm.astype("Float64") ) - expected.scaled_flipper_length_mm = expected.scaled_flipper_length_mm.astype( - "Float64" + expected.standard_scaled_flipper_length_mm = ( + expected.standard_scaled_flipper_length_mm.astype("Float64") ) pandas.testing.assert_frame_equal(result, expected, rtol=1e-3) diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 87664b4c3d4..c69a00b81ca 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -566,10 +566,15 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind "species", ), ( - "scale", + "standard_scale", preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "max_abs_scale", + preprocessing.MaxAbsScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -637,6 +642,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "max_abs_scale", + preprocessing.MaxAbsScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -660,30 +670,26 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id assert isinstance(pl_loaded._transform, compose.ColumnTransformer) transformers = pl_loaded._transform.transformers_ - assert len(transformers) == 4 - - assert transformers[0][0] == "ont_hot_encoder" - assert isinstance(transformers[0][1], preprocessing.OneHotEncoder) - one_hot_encoder = transformers[0][1] - assert one_hot_encoder.drop == "most_frequent" - assert one_hot_encoder.min_frequency == 5 - assert one_hot_encoder.max_categories == 100 - assert transformers[0][2] == "species" - - assert transformers[1][0] == "label_encoder" - assert isinstance(transformers[1][1], preprocessing.LabelEncoder) - one_hot_encoder = transformers[1][1] - assert one_hot_encoder.min_frequency == 0 - assert one_hot_encoder.max_categories == 1000001 - assert transformers[1][2] == "species" - - assert transformers[2][0] == "standard_scaler" - assert isinstance(transformers[2][1], preprocessing.StandardScaler) - assert transformers[2][2] == "culmen_length_mm" + expected = [ + ( + "ont_hot_encoder", + preprocessing.OneHotEncoder( + drop="most_frequent", max_categories=100, min_frequency=5 + ), + "species", + ), + ( + "label_encoder", + preprocessing.LabelEncoder(max_categories=1000001, min_frequency=0), + "species", + ), + ("standard_scaler", preprocessing.StandardScaler(), "culmen_length_mm"), + ("max_abs_encoder", preprocessing.MaxAbsScaler(), "culmen_length_mm"), + ("standard_scaler", preprocessing.StandardScaler(), "flipper_length_mm"), + ("max_abs_encoder", preprocessing.MaxAbsScaler(), "flipper_length_mm"), + ] - assert transformers[3][0] == "standard_scaler" - assert isinstance(transformers[2][1], preprocessing.StandardScaler) - assert transformers[3][2] == "flipper_length_mm" + assert transformers == expected assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) assert pl_loaded._estimator.fit_intercept is False @@ -717,6 +723,34 @@ def test_pipeline_standard_scaler_to_gbq(penguins_df_default_index, dataset_id): assert pl_loaded._estimator.fit_intercept is False +def test_pipeline_max_abs_scaler_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ("transform", preprocessing.MaxAbsScaler()), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_standard_scaler", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.MaxAbsScaler) + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False + + def test_pipeline_one_hot_encoder_to_gbq(penguins_df_default_index, dataset_id): pl = pipeline.Pipeline( [ diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index 1f08ef2c2c9..7779eb8f6ea 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -47,9 +47,9 @@ def test_standard_scaler_normalizes(penguins_df_default_index, new_penguins_df): expected = pd.DataFrame( { - "scaled_culmen_depth_mm": [0.836148, 0.024748, 0.48116], - "scaled_culmen_length_mm": [-0.81112, -0.994552, -1.104611], - "scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], + "standard_scaled_culmen_depth_mm": [0.836148, 0.024748, 0.48116], + "standard_scaled_culmen_length_mm": [-0.81112, -0.994552, -1.104611], + "standard_scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198], }, dtype="Float64", index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), @@ -76,9 +76,9 @@ def test_standard_scaler_normalizeds_fit_transform(new_penguins_df): expected = pd.DataFrame( { - "scaled_culmen_depth_mm": [1.17072, -1.272416, 0.101848], - "scaled_culmen_length_mm": [1.313249, -0.20198, -1.111118], - "scaled_flipper_length_mm": [1.251089, -1.196588, -0.054338], + "standard_scaled_culmen_depth_mm": [1.17072, -1.272416, 0.101848], + "standard_scaled_culmen_length_mm": [1.313249, -0.20198, -1.111118], + "standard_scaled_flipper_length_mm": [1.251089, -1.196588, -0.054338], }, dtype="Float64", index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), @@ -107,7 +107,7 @@ def test_standard_scaler_series_normalizes(penguins_df_default_index, new_pengui expected = pd.DataFrame( { - "scaled_culmen_length_mm": [ + "standard_scaled_culmen_length_mm": [ -0.811119671289163, -0.9945520581113803, -1.104611490204711, @@ -120,6 +120,97 @@ def test_standard_scaler_series_normalizes(penguins_df_default_index, new_pengui pd.testing.assert_frame_equal(result, expected, rtol=1e-3) +def test_max_abs_scaler_normalizes(penguins_df_default_index, new_penguins_df): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + scaler = bigframes.ml.preprocessing.MaxAbsScaler() + scaler.fit( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + result = scaler.transform( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ).to_pandas() + + # If maxabs-scaled correctly, max should be 1.0 + for column in result.columns: + assert math.isclose(result[column].max(), 1.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "max_abs_scaled_culmen_depth_mm": [0.874419, 0.8, 0.84186], + "max_abs_scaled_culmen_length_mm": [0.662752, 0.645973, 0.635906], + "max_abs_scaled_flipper_length_mm": [0.848485, 0.78355, 0.813853], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_max_abs_scaler_normalizeds_fit_transform(new_penguins_df): + scaler = bigframes.ml.preprocessing.MaxAbsScaler() + result = scaler.fit_transform( + new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] + ).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "max_abs_scaled_culmen_depth_mm": [1.0, 0.914894, 0.962766], + "max_abs_scaled_culmen_length_mm": [1.0, 0.974684, 0.959494], + "max_abs_scaled_flipper_length_mm": [1.0, 0.923469, 0.959184], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_max_abs_scaler_series_normalizes(penguins_df_default_index, new_penguins_df): + scaler = bigframes.ml.preprocessing.MaxAbsScaler() + scaler.fit(penguins_df_default_index["culmen_length_mm"]) + + result = scaler.transform(penguins_df_default_index["culmen_length_mm"]).to_pandas() + + # If maxabs-scaled correctly, max should be 1.0 + for column in result.columns: + assert math.isclose(result[column].max(), 1.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "max_abs_scaled_culmen_length_mm": [0.662752, 0.645973, 0.635906], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + def test_one_hot_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() encoder.fit(new_penguins_df[["species", "sex"]]) diff --git a/tests/unit/ml/test_compose.py b/tests/unit/ml/test_compose.py index c5b3b508762..24cf0a333e7 100644 --- a/tests/unit/ml/test_compose.py +++ b/tests/unit/ml/test_compose.py @@ -12,45 +12,60 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import sklearn.compose as sklearn_compose # type: ignore import sklearn.preprocessing as sklearn_preprocessing # type: ignore -import bigframes.ml.compose -import bigframes.ml.preprocessing +from bigframes.ml import compose, preprocessing def test_columntransformer_init_expectedtransforms(): - onehot_transformer = bigframes.ml.preprocessing.OneHotEncoder() - scaler_transformer = bigframes.ml.preprocessing.StandardScaler() - label_transformer = bigframes.ml.preprocessing.LabelEncoder() - column_transformer = bigframes.ml.compose.ColumnTransformer( + onehot_transformer = preprocessing.OneHotEncoder() + standard_scaler_transformer = preprocessing.StandardScaler() + max_abs_scaler_transformer = preprocessing.MaxAbsScaler() + label_transformer = preprocessing.LabelEncoder() + column_transformer = compose.ColumnTransformer( [ ("onehot", onehot_transformer, "species"), - ("scale", scaler_transformer, ["culmen_length_mm", "flipper_length_mm"]), - ("onehot", label_transformer, "species"), + ( + "standard_scale", + standard_scaler_transformer, + ["culmen_length_mm", "flipper_length_mm"], + ), + ( + "max_abs_scale", + max_abs_scaler_transformer, + ["culmen_length_mm", "flipper_length_mm"], + ), + ("label", label_transformer, "species"), ] ) assert column_transformer.transformers_ == [ ("onehot", onehot_transformer, "species"), - ("scale", scaler_transformer, "culmen_length_mm"), - ("scale", scaler_transformer, "flipper_length_mm"), - ("onehot", label_transformer, "species"), + ("standard_scale", standard_scaler_transformer, "culmen_length_mm"), + ("standard_scale", standard_scaler_transformer, "flipper_length_mm"), + ("max_abs_scale", max_abs_scaler_transformer, "culmen_length_mm"), + ("max_abs_scale", max_abs_scaler_transformer, "flipper_length_mm"), + ("label", label_transformer, "species"), ] def test_columntransformer_repr(): - column_transformer = bigframes.ml.compose.ColumnTransformer( + column_transformer = compose.ColumnTransformer( [ ( "onehot", - bigframes.ml.preprocessing.OneHotEncoder(), + preprocessing.OneHotEncoder(), "species", ), ( - "scale", - bigframes.ml.preprocessing.StandardScaler(), + "standard_scale", + preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ( + "max_abs_scale", + preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), ] @@ -59,23 +74,29 @@ def test_columntransformer_repr(): assert ( column_transformer.__repr__() == """ColumnTransformer(transformers=[('onehot', OneHotEncoder(), 'species'), - ('scale', StandardScaler(), + ('standard_scale', StandardScaler(), + ['culmen_length_mm', 'flipper_length_mm']), + ('max_abs_scale', MaxAbsScaler(), ['culmen_length_mm', 'flipper_length_mm'])])""" ) -@pytest.mark.skipif(sklearn_compose is None, reason="requires sklearn") def test_columntransformer_repr_matches_sklearn(): - bf_column_transformer = bigframes.ml.compose.ColumnTransformer( + bf_column_transformer = compose.ColumnTransformer( [ ( "onehot", - bigframes.ml.preprocessing.OneHotEncoder(), + preprocessing.OneHotEncoder(), "species", ), ( - "scale", - bigframes.ml.preprocessing.StandardScaler(), + "standard_scale", + preprocessing.StandardScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), + ( + "max_abs_scale", + preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), ] @@ -88,10 +109,15 @@ def test_columntransformer_repr_matches_sklearn(): "species", ), ( - "scale", + "standard_scale", sklearn_preprocessing.StandardScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "max_abs_scale", + sklearn_preprocessing.MaxAbsScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 23b68aa150a..c1b29c5e52d 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -69,6 +69,13 @@ def test_standard_scaler_produces_correct_sql( assert sql == "ML.STANDARD_SCALER(col_a) OVER() AS scaled_col_a" +def test_max_abs_scaler_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_max_abs_scaler("col_a", "scaled_col_a") + assert sql == "ML.MAX_ABS_SCALER(col_a) OVER() AS scaled_col_a" + + def test_one_hot_encoder_produces_correct_sql( base_sql_generator: ml_sql.BaseSqlGenerator, ): diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index 89981e34c03..40b4f76ab7f 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -43,13 +43,16 @@ class StandardScaler(BaseEstimator, TransformerMixin): print(scaler.transform(bpd.DataFrame({"a": [2], "b":[2]}))) """ - def fit(self, X): + def fit(self, X, y=None): """Compute the mean and std to be used for later scaling. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): The Dataframe or Series with training data. + y (default None): + Ignored. + Returns: StandardScaler: Fitted scaler. """ @@ -66,3 +69,40 @@ def transform(self, X): bigframes.dataframe.DataFrame: Transformed result. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + +class MaxAbsScaler(BaseEstimator, TransformerMixin): + """Scale each feature by its maximum absolute value. + + This estimator scales and translates each feature individually such + that the maximal absolute value of each feature in the + training set will be 1.0. It does not shift/center the data, and + thus does not destroy any sparsity. + """ + + def fit(self, X, y=None): + """Compute the maximum absolute value to be used for later scaling. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Dataframe or Series with training data. + + y (default None): + Ignored. + + Returns: + MaxAbsScaler: Fitted scaler. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transform(self, X): + """Scale the data. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. + + Returns: + bigframes.dataframe.DataFrame: Transformed result. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py index b0f0df8e157..cf660ece5d4 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -54,13 +54,16 @@ class OneHotEncoder(BaseEstimator): print(enc.transform(bpd.DataFrame({"a": ["Female", "Male"], "b": ["1", "4"]}))) """ - def fit(self, X): + def fit(self, X, y=None): """Fit OneHotEncoder to X. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): The DataFrame or Series with training data. + y (default None): + Ignored. + Returns: OneHotEncoder: Fitted encoder. """ From 9cf99721ed83704e6ee28b15c699326c431eb252 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Mon, 25 Sep 2023 15:46:10 -0700 Subject: [PATCH 35/80] feat: add axis param to simple df aggregations (#52) --- bigframes/core/blocks.py | 50 ++++++++++++--- bigframes/dataframe.py | 63 ++++++++++++------- tests/system/small/test_dataframe.py | 30 ++++++++- .../bigframes_vendored/pandas/core/frame.py | 42 ++++++++++--- 4 files changed, 143 insertions(+), 42 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index c4127c5fd52..b53c2212c10 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -822,22 +822,54 @@ def filter(self, column_id: str, keep_null: bool = False): index_labels=self.index.names, ) - def aggregate_all_and_pivot( + def aggregate_all_and_stack( self, operation: agg_ops.AggregateOp, *, + axis: int | str = 0, value_col_id: str = "values", dropna: bool = True, dtype=pd.Float64Dtype(), ) -> Block: - aggregations = [(col_id, operation, col_id) for col_id in self.value_columns] - result_expr = self.expr.aggregate(aggregations, dropna=dropna).unpivot( - row_labels=self.column_labels.to_list(), - index_col_id="index", - unpivot_columns=[(value_col_id, self.value_columns)], - dtype=dtype, - ) - return Block(result_expr, index_columns=["index"], column_labels=[None]) + axis_n = utils.get_axis_number(axis) + if axis_n == 0: + aggregations = [ + (col_id, operation, col_id) for col_id in self.value_columns + ] + result_expr = self.expr.aggregate(aggregations, dropna=dropna).unpivot( + row_labels=self.column_labels.to_list(), + index_col_id="index", + unpivot_columns=[(value_col_id, self.value_columns)], + dtype=dtype, + ) + return Block(result_expr, index_columns=["index"], column_labels=[None]) + else: # axis_n == 1 + # using offsets as identity to group on. + # TODO: Allow to promote identity/total_order columns instead for better perf + expr_with_offsets, offset_col = self.expr.promote_offsets() + stacked_expr = expr_with_offsets.unpivot( + row_labels=self.column_labels.to_list(), + index_col_id=guid.generate_guid(), + unpivot_columns=[(value_col_id, self.value_columns)], + passthrough_columns=[*self.index_columns, offset_col], + dtype=dtype, + ) + index_aggregations = [ + (col_id, agg_ops.AnyValueOp(), col_id) + for col_id in [*self.index_columns] + ] + main_aggregation = (value_col_id, operation, value_col_id) + result_expr = stacked_expr.aggregate( + [*index_aggregations, main_aggregation], + by_column_ids=[offset_col], + dropna=dropna, + ) + return Block( + result_expr.drop_columns([offset_col]), + self.index_columns, + column_labels=[None], + index_labels=self.index_labels, + ) def select_column(self, id: str) -> Block: return self.select_columns([id]) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 0d357e7c3d3..e4e22e03066 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1462,41 +1462,48 @@ def dropna( def any( self, *, + axis: typing.Union[str, int] = 0, bool_only: bool = False, ) -> bigframes.series.Series: if not bool_only: frame = self._raise_on_non_boolean("any") else: frame = self._drop_non_bool() - block = frame._block.aggregate_all_and_pivot( - agg_ops.any_op, dtype=pandas.BooleanDtype() + block = frame._block.aggregate_all_and_stack( + agg_ops.any_op, dtype=pandas.BooleanDtype(), axis=axis ) return bigframes.series.Series(block.select_column("values")) - def all(self, *, bool_only: bool = False) -> bigframes.series.Series: + def all( + self, axis: typing.Union[str, int] = 0, *, bool_only: bool = False + ) -> bigframes.series.Series: if not bool_only: frame = self._raise_on_non_boolean("all") else: frame = self._drop_non_bool() - block = frame._block.aggregate_all_and_pivot( - agg_ops.all_op, dtype=pandas.BooleanDtype() + block = frame._block.aggregate_all_and_stack( + agg_ops.all_op, dtype=pandas.BooleanDtype(), axis=axis ) return bigframes.series.Series(block.select_column("values")) - def sum(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def sum( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("sum") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.sum_op) + block = frame._block.aggregate_all_and_stack(agg_ops.sum_op, axis=axis) return bigframes.series.Series(block.select_column("values")) - def mean(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def mean( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("mean") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.mean_op) + block = frame._block.aggregate_all_and_stack(agg_ops.mean_op, axis=axis) return bigframes.series.Series(block.select_column("values")) def median( @@ -1510,47 +1517,57 @@ def median( frame = self._raise_on_non_numeric("median") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.median_op) + block = frame._block.aggregate_all_and_stack(agg_ops.median_op) return bigframes.series.Series(block.select_column("values")) - def std(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def std( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("std") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.std_op) + block = frame._block.aggregate_all_and_stack(agg_ops.std_op, axis=axis) return bigframes.series.Series(block.select_column("values")) - def var(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def var( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("var") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.var_op) + block = frame._block.aggregate_all_and_stack(agg_ops.var_op, axis=axis) return bigframes.series.Series(block.select_column("values")) - def min(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def min( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("min") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.min_op) + block = frame._block.aggregate_all_and_stack(agg_ops.min_op, axis=axis) return bigframes.series.Series(block.select_column("values")) - def max(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def max( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("max") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.max_op) + block = frame._block.aggregate_all_and_stack(agg_ops.max_op, axis=axis) return bigframes.series.Series(block.select_column("values")) - def prod(self, *, numeric_only: bool = False) -> bigframes.series.Series: + def prod( + self, axis: typing.Union[str, int] = 0, *, numeric_only: bool = False + ) -> bigframes.series.Series: if not numeric_only: frame = self._raise_on_non_numeric("prod") else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.product_op) + block = frame._block.aggregate_all_and_stack(agg_ops.product_op, axis=axis) return bigframes.series.Series(block.select_column("values")) product = prod @@ -1560,11 +1577,11 @@ def count(self, *, numeric_only: bool = False) -> bigframes.series.Series: frame = self else: frame = self._drop_non_numeric() - block = frame._block.aggregate_all_and_pivot(agg_ops.count_op) + block = frame._block.aggregate_all_and_stack(agg_ops.count_op) return bigframes.series.Series(block.select_column("values")) def nunique(self) -> bigframes.series.Series: - block = self._block.aggregate_all_and_pivot(agg_ops.nunique_op) + block = self._block.aggregate_all_and_stack(agg_ops.nunique_op) return bigframes.series.Series(block.select_column("values")) def agg( @@ -1587,7 +1604,7 @@ def agg( ) else: return bigframes.series.Series( - self._block.aggregate_all_and_pivot( + self._block.aggregate_all_and_stack( agg_ops.lookup_agg_func(typing.cast(str, func)) ) ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index b6ca958c038..adf17848ee6 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1999,6 +1999,29 @@ def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op): pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False) +@pytest.mark.parametrize( + ("op"), + [ + (lambda x: x.sum(axis=1, numeric_only=True)), + (lambda x: x.mean(axis=1, numeric_only=True)), + (lambda x: x.min(axis=1, numeric_only=True)), + (lambda x: x.max(axis=1, numeric_only=True)), + (lambda x: x.std(axis=1, numeric_only=True)), + (lambda x: x.var(axis=1, numeric_only=True)), + ], + ids=["sum", "mean", "min", "max", "std", "var"], +) +def test_dataframe_aggregates_axis_1(scalars_df_index, scalars_pandas_df_index, op): + col_names = ["int64_too", "int64_col", "float64_col", "bool_col", "string_col"] + bf_result = op(scalars_df_index[col_names]).to_pandas() + pd_result = op(scalars_pandas_df_index[col_names]) + + # Pandas may produce narrower numeric types, but bigframes always produces Float64 + pd_result = pd_result.astype("Float64") + # Pandas has object index type + pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) + + def test_dataframe_aggregates_median(scalars_df_index, scalars_pandas_df_index): col_names = ["int64_too", "float64_col", "int64_col", "bool_col"] bf_result = scalars_df_index[col_names].median(numeric_only=True).to_pandas() @@ -2019,11 +2042,16 @@ def test_dataframe_aggregates_median(scalars_df_index, scalars_pandas_df_index): [ (lambda x: x.all(bool_only=True)), (lambda x: x.any(bool_only=True)), + (lambda x: x.all(axis=1, bool_only=True)), + (lambda x: x.any(axis=1, bool_only=True)), ], - ids=["all", "any"], + ids=["all_axis0", "any_axis0", "all_axis1", "any_axis1"], ) def test_dataframe_bool_aggregates(scalars_df_index, scalars_pandas_df_index, op): # Pandas will drop nullable 'boolean' dtype so we convert first to bool, then cast back later + scalars_df_index = scalars_df_index.assign( + bool_col=scalars_df_index.bool_col.fillna(False) + ) scalars_pandas_df_index = scalars_pandas_df_index.assign( bool_col=scalars_pandas_df_index.bool_col.fillna(False).astype("bool") ) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 9d26938e080..6ce11cd7e98 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -11,7 +11,7 @@ """ from __future__ import annotations -from typing import Iterable, Literal, Mapping, Optional, Sequence, Union +from typing import Literal, Mapping, Optional, Sequence, Union import numpy as np @@ -1457,7 +1457,7 @@ def apply(self, func, *, args=(), **kwargs): # ---------------------------------------------------------------------- # ndarray-like stats methods - def any(self, *, bool_only: bool = False): + def any(self, *, axis=0, bool_only: bool = False): """ Return whether any element is True, potentially over an axis. @@ -1466,6 +1466,9 @@ def any(self, *, bool_only: bool = False): non-empty). Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. bool_only (bool. default False): Include only boolean columns. @@ -1474,7 +1477,7 @@ def any(self, *, bool_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def all(self, *, bool_only: bool = False): + def all(self, axis=0, *, bool_only: bool = False): """ Return whether all elements are True, potentially over an axis. @@ -1483,6 +1486,9 @@ def all(self, *, bool_only: bool = False): empty). Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. bool_only (bool. default False): Include only boolean columns. @@ -1491,11 +1497,14 @@ def all(self, *, bool_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def prod(self, *, numeric_only: bool = False): + def prod(self, axis=0, *, numeric_only: bool = False): """ Return the product of the values over the requested axis. Args: + aßxis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool. default False): Include only float, int, boolean columns. @@ -1504,13 +1513,16 @@ def prod(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def min(self, *, numeric_only: bool = False): + def min(self, axis=0, *, numeric_only: bool = False): """Return the minimum of the values over the requested axis. If you want the *index* of the minimum, use ``idxmin``. This is the equivalent of the ``numpy.ndarray`` method ``argmin``. Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool, default False): Default False. Include only float, int, boolean columns. @@ -1519,13 +1531,16 @@ def min(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def max(self, *, numeric_only: bool = False): + def max(self, axis=0, *, numeric_only: bool = False): """Return the maximum of the values over the requested axis. If you want the *index* of the maximum, use ``idxmax``. This is the equivalent of the ``numpy.ndarray`` method ``argmax``. Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool. default False): Default False. Include only float, int, boolean columns. @@ -1534,12 +1549,15 @@ def max(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def sum(self, *, numeric_only: bool = False): + def sum(self, axis=0, *, numeric_only: bool = False): """Return the sum of the values over the requested axis. This is equivalent to the method ``numpy.sum``. Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool. default False): Default False. Include only float, int, boolean columns. @@ -1548,10 +1566,13 @@ def sum(self, *, numeric_only: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def mean(self, *, numeric_only: bool = False): + def mean(self, axis=0, *, numeric_only: bool = False): """Return the mean of the values over the requested axis. Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool. default False): Default False. Include only float, int, boolean columns. @@ -1575,12 +1596,15 @@ def median(self, *, numeric_only: bool = False, exact: bool = False): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def var(self, *, numeric_only: bool = False): + def var(self, axis=0, *, numeric_only: bool = False): """Return unbiased variance over requested axis. Normalized by N-1 by default. Args: + axis ({index (0), columns (1)}): + Axis for the function to be applied on. + For Series this parameter is unused and defaults to 0. numeric_only (bool. default False): Default False. Include only float, int, boolean columns. From d56258cbfcda168cb9e437a021e282818d622d6a Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 26 Sep 2023 01:03:24 +0000 Subject: [PATCH 36/80] fix: Fix header skipping logic in `read_csv` (#49) Change-Id: Ib575e2c2b07f819d1dc499a271fea91107fbb8b4 --- bigframes/session.py | 7 +++---- tests/system/small/test_session.py | 18 ++++++++++++------ 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/bigframes/session.py b/bigframes/session.py index 04ae6ba4541..7b827c7dcfc 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -1050,11 +1050,10 @@ def read_csv( # We want to match pandas behavior. If header is 0, no rows should be skipped, so we # do not need to set `skip_leading_rows`. If header is None, then there is no header. # Setting skip_leading_rows to 0 does that. If header=N and N>0, we want to skip N rows. - # `skip_leading_rows` skips N-1 rows, so we set it to header+1. - if header is not None and header > 0: - job_config.skip_leading_rows = header + 1 - elif header is None: + if header is None: job_config.skip_leading_rows = 0 + elif header > 0: + job_config.skip_leading_rows = header return self._read_bigquery_load_job( filepath_or_buffer, diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index b7bee16ffd4..614c9537649 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -578,9 +578,12 @@ def test_read_csv_gcs_bq_engine_w_header(session, scalars_df_index, gcs_folder): path = gcs_folder + "test_read_csv_gcs_bq_engine_w_header*.csv" scalars_df_index.to_csv(path, index=False) - # Skip the header and the first 2 data rows. Without provided schema, the column names - # would be like `bool_field_0`, `string_field_1` and etc. - df = session.read_csv(path, header=2, engine="bigquery") + # Skip the header and the first 2 data rows. Note that one line of header + # also got added while writing the csv through `to_csv`, so we would have to + # pass headers=3 in the `read_csv` to skip reading the header and two rows. + # Without provided schema, the column names would be like `bool_field_0`, + # `string_field_1` and etc. + df = session.read_csv(path, header=3, engine="bigquery") assert df.shape[0] == scalars_df_index.shape[0] - 2 assert len(df.columns) == len(scalars_df_index.columns) @@ -609,9 +612,12 @@ def test_read_csv_local_bq_engine_w_header(session, scalars_pandas_df_index): # Using the pandas to_csv method because the BQ one does not support local write. scalars_pandas_df_index.to_csv(path, index=False) - # Skip the header and the first 2 data rows. Without provided schema, the column names - # would be like `bool_field_0`, `string_field_1` and etc. - df = session.read_csv(path, header=2, engine="bigquery") + # Skip the header and the first 2 data rows. Note that one line of + # header also got added while writing the csv through `to_csv`, so we + # would have to pass headers=3 in the `read_csv` to skip reading the + # header and two rows. Without provided schema, the column names would + # be like `bool_field_0`, `string_field_1` and etc. + df = session.read_csv(path, header=3, engine="bigquery") assert df.shape[0] == scalars_pandas_df_index.shape[0] - 2 assert len(df.columns) == len(scalars_pandas_df_index.columns) From 632caec420a7e23188f01b96a00c354d205da74e Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Mon, 25 Sep 2023 21:17:11 -0700 Subject: [PATCH 37/80] fix: LabelEncoder params consistent with Sklearn (#60) * fix: LabelEncoder params consistent with Sklearn * fix:add LabelTransformer * fix: address comments for base LabelTransformer * fix: type for params --- bigframes/ml/base.py | 20 +++++++++++++ bigframes/ml/preprocessing.py | 17 +++++------ tests/system/small/ml/test_preprocessing.py | 30 ++++--------------- .../sklearn/preprocessing/_label.py | 12 ++++---- 4 files changed, 39 insertions(+), 40 deletions(-) diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py index f07274f8fcb..f899ac71193 100644 --- a/bigframes/ml/base.py +++ b/bigframes/ml/base.py @@ -195,3 +195,23 @@ def fit_transform( y: Optional[Union[bpd.DataFrame, bpd.Series]] = None, ) -> bpd.DataFrame: return self.fit(X, y).transform(X) + + +class LabelTransformer(BaseEstimator): + """A BigQuery DataFrames Label Transformer base class that transforms data. + + Also the transformers can be attached to a pipeline with a predictor.""" + + @abc.abstractmethod + def fit(self, y): + pass + + @abc.abstractmethod + def transform(self, y): + pass + + def fit_transform( + self, + y: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: + return self.fit(y).transform(y) diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index f4f5446651c..ed0b36deef9 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -315,7 +315,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: class LabelEncoder( - base.Transformer, + base.LabelTransformer, third_party.bigframes_vendored.sklearn.preprocessing._label.LabelEncoder, ): # BQML max value https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder#syntax @@ -401,16 +401,15 @@ def _parse_from_sql(cls, sql: str) -> tuple[LabelEncoder, str]: def fit( self, - X: Union[bpd.DataFrame, bpd.Series], - y=None, # ignored + y: Union[bpd.DataFrame, bpd.Series], ) -> LabelEncoder: - (X,) = utils.convert_to_dataframe(X) + (y,) = utils.convert_to_dataframe(y) - compiled_transforms = self._compile_to_sql(X.columns.tolist()) + compiled_transforms = self._compile_to_sql(y.columns.tolist()) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] self._bqml_model = self._bqml_model_factory.create_model( - X, + y, options={"model_type": "transform_only"}, transforms=transform_sqls, ) @@ -419,13 +418,13 @@ def fit( self._output_names = [name for _, name in compiled_transforms] return self - def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + def transform(self, y: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("Must be fitted before transform") - (X,) = utils.convert_to_dataframe(X) + (y,) = utils.convert_to_dataframe(y) - df = self._bqml_model.transform(X) + df = self._bqml_model.transform(y) return typing.cast( bpd.DataFrame, df[self._output_names], diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index 7779eb8f6ea..61bddb144d6 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -357,9 +357,9 @@ def test_one_hot_encoder_different_data(penguins_df_default_index, new_penguins_ def test_label_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.LabelEncoder() - encoder.fit(new_penguins_df[["species", "sex"]]) + encoder.fit(new_penguins_df["species"]) - result = encoder.transform(new_penguins_df).to_pandas() + result = encoder.transform(new_penguins_df["species"]).to_pandas() # TODO: bug? feature columns seem to be in nondeterministic random order # workaround: sort columns by name. Can't repro it in pantheon, so could @@ -368,11 +368,6 @@ def test_label_encoder_default_params(new_penguins_df): expected = pd.DataFrame( { - "labelencoded_sex": [ - 2, - 1, - 1, - ], "labelencoded_species": [ 1, 1, @@ -389,7 +384,7 @@ def test_label_encoder_default_params(new_penguins_df): def test_label_encoder_default_params_fit_transform(new_penguins_df): encoder = bigframes.ml.preprocessing.LabelEncoder() - result = encoder.fit_transform(new_penguins_df[["species", "sex"]]).to_pandas() + result = encoder.fit_transform(new_penguins_df[["species"]]).to_pandas() # TODO: bug? feature columns seem to be in nondeterministic random order # workaround: sort columns by name. Can't repro it in pantheon, so could @@ -398,11 +393,6 @@ def test_label_encoder_default_params_fit_transform(new_penguins_df): expected = pd.DataFrame( { - "labelencoded_sex": [ - 2, - 1, - 1, - ], "labelencoded_species": [ 1, 1, @@ -444,7 +434,7 @@ def test_label_encoder_series_default_params(new_penguins_df): def test_label_encoder_params(new_penguins_df): encoder = bigframes.ml.preprocessing.LabelEncoder(100, 2) - encoder.fit(new_penguins_df[["species", "sex"]]) + encoder.fit(new_penguins_df[["species"]]) result = encoder.transform(new_penguins_df).to_pandas() @@ -455,11 +445,6 @@ def test_label_encoder_params(new_penguins_df): expected = pd.DataFrame( { - "labelencoded_sex": [ - 0, - 0, - 0, - ], "labelencoded_species": [ 0, 0, @@ -475,7 +460,7 @@ def test_label_encoder_params(new_penguins_df): def test_label_encoder_different_data(penguins_df_default_index, new_penguins_df): encoder = bigframes.ml.preprocessing.LabelEncoder() - encoder.fit(penguins_df_default_index[["species", "sex"]]) + encoder.fit(penguins_df_default_index[["species"]]) result = encoder.transform(new_penguins_df).to_pandas() @@ -486,11 +471,6 @@ def test_label_encoder_different_data(penguins_df_default_index, new_penguins_df expected = pd.DataFrame( { - "labelencoded_sex": [ - 3, - 2, - 2, - ], "labelencoded_species": [ 1, 1, diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_label.py b/third_party/bigframes_vendored/sklearn/preprocessing/_label.py index 7e60c846d45..83f8eb0f9c2 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_label.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_label.py @@ -28,11 +28,11 @@ class LabelEncoder(BaseEstimator): Default None, set limit to 1,000,000. """ - def fit(self, X): - """Fit LabelEncoder to X. + def fit(self, y): + """Fit label encoder. Args: - X (bigframes.dataframe.DataFrame or bigframes.series.Series): + y (bigframes.dataframe.DataFrame or bigframes.series.Series): The DataFrame or Series with training data. Returns: @@ -40,11 +40,11 @@ def fit(self, X): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) - def transform(self, X): - """Transform X using label encoding. + def transform(self, y): + """Transform y using label encoding. Args: - X (bigframes.dataframe.DataFrame or bigframes.series.Series): + y (bigframes.dataframe.DataFrame or bigframes.series.Series): The DataFrame or Series to be transformed. Returns: From 3502f835b35c437933430698e7a1c9badaddcb99 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 26 Sep 2023 12:09:25 -0700 Subject: [PATCH 38/80] feat: support casting string to integer or float (#59) --- bigframes/dtypes.py | 2 +- tests/system/small/test_series.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 271b8aa2f24..59d3007fab0 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -313,7 +313,7 @@ def cast_ibis_value( ibis_dtypes.string, ), ibis_dtypes.float64: (ibis_dtypes.string, ibis_dtypes.int64), - ibis_dtypes.string: (), + ibis_dtypes.string: (ibis_dtypes.int64, ibis_dtypes.float64), ibis_dtypes.date: (), ibis_dtypes.time: (), ibis_dtypes.timestamp: (ibis_dtypes.Timestamp(timezone="UTC"),), diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index d702049e684..588dcc2c83a 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2389,6 +2389,29 @@ def test_astype(scalars_df_index, scalars_pandas_df_index, column, to_type): pd.testing.assert_series_equal(bf_result, pd_result) +def test_string_astype_int(): + pd_series = pd.Series(["4", "-7", "0", " -03"]) + bf_series = series.Series(pd_series) + + pd_result = pd_series.astype("Int64") + bf_result = bf_series.astype("Int64").to_pandas() + + pd.testing.assert_series_equal(bf_result, pd_result, check_index_type=False) + + +def test_string_astype_float(): + pd_series = pd.Series( + ["1", "-1", "-0", "000", " -03.235", "naN", "-inf", "INf", ".33", "7.235e-8"] + ) + + bf_series = series.Series(pd_series) + + pd_result = pd_series.astype("Float64") + bf_result = bf_series.astype("Float64").to_pandas() + + pd.testing.assert_series_equal(bf_result, pd_result, check_index_type=False) + + @pytest.mark.parametrize( "index", [0, 5, -2], From a6e32aa875370063c48ce7922c2aa369a770bd30 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 26 Sep 2023 13:40:07 -0700 Subject: [PATCH 39/80] feat: Add more index methods (#54) --- bigframes/core/block_transforms.py | 10 +- bigframes/core/indexes/index.py | 118 +++++++++++++-- bigframes/dataframe.py | 7 +- bigframes/series.py | 6 +- tests/system/small/test_index.py | 120 ++++++++++++++++ .../pandas/core/indexes/base.py | 135 ++++++++++++++++++ 6 files changed, 373 insertions(+), 23 deletions(-) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index da6ba65b8a9..d22112417cc 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -218,13 +218,17 @@ def rank( return block.select_columns(rownum_col_ids).with_column_labels(labels) -def dropna(block: blocks.Block, how: typing.Literal["all", "any"] = "any"): +def dropna( + block: blocks.Block, + column_ids: typing.Sequence[str], + how: typing.Literal["all", "any"] = "any", +): """ Drop na entries from block """ if how == "any": filtered_block = block - for column in block.value_columns: + for column in column_ids: filtered_block, result_id = filtered_block.apply_unary_op( column, ops.notnull_op ) @@ -234,7 +238,7 @@ def dropna(block: blocks.Block, how: typing.Literal["all", "any"] = "any"): else: # "all" filtered_block = block predicate = None - for column in block.value_columns: + for column in column_ids: filtered_block, partial_predicate = filtered_block.apply_unary_op( column, ops.notnull_op ) diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index c08c851c91d..f211afe4d56 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -24,8 +24,10 @@ import bigframes.constants as constants import bigframes.core as core +import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks import bigframes.core.joins as joins +import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.dtypes import bigframes.dtypes as bf_dtypes @@ -149,6 +151,27 @@ def has_duplicates(self) -> bool: def _block(self) -> blocks.Block: return self._data._get_block() + @property + def T(self) -> Index: + return self.transpose() + + def transpose(self) -> Index: + return self + + def sort_values(self, *, ascending: bool = True, na_position: str = "last"): + if na_position not in ["first", "last"]: + raise ValueError("Param na_position must be one of 'first' or 'last'") + direction = ( + order.OrderingDirection.ASC if ascending else order.OrderingDirection.DESC + ) + na_last = na_position == "last" + index_columns = self._block.index_columns + ordering = [ + order.OrderingColumnReference(column, direction=direction, na_last=na_last) + for column in index_columns + ] + return Index._from_block(self._block.order_by(ordering)) + def astype( self, dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], @@ -176,6 +199,57 @@ def max(self) -> typing.Any: def min(self) -> typing.Any: return self._apply_aggregation(agg_ops.min_op) + def argmax(self) -> int: + block, row_nums = self._block.promote_offsets() + block = block.order_by( + [ + *[ + order.OrderingColumnReference( + col, direction=order.OrderingDirection.DESC + ) + for col in self._block.index_columns + ], + order.OrderingColumnReference(row_nums), + ] + ) + import bigframes.series as series + + return typing.cast(int, series.Series(block.select_column(row_nums)).iloc[0]) + + def argmin(self) -> int: + block, row_nums = self._block.promote_offsets() + block = block.order_by( + [ + *[ + order.OrderingColumnReference(col) + for col in self._block.index_columns + ], + order.OrderingColumnReference(row_nums), + ] + ) + import bigframes.series as series + + return typing.cast(int, series.Series(block.select_column(row_nums)).iloc[0]) + + def value_counts( + self, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + *, + dropna: bool = True, + ): + block = block_ops.value_counts( + self._block, + self._block.index_columns, + normalize=normalize, + ascending=ascending, + dropna=dropna, + ) + import bigframes.series as series + + return series.Series(block) + def fillna(self, value=None) -> Index: if self.nlevels > 1: raise TypeError("Multiindex does not support 'fillna'") @@ -185,10 +259,7 @@ def rename(self, name: Union[str, Sequence[str]]) -> Index: names = [name] if isinstance(name, str) else list(name) if len(names) != self.nlevels: raise ValueError("'name' must be same length as levels") - - import bigframes.dataframe as df - - return Index(df.DataFrame(self._block.with_index_labels(names))) + return Index._from_block(self._block.with_index_labels(names)) def drop( self, @@ -210,9 +281,28 @@ def drop( ) block = block.filter(condition_id, keep_null=True) block = block.drop_columns([condition_id]) - import bigframes.dataframe as df + return Index._from_block(block) + + def dropna(self, how: str = "any") -> Index: + if how not in ("any", "all"): + raise ValueError("'how' must be one of 'any', 'all'") + result = block_ops.dropna(self._block, self._block.index_columns, how=how) # type: ignore + return Index._from_block(result) + + def drop_duplicates(self, *, keep: str = "first") -> Index: + block = block_ops.drop_duplicates(self._block, self._block.index_columns, keep) + return Index._from_block(block) + + def isin(self, values) -> Index: + if not utils.is_list_like(values): + raise TypeError( + "only list-like objects are allowed to be passed to " + f"isin(), you passed a [{type(values).__name__}]" + ) - return Index(df.DataFrame(block.select_columns([]))) + return self._apply_unary_op(ops.IsInOp(values, match_nulls=True)).fillna( + value=False + ) def _apply_unary_op( self, @@ -226,9 +316,7 @@ def _apply_unary_op( result_ids.append(result_id) block = block.set_index(result_ids, index_labels=self._block.index_labels) - import bigframes.dataframe as df - - return Index(df.DataFrame(block)) + return Index._from_block(block) def _apply_aggregation(self, op: agg_ops.AggregateOp) -> typing.Any: if self.nlevels > 1: @@ -262,6 +350,12 @@ def to_numpy(self, dtype=None, **kwargs) -> np.ndarray: def __len__(self): return self.shape[0] + @classmethod + def _from_block(cls, block: blocks.Block) -> Index: + import bigframes.dataframe as df + + return Index(df.DataFrame(block)) + class IndexValue: """An immutable index.""" @@ -356,12 +450,6 @@ def resolve_level_name(self: IndexValue, label: blocks.Label) -> str: def is_uniquely_named(self: IndexValue): return len(set(self.names)) == len(self.names) - def _set_block(self, block: blocks.Block): - self._block = block - - def _get_block(self) -> blocks.Block: - return self._block - def join_mono_indexed( left: IndexValue, diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index e4e22e03066..113355589b1 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1440,7 +1440,7 @@ def dropna( axis_n = utils.get_axis_number(axis) if axis_n == 0: - result = block_ops.dropna(self._block, how=how) # type: ignore + result = block_ops.dropna(self._block, self._block.value_columns, how=how) # type: ignore if ignore_index: result = result.reset_index() return DataFrame(result) @@ -1674,7 +1674,10 @@ def pivot( def stack(self): # TODO: support 'level' param by simply reordering levels such that selected level is last before passing to Block.stack. # TODO: match impl to pandas future_stack as described in pandas 2.1 release notes - result_block = block_ops.dropna(self._block.stack(), how="all") + stack_block = self._block.stack() + result_block = block_ops.dropna( + stack_block, stack_block.value_columns, how="all" + ) if not isinstance(self.columns, pandas.MultiIndex): return bigframes.series.Series(result_block) return DataFrame(result_block) diff --git a/bigframes/series.py b/bigframes/series.py index c1c0cb05372..47298d59f50 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -459,7 +459,7 @@ def dropna( ) -> Series: if inplace: raise NotImplementedError("'inplace'=True not supported") - result = block_ops.dropna(self._block, how="any") + result = block_ops.dropna(self._block, [self._value_column], how="any") if ignore_index: result = result.reset_index() return Series(result) @@ -856,7 +856,7 @@ def clip(self, lower, upper): ) return Series(block.select_column(result_id).with_column_labels([self.name])) - def argmax(self) -> scalars.Scalar: + def argmax(self) -> int: block, row_nums = self._block.promote_offsets() block = block.order_by( [ @@ -870,7 +870,7 @@ def argmax(self) -> scalars.Scalar: scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] ) - def argmin(self) -> scalars.Scalar: + def argmin(self) -> int: block, row_nums = self._block.promote_offsets() block = block.order_by( [ diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py index 7f09e3a9d51..f7fa0f0855e 100644 --- a/tests/system/small/test_index.py +++ b/tests/system/small/test_index.py @@ -14,6 +14,7 @@ import numpy import pandas as pd +import pytest from tests.system.utils import assert_pandas_index_equal_ignore_index_type @@ -174,3 +175,122 @@ def test_is_monotonic_decreasing(scalars_df_index, scalars_pandas_df_index): scalars_df_index.index.is_monotonic_increasing == scalars_pandas_df_index.index.is_monotonic_increasing ) + + +def test_index_argmin(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("1."): + pytest.skip("doesn't work in pandas 1.x.") + bf_result = scalars_df_index.set_index(["int64_too", "rowindex_2"]).index.argmin() + pd_result = scalars_pandas_df_index.set_index( + ["int64_too", "rowindex_2"] + ).index.argmin() + assert bf_result == pd_result + + +def test_index_argmax(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("1."): + pytest.skip("doesn't work in pandas 1.x.") + bf_result = scalars_df_index.set_index(["int64_too", "rowindex_2"]).index.argmax() + pd_result = scalars_pandas_df_index.set_index( + ["int64_too", "rowindex_2"] + ).index.argmax() + assert bf_result == pd_result + + +@pytest.mark.parametrize( + ("ascending", "na_position"), + [ + (True, "first"), + (True, "last"), + (False, "first"), + (False, "last"), + ], +) +def test_index_sort_values( + scalars_df_index, scalars_pandas_df_index, ascending, na_position +): + # Test needs values to be unique + bf_result = ( + scalars_df_index.set_index(["int64_too", "rowindex_2"]) + .index.sort_values(ascending=ascending, na_position=na_position) + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index( + ["int64_too", "rowindex_2"] + ).index.sort_values(ascending=ascending, na_position=na_position) + + pd.testing.assert_index_equal( + bf_result, + pd_result, + ) + + +def test_index_value_counts(scalars_df_index, scalars_pandas_df_index): + if pd.__version__.startswith("1."): + pytest.skip("value_counts results different in pandas 1.x.") + bf_result = ( + scalars_df_index.set_index(["int64_too", "rowindex_2"]) + .index.value_counts() + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index( + ["int64_too", "rowindex_2"] + ).index.value_counts() + + pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("how",), + [ + ("any",), + ("all",), + ], +) +def test_index_dropna(scalars_df_index, scalars_pandas_df_index, how): + bf_result = ( + scalars_df_index.set_index(["int64_col", "float64_col"]) + .index.dropna(how=how) + .to_pandas() + ) + pd_result = scalars_pandas_df_index.set_index( + ["int64_col", "float64_col"] + ).index.dropna(how=how) + pd.testing.assert_index_equal(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("keep",), + [ + ("first",), + ("last",), + (False,), + ], +) +def test_index_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep): + bf_series = ( + scalars_df_index.set_index("int64_col") + .index.drop_duplicates(keep=keep) + .to_pandas() + ) + pd_series = scalars_pandas_df_index.set_index("int64_col").index.drop_duplicates( + keep=keep + ) + pd.testing.assert_index_equal( + pd_series, + bf_series, + ) + + +def test_index_isin(scalars_df_index, scalars_pandas_df_index): + bf_series = ( + scalars_df_index.set_index("int64_col").index.isin([2, 55555, 4]).to_pandas() + ) + pd_result_array = scalars_pandas_df_index.set_index("int64_col").index.isin( + [2, 55555, 4] + ) + pd.testing.assert_index_equal( + pd.Index(pd_result_array), + bf_series, + check_names=False, + ) diff --git a/third_party/bigframes_vendored/pandas/core/indexes/base.py b/third_party/bigframes_vendored/pandas/core/indexes/base.py index f89964e2208..e8737341a34 100644 --- a/third_party/bigframes_vendored/pandas/core/indexes/base.py +++ b/third_party/bigframes_vendored/pandas/core/indexes/base.py @@ -53,6 +53,20 @@ def dtypes(self): """Return the dtypes as a Series for the underlying MultiIndex.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property + def T(self) -> Index: + """Return the transpose, which is by definition self.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transpose(self) -> Index: + """ + Return the transpose, which is by definition self. + + Returns: + Index + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def astype(self, dtype): """Create an Index with values cast to dtypes. @@ -67,6 +81,23 @@ def astype(self, dtype): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def isin(self, values): + """ + Return a boolean array where the index values are in `values`. + + Compute boolean array of whether each index value is found in the + passed set of values. The length of the returned boolean array matches + the length of the index. + + Args: + values (set or list-like): + Sought values. + + Returns: + Series: Series of boolean values. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def all(self) -> bool: """Return whether all elements are Truthy. @@ -99,6 +130,30 @@ def max(self): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def argmin(self) -> int: + """ + Return int position of the smallest value in the Series. + + If the minimum is achieved in multiple locations, + the first row position is returned. + + Returns: + int: Row position of the minimum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def argmax(self) -> int: + """ + Return int position of the largest value in the Series. + + If the maximum is achieved in multiple locations, + the first row position is returned. + + Returns: + int: Row position of the maximum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def nunique(self) -> int: """Return number of unique elements in the object. @@ -109,6 +164,57 @@ def nunique(self) -> int: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def sort_values( + self, *, ascending: bool = True, na_position: str = "last" + ) -> Index: + """ + Return a sorted copy of the index. + + Return a sorted copy of the index, and optionally return the indices + that sorted the index itself. + + Args: + ascending (bool, default True): + Should the index values be sorted in an ascending order. + na_position ({'first' or 'last'}, default 'last'): + Argument 'first' puts NaNs at the beginning, 'last' puts NaNs at + the end. + + Returns: + pandas.Index: Sorted copy of the index. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def value_counts( + self, + normalize: bool = True, + sort: bool = True, + ascending: bool = False, + *, + dropna: bool = True, + ): + """Return a Series containing counts of unique values. + + The resulting object will be in descending order so that the + first element is the most frequently-occurring element. + Excludes NA values by default. + + Args: + normalize (bool, default False): + If True then the object returned will contain the relative + frequencies of the unique values. + sort (bool, default True): + Sort by frequencies. + ascending (bool, default False): + Sort in ascending order. + dropna (bool, default True): + Don't include counts of NaN. + + Returns: + Series + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def fillna(self, value) -> Index: """ Fill NA/NaN values with the specified value. @@ -151,6 +257,35 @@ def drop(self, labels) -> Index: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def dropna(self, how: str = "any"): + """Return Index without NA/NaN values. + + Args: + how ({'any', 'all'}, default 'any'): + If the Index is a MultiIndex, drop the value when any or all levels + are NaN. + + Returns: + Index + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def drop_duplicates(self, *, keep: str = "first"): + """ + Return Index with duplicate values removed. + + Args: + keep ({'first', 'last', ``False``}, default 'first'): + One of: + 'first' : Drop duplicates except for the first occurrence. + 'last' : Drop duplicates except for the last occurrence. + ``False`` : Drop all duplicates. + + Returns: + Index + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def to_numpy(self, dtype): """ A NumPy ndarray representing the values in this Series or Index. From e804e130c218d8ac81a8fc0a853eeb7a93884a50 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Tue, 26 Sep 2023 14:42:40 -0700 Subject: [PATCH 40/80] refactor: push down SQL generate logic in core.BqmlModelFactory (#62) --- bigframes/ml/core.py | 33 ++++++------------------ bigframes/ml/sql.py | 23 ++++++++++++----- tests/unit/ml/test_sql.py | 54 +++++++++++++++++++++++++++++---------- 3 files changed, 64 insertions(+), 46 deletions(-) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 70be0d35ee1..37478d8bafd 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -251,17 +251,10 @@ def create_model( session = X_train._session - source_sql = input_data.sql - options_sql = self._model_creation_sql_generator.options(**options) - transform_sql = ( - self._model_creation_sql_generator.transform(*transforms) - if transforms is not None - else None - ) sql = self._model_creation_sql_generator.create_model( - source_sql=source_sql, - transform_sql=transform_sql, - options_sql=options_sql, + source=input_data, + transforms=transforms, + options=options, ) return self._create_model_with_sql(session=session, sql=sql) @@ -287,18 +280,10 @@ def create_time_series_model( session = X_train._session - source_sql = input_data.sql - options_sql = self._model_creation_sql_generator.options(**options) - - transform_sql = ( - self._model_creation_sql_generator.transform(*transforms) - if transforms is not None - else None - ) sql = self._model_creation_sql_generator.create_model( - source_sql=source_sql, - transform_sql=transform_sql, - options_sql=options_sql, + source=input_data, + transforms=transforms, + options=options, ) return self._create_model_with_sql(session=session, sql=sql) @@ -320,10 +305,9 @@ def create_remote_model( Returns: BqmlModel: a BqmlModel wrapping a trained model in BigQuery """ - options_sql = self._model_creation_sql_generator.options(**options) sql = self._model_creation_sql_generator.create_remote_model( connection_name=connection_name, - options_sql=options_sql, + options=options, ) return self._create_model_with_sql(session=session, sql=sql) @@ -341,9 +325,8 @@ def create_imported_model( Returns: a BqmlModel, wrapping a trained model in BigQuery """ - options_sql = self._model_creation_sql_generator.options(**options) sql = self._model_creation_sql_generator.create_imported_model( - options_sql=options_sql, + options=options, ) return self._create_model_with_sql(session=session, sql=sql) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index a54d39e6b2f..c1b4f46e185 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -16,9 +16,10 @@ Generates SQL queries needed for BigQuery DataFrames ML """ -from typing import Iterable, Optional, Union +from typing import Iterable, Mapping, Optional, Union import bigframes.constants as constants +import bigframes.pandas as bpd class BaseSqlGenerator: @@ -113,11 +114,15 @@ def __init__(self, model_id: str): # Model create and alter def create_model( self, - source_sql: str, - transform_sql: Optional[str] = None, - options_sql: Optional[str] = None, + source: bpd.DataFrame, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, + transforms: Optional[Iterable[str]] = None, ) -> str: """Encode the CREATE TEMP MODEL statement for BQML""" + source_sql = source.sql + transform_sql = self.transform(*transforms) if transforms is not None else None + options_sql = self.options(**options) + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] if transform_sql: parts.append(transform_sql) @@ -129,9 +134,11 @@ def create_model( def create_remote_model( self, connection_name: str, - options_sql: Optional[str] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + options_sql = self.options(**options) + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] parts.append(self.connection(connection_name)) if options_sql: @@ -140,9 +147,11 @@ def create_remote_model( def create_imported_model( self, - options_sql: Optional[str] = None, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + options_sql = self.options(**options) + parts = [f"CREATE TEMP MODEL `{self._model_id}`"] if options_sql: parts.append(options_sql) @@ -150,7 +159,7 @@ def create_imported_model( class ModelManipulationSqlGenerator(BaseSqlGenerator): - """Sql generator for manipulating a model entity. Model name is the fully model path of project_id.dataset_id.model_id.""" + """Sql generator for manipulating a model entity. Model name is the full model path of project_id.dataset_id.model_id.""" def __init__(self, model_name: str): self._model_name = model_name diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index c1b29c5e52d..b88523c7efa 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + import pytest import bigframes.ml.sql as ml_sql +import bigframes.pandas as bpd @pytest.fixture(scope="session") @@ -34,6 +37,14 @@ def model_manipulation_sql_generator() -> ml_sql.ModelManipulationSqlGenerator: ) +@pytest.fixture(scope="session") +def mock_df(): + mock_df = mock.create_autospec(spec=bpd.DataFrame) + mock_df.sql = "input_X_y_sql" + + return mock_df + + def test_options_produces_correct_sql(base_sql_generator: ml_sql.BaseSqlGenerator): sql = base_sql_generator.options( model_type="lin_reg", input_label_cols=["col_a"], l1_reg=0.6 @@ -96,33 +107,44 @@ def test_label_encoder_produces_correct_sql( def test_create_model_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, + mock_df: bpd.DataFrame, ): sql = model_creation_sql_generator.create_model( - source_sql="my_source_sql", - options_sql="my_options_sql", + source=mock_df, + options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql == """CREATE TEMP MODEL `my_model_id` -my_options_sql -AS my_source_sql""" +OPTIONS( + option_key1="option_value1", + option_key2=2) +AS input_X_y_sql""" ) def test_create_model_transform_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, + mock_df: bpd.DataFrame, ): sql = model_creation_sql_generator.create_model( - source_sql="my_source_sql", - options_sql="my_options_sql", - transform_sql="my_transform_sql", + source=mock_df, + options={"option_key1": "option_value1", "option_key2": 2}, + transforms=[ + "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", + "ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b", + ], ) assert ( sql == """CREATE TEMP MODEL `my_model_id` -my_transform_sql -my_options_sql -AS my_source_sql""" +TRANSFORM( + ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a, + ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b) +OPTIONS( + option_key1="option_value1", + option_key2=2) +AS input_X_y_sql""" ) @@ -131,13 +153,15 @@ def test_create_remote_model_produces_correct_sql( ): sql = model_creation_sql_generator.create_remote_model( connection_name="my_project.us.my_connection", - options_sql="my_options_sql", + options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql == """CREATE TEMP MODEL `my_model_id` REMOTE WITH CONNECTION `my_project.us.my_connection` -my_options_sql""" +OPTIONS( + option_key1="option_value1", + option_key2=2)""" ) @@ -145,12 +169,14 @@ def test_create_imported_model_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, ): sql = model_creation_sql_generator.create_imported_model( - options_sql="my_options_sql", + options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql == """CREATE TEMP MODEL `my_model_id` -my_options_sql""" +OPTIONS( + option_key1="option_value1", + option_key2=2)""" ) From 392113b70d6a8c407accbb6684d75b31261e3741 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 26 Sep 2023 16:25:59 -0700 Subject: [PATCH 41/80] feat: add ml.preprocessing.MinMaxScaler (#64) * feat: add ml.preprocessing.MinMaxScaler * fix comments and typo * add test check for min value * nit fix --- bigframes/ml/compose.py | 1 + bigframes/ml/pipeline.py | 12 ++- bigframes/ml/preprocessing.py | 84 ++++++++++++++++- bigframes/ml/sql.py | 4 + tests/system/large/ml/test_pipeline.py | 46 ++++++++- tests/system/small/ml/test_preprocessing.py | 93 +++++++++++++++++++ tests/unit/ml/test_compose.py | 25 +++++ tests/unit/ml/test_sql.py | 7 ++ .../sklearn/preprocessing/_data.py | 36 +++++++ 9 files changed, 302 insertions(+), 6 deletions(-) diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index a1075c23981..9effbf1968c 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -30,6 +30,7 @@ preprocessing.OneHotEncoder, preprocessing.StandardScaler, preprocessing.MaxAbsScaler, + preprocessing.MinMaxScaler, preprocessing.LabelEncoder, ] diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index 86b20996194..ac02c391129 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -51,6 +51,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): preprocessing.StandardScaler, preprocessing.OneHotEncoder, preprocessing.MaxAbsScaler, + preprocessing.MinMaxScaler, preprocessing.LabelEncoder, ), ): @@ -149,6 +150,7 @@ def _extract_as_column_transformer( preprocessing.OneHotEncoder, preprocessing.StandardScaler, preprocessing.MaxAbsScaler, + preprocessing.MinMaxScaler, preprocessing.LabelEncoder, ], Union[str, List[str]], @@ -177,10 +179,17 @@ def _extract_as_column_transformer( elif transform_sql.startswith("ML.MAX_ABS_SCALER"): transformers.append( ( - "max_abs_encoder", + "max_abs_scaler", *preprocessing.MaxAbsScaler._parse_from_sql(transform_sql), ) ) + elif transform_sql.startswith("ML.MIN_MAX_SCALER"): + transformers.append( + ( + "min_max_scaler", + *preprocessing.MinMaxScaler._parse_from_sql(transform_sql), + ) + ) elif transform_sql.startswith("ML.LABEL_ENCODER"): transformers.append( ( @@ -203,6 +212,7 @@ def _merge_column_transformer( preprocessing.StandardScaler, preprocessing.OneHotEncoder, preprocessing.MaxAbsScaler, + preprocessing.MinMaxScaler, preprocessing.LabelEncoder, ]: """Try to merge the column transformer to a simple transformer.""" diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index ed0b36deef9..caf4657a633 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -144,13 +144,13 @@ def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: @classmethod def _parse_from_sql(cls, sql: str) -> tuple[MaxAbsScaler, str]: - """Parse SQL to tuple(StandardScaler, column_label). + """Parse SQL to tuple(MaxAbsScaler, column_label). Args: sql: SQL string of format "ML.MAX_ABS_SCALER({col_label}) OVER()" Returns: - tuple(StandardScaler, column_label)""" + tuple(MaxAbsScaler, column_label)""" col_label = sql[sql.find("(") + 1 : sql.find(")")] return cls(), col_label @@ -187,6 +187,86 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +class MinMaxScaler( + base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._data.MinMaxScaler, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._data.MinMaxScaler.__doc__ + ) + + def __init__(self): + self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return type(other) is MinMaxScaler and self._bqml_model == other._bqml_model + + def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: a list of column names to transform + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + ( + self._base_sql_generator.ml_min_max_scaler( + column, f"min_max_scaled_{column}" + ), + f"min_max_scaled_{column}", + ) + for column in columns + ] + + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[MinMaxScaler, str]: + """Parse SQL to tuple(MinMaxScaler, column_label). + + Args: + sql: SQL string of format "ML.MIN_MAX_SCALER({col_label}) OVER()" + + Returns: + tuple(MinMaxScaler, column_label)""" + col_label = sql[sql.find("(") + 1 : sql.find(")")] + return cls(), col_label + + def fit( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> MinMaxScaler: + (X,) = utils.convert_to_dataframe(X) + + compiled_transforms = self._compile_to_sql(X.columns.tolist()) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = self._bqml_model_factory.create_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + return self + + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.transform(X) + return typing.cast( + bpd.DataFrame, + df[self._output_names], + ) + + class OneHotEncoder( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index c1b4f46e185..3897d1be390 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -81,6 +81,10 @@ def ml_max_abs_scaler(self, numeric_expr_sql: str, name: str) -> str: """Encode ML.MAX_ABS_SCALER for BQML""" return f"""ML.MAX_ABS_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_min_max_scaler(self, numeric_expr_sql: str, name: str) -> str: + """Encode ML.MIN_MAX_SCALER for BQML""" + return f"""ML.MIN_MAX_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_one_hot_encoder( self, numeric_expr_sql: str, diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index c69a00b81ca..34a2ca01011 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -575,6 +575,11 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + preprocessing.MinMaxScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -647,6 +652,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + preprocessing.MinMaxScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -684,9 +694,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id "species", ), ("standard_scaler", preprocessing.StandardScaler(), "culmen_length_mm"), - ("max_abs_encoder", preprocessing.MaxAbsScaler(), "culmen_length_mm"), + ("max_abs_scaler", preprocessing.MaxAbsScaler(), "culmen_length_mm"), + ("min_max_scaler", preprocessing.MinMaxScaler(), "culmen_length_mm"), ("standard_scaler", preprocessing.StandardScaler(), "flipper_length_mm"), - ("max_abs_encoder", preprocessing.MaxAbsScaler(), "flipper_length_mm"), + ("max_abs_scaler", preprocessing.MaxAbsScaler(), "flipper_length_mm"), + ("min_max_scaler", preprocessing.MinMaxScaler(), "flipper_length_mm"), ] assert transformers == expected @@ -743,7 +755,7 @@ def test_pipeline_max_abs_scaler_to_gbq(penguins_df_default_index, dataset_id): pl.fit(X_train, y_train) pl_loaded = pl.to_gbq( - f"{dataset_id}.test_penguins_pipeline_standard_scaler", replace=True + f"{dataset_id}.test_penguins_pipeline_min_max_scaler", replace=True ) assert isinstance(pl_loaded._transform, preprocessing.MaxAbsScaler) @@ -751,6 +763,34 @@ def test_pipeline_max_abs_scaler_to_gbq(penguins_df_default_index, dataset_id): assert pl_loaded._estimator.fit_intercept is False +def test_pipeline_min_max_scaler_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ("transform", preprocessing.MinMaxScaler()), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_min_max_scaler", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.MinMaxScaler) + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False + + def test_pipeline_one_hot_encoder_to_gbq(penguins_df_default_index, dataset_id): pl = pipeline.Pipeline( [ diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index 61bddb144d6..fc8f3251bda 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -211,6 +211,99 @@ def test_max_abs_scaler_series_normalizes(penguins_df_default_index, new_penguin pd.testing.assert_frame_equal(result, expected, rtol=1e-3) +def test_min_max_scaler_normalizeds_fit_transform(new_penguins_df): + scaler = bigframes.ml.preprocessing.MinMaxScaler() + result = scaler.fit_transform( + new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] + ).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "min_max_scaled_culmen_depth_mm": [1.0, 0.0, 0.5625], + "min_max_scaled_culmen_length_mm": [1.0, 0.375, 0.0], + "min_max_scaled_flipper_length_mm": [1.0, 0.0, 0.466667], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_min_max_scaler_series_normalizes(penguins_df_default_index, new_penguins_df): + scaler = bigframes.ml.preprocessing.MinMaxScaler() + scaler.fit(penguins_df_default_index["culmen_length_mm"]) + + result = scaler.transform(penguins_df_default_index["culmen_length_mm"]).to_pandas() + + # If minmax-scaled correctly, min should be 0 and max should be 1. + for column in result.columns: + assert math.isclose(result[column].max(), 1.0, abs_tol=1e-3) + assert math.isclose(result[column].min(), 0.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "min_max_scaled_culmen_length_mm": [0.269091, 0.232727, 0.210909], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_min_max_scaler_normalizes(penguins_df_default_index, new_penguins_df): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + scaler = bigframes.ml.preprocessing.MinMaxScaler() + scaler.fit( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + result = scaler.transform( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ).to_pandas() + + # If minmax-scaled correctly, min should be 0 and max should be 1. + for column in result.columns: + assert math.isclose(result[column].max(), 1.0, abs_tol=1e-3) + assert math.isclose(result[column].min(), 0.0, abs_tol=1e-3) + + result = scaler.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "min_max_scaled_culmen_depth_mm": [0.678571, 0.4880952, 0.595238], + "min_max_scaled_culmen_length_mm": [0.269091, 0.232727, 0.210909], + "min_max_scaled_flipper_length_mm": [0.40678, 0.152542, 0.271186], + }, + dtype="Float64", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + def test_one_hot_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() encoder.fit(new_penguins_df[["species", "sex"]]) diff --git a/tests/unit/ml/test_compose.py b/tests/unit/ml/test_compose.py index 24cf0a333e7..8c8fbd6ab54 100644 --- a/tests/unit/ml/test_compose.py +++ b/tests/unit/ml/test_compose.py @@ -22,6 +22,7 @@ def test_columntransformer_init_expectedtransforms(): onehot_transformer = preprocessing.OneHotEncoder() standard_scaler_transformer = preprocessing.StandardScaler() max_abs_scaler_transformer = preprocessing.MaxAbsScaler() + min_max_scaler_transformer = preprocessing.MinMaxScaler() label_transformer = preprocessing.LabelEncoder() column_transformer = compose.ColumnTransformer( [ @@ -36,6 +37,11 @@ def test_columntransformer_init_expectedtransforms(): max_abs_scaler_transformer, ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + min_max_scaler_transformer, + ["culmen_length_mm", "flipper_length_mm"], + ), ("label", label_transformer, "species"), ] ) @@ -46,6 +52,8 @@ def test_columntransformer_init_expectedtransforms(): ("standard_scale", standard_scaler_transformer, "flipper_length_mm"), ("max_abs_scale", max_abs_scaler_transformer, "culmen_length_mm"), ("max_abs_scale", max_abs_scaler_transformer, "flipper_length_mm"), + ("min_max_scale", min_max_scaler_transformer, "culmen_length_mm"), + ("min_max_scale", min_max_scaler_transformer, "flipper_length_mm"), ("label", label_transformer, "species"), ] @@ -68,6 +76,11 @@ def test_columntransformer_repr(): preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + preprocessing.MinMaxScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) @@ -77,6 +90,8 @@ def test_columntransformer_repr(): ('standard_scale', StandardScaler(), ['culmen_length_mm', 'flipper_length_mm']), ('max_abs_scale', MaxAbsScaler(), + ['culmen_length_mm', 'flipper_length_mm']), + ('min_max_scale', MinMaxScaler(), ['culmen_length_mm', 'flipper_length_mm'])])""" ) @@ -99,6 +114,11 @@ def test_columntransformer_repr_matches_sklearn(): preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + preprocessing.MinMaxScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) sk_column_transformer = sklearn_compose.ColumnTransformer( @@ -118,6 +138,11 @@ def test_columntransformer_repr_matches_sklearn(): sklearn_preprocessing.MaxAbsScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "min_max_scale", + sklearn_preprocessing.MinMaxScaler(), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index b88523c7efa..f461dc76df8 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -87,6 +87,13 @@ def test_max_abs_scaler_produces_correct_sql( assert sql == "ML.MAX_ABS_SCALER(col_a) OVER() AS scaled_col_a" +def test_min_max_scaler_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_min_max_scaler("col_a", "scaled_col_a") + assert sql == "ML.MIN_MAX_SCALER(col_a) OVER() AS scaled_col_a" + + def test_one_hot_encoder_produces_correct_sql( base_sql_generator: ml_sql.BaseSqlGenerator, ): diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index 40b4f76ab7f..58e16e135b7 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -106,3 +106,39 @@ def transform(self, X): bigframes.dataframe.DataFrame: Transformed result. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + +class MinMaxScaler(BaseEstimator, TransformerMixin): + """Transform features by scaling each feature to a given range. + + This estimator scales and translates each feature individually such + that it is in the given range on the training set, e.g. between + zero and one. + """ + + def fit(self, X, y=None): + """Compute the minimum and maximum to be used for later scaling. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Dataframe or Series with training data. + + y (default None): + Ignored. + + Returns: + MaxAbsScaler: Fitted scaler. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transform(self, X): + """Scale the data. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. + + Returns: + bigframes.dataframe.DataFrame: Transformed result. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 61200bd3ae08ddafcc5e59131ac0295188e81f53 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Wed, 27 Sep 2023 11:14:19 -0700 Subject: [PATCH 42/80] refactor: push down SQL generate logic in core.BqmlModel (#66) --- bigframes/ml/core.py | 41 ++++++++++--------------------- bigframes/ml/sql.py | 43 ++++++++++++++++++++++++--------- tests/unit/ml/test_sql.py | 51 +++++++++++++++++++++++++++------------ 3 files changed, 80 insertions(+), 55 deletions(-) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 37478d8bafd..4c5a48cf625 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -58,7 +58,7 @@ def model(self) -> bigquery.Model: def _apply_sql( self, input_data: bpd.DataFrame, - func: Callable[[str], str], + func: Callable[[bpd.DataFrame], str], ) -> bpd.DataFrame: """Helper to wrap a dataframe in a SQL query, keeping the index intact. @@ -74,11 +74,9 @@ def _apply_sql( string from which to construct the output dataframe. It must include the index columns of the input SQL. """ - source_sql, index_col_ids, index_labels = input_data._to_sql_query( - include_index=True - ) + _, index_col_ids, index_labels = input_data._to_sql_query(include_index=True) - sql = func(source_sql) + sql = func(input_data) df = self._session.read_gbq(sql, index_col=index_col_ids) df.index.names = index_labels @@ -106,11 +104,9 @@ def generate_text( # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text( - source_sql=source_sql, - struct_options=self._model_manipulation_sql_generator.struct_options( - **options - ), + lambda source_df: self._model_manipulation_sql_generator.ml_generate_text( + source_df=source_df, + struct_options=options, ), ) @@ -122,11 +118,9 @@ def generate_text_embedding( # TODO: validate input data schema return self._apply_sql( input_data, - lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text_embedding( - source_sql=source_sql, - struct_options=self._model_manipulation_sql_generator.struct_options( - **options - ), + lambda source_df: self._model_manipulation_sql_generator.ml_generate_text_embedding( + source_df=source_df, + struct_options=options, ), ) @@ -136,13 +130,7 @@ def forecast(self) -> bpd.DataFrame: def evaluate(self, input_data: Optional[bpd.DataFrame] = None): # TODO: validate input data schema - # Note: don't need index as evaluate returns a new table - source_sql, _, _ = ( - input_data._to_sql_query(include_index=False) - if (input_data is not None) - else (None, None, None) - ) - sql = self._model_manipulation_sql_generator.ml_evaluate(source_sql) + sql = self._model_manipulation_sql_generator.ml_evaluate(input_data) return self._session.read_gbq(sql) @@ -188,11 +176,8 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: # truncate as Vertex ID only accepts 63 characters, easily exceeding the limit for temp models. # The possibility of conflicts should be low. vertex_ai_model_id = vertex_ai_model_id[:63] - options_sql = self._model_manipulation_sql_generator.options( - **{"vertex_ai_model_id": vertex_ai_model_id} - ) sql = self._model_manipulation_sql_generator.alter_model( - options_sql=options_sql + options={"vertex_ai_model_id": vertex_ai_model_id} ) # Register the model and wait it to finish self._session._start_query(sql) @@ -252,7 +237,7 @@ def create_model( session = X_train._session sql = self._model_creation_sql_generator.create_model( - source=input_data, + source_df=input_data, transforms=transforms, options=options, ) @@ -281,7 +266,7 @@ def create_time_series_model( session = X_train._session sql = self._model_creation_sql_generator.create_model( - source=input_data, + source_df=input_data, transforms=transforms, options=options, ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 3897d1be390..57c8ba672af 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -118,12 +118,12 @@ def __init__(self, model_id: str): # Model create and alter def create_model( self, - source: bpd.DataFrame, + source_df: bpd.DataFrame, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, transforms: Optional[Iterable[str]] = None, ) -> str: """Encode the CREATE TEMP MODEL statement for BQML""" - source_sql = source.sql + source_sql = source_df.sql transform_sql = self.transform(*transforms) if transforms is not None else None options_sql = self.options(**options) @@ -168,39 +168,58 @@ class ModelManipulationSqlGenerator(BaseSqlGenerator): def __init__(self, model_name: str): self._model_name = model_name + def _source_sql(self, source_df: bpd.DataFrame) -> str: + """Return DataFrame sql with index columns.""" + _source_sql, _, _ = source_df._to_sql_query(include_index=True) + return _source_sql + # Alter model def alter_model( self, - options_sql: str, + options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: """Encode the ALTER MODEL statement for BQML""" + options_sql = self.options(**options) + parts = [f"ALTER MODEL `{self._model_name}`"] parts.append(f"SET {options_sql}") return "\n".join(parts) # ML prediction TVFs - def ml_predict(self, source_sql: str) -> str: + def ml_predict(self, source_df: bpd.DataFrame) -> str: """Encode ML.PREDICT for BQML""" return f"""SELECT * FROM ML.PREDICT(MODEL `{self._model_name}`, - ({source_sql}))""" + ({self._source_sql(source_df)}))""" def ml_forecast(self) -> str: """Encode ML.FORECAST for BQML""" return f"""SELECT * FROM ML.FORECAST(MODEL `{self._model_name}`)""" - def ml_generate_text(self, source_sql: str, struct_options: str) -> str: + def ml_generate_text( + self, source_df: bpd.DataFrame, struct_options: Mapping[str, Union[int, float]] + ) -> str: """Encode ML.GENERATE_TEXT for BQML""" + struct_options_sql = self.struct_options(**struct_options) return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL `{self._model_name}`, - ({source_sql}), {struct_options})""" + ({self._source_sql(source_df)}), {struct_options_sql})""" - def ml_generate_text_embedding(self, source_sql: str, struct_options: str) -> str: + def ml_generate_text_embedding( + self, source_df: bpd.DataFrame, struct_options: Mapping[str, Union[int, float]] + ) -> str: """Encode ML.GENERATE_TEXT_EMBEDDING for BQML""" + struct_options_sql = self.struct_options(**struct_options) return f"""SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `{self._model_name}`, - ({source_sql}), {struct_options})""" + ({self._source_sql(source_df)}), {struct_options_sql})""" # ML evaluation TVFs - def ml_evaluate(self, source_sql: Optional[str] = None) -> str: + def ml_evaluate(self, source_df: Optional[bpd.DataFrame] = None) -> str: """Encode ML.EVALUATE for BQML""" + if source_df is None: + source_sql = None + else: + # Note: don't need index as evaluate returns a new table + source_sql, _, _ = source_df._to_sql_query(include_index=False) + if source_sql is None: return f"""SELECT * FROM ML.EVALUATE(MODEL `{self._model_name}`)""" else: @@ -222,7 +241,7 @@ def ml_principal_component_info(self) -> str: ) # ML transform TVF, that require a transform_only type model - def ml_transform(self, source_sql: str) -> str: + def ml_transform(self, source_df: bpd.DataFrame) -> str: """Encode ML.TRANSFORM for BQML""" return f"""SELECT * FROM ML.TRANSFORM(MODEL `{self._model_name}`, - ({source_sql}))""" + ({self._source_sql(source_df)}))""" diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index f461dc76df8..a3338e762d4 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -41,6 +41,7 @@ def model_manipulation_sql_generator() -> ml_sql.ModelManipulationSqlGenerator: def mock_df(): mock_df = mock.create_autospec(spec=bpd.DataFrame) mock_df.sql = "input_X_y_sql" + mock_df._to_sql_query.return_value = "input_X_sql", None, None return mock_df @@ -117,7 +118,7 @@ def test_create_model_produces_correct_sql( mock_df: bpd.DataFrame, ): sql = model_creation_sql_generator.create_model( - source=mock_df, + source_df=mock_df, options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( @@ -135,7 +136,7 @@ def test_create_model_transform_produces_correct_sql( mock_df: bpd.DataFrame, ): sql = model_creation_sql_generator.create_model( - source=mock_df, + source_df=mock_df, options={"option_key1": "option_value1", "option_key2": 2}, transforms=[ "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", @@ -191,38 +192,38 @@ def test_alter_model_correct_sql( model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, ): sql = model_manipulation_sql_generator.alter_model( - options_sql="my_options_sql", + options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql == """ALTER MODEL `my_project_id.my_dataset_id.my_model_id` -SET my_options_sql""" +SET OPTIONS( + option_key1="option_value1", + option_key2=2)""" ) def test_ml_predict_produces_correct_sql( model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, + mock_df: bpd.DataFrame, ): - sql = model_manipulation_sql_generator.ml_predict( - source_sql="SELECT * FROM my_table" - ) + sql = model_manipulation_sql_generator.ml_predict(source_df=mock_df) assert ( sql == """SELECT * FROM ML.PREDICT(MODEL `my_project_id.my_dataset_id.my_model_id`, - (SELECT * FROM my_table))""" + (input_X_sql))""" ) def test_ml_evaluate_produces_correct_sql( model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, + mock_df: bpd.DataFrame, ): - sql = model_manipulation_sql_generator.ml_evaluate( - source_sql="SELECT * FROM my_table" - ) + sql = model_manipulation_sql_generator.ml_evaluate(source_df=mock_df) assert ( sql == """SELECT * FROM ML.EVALUATE(MODEL `my_project_id.my_dataset_id.my_model_id`, - (SELECT * FROM my_table))""" + (input_X_sql))""" ) @@ -248,15 +249,35 @@ def test_ml_centroids_produces_correct_sql( def test_ml_generate_text_produces_correct_sql( model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, + mock_df: bpd.DataFrame, ): sql = model_manipulation_sql_generator.ml_generate_text( - source_sql="SELECT * FROM my_table", - struct_options="STRUCT(value AS item)", + source_df=mock_df, + struct_options={"option_key1": 1, "option_key2": 2.2}, ) assert ( sql == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project_id.my_dataset_id.my_model_id`, - (SELECT * FROM my_table), STRUCT(value AS item))""" + (input_X_sql), STRUCT( + 1 AS option_key1, + 2.2 AS option_key2))""" + ) + + +def test_ml_generate_text_embedding_produces_correct_sql( + model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator, + mock_df: bpd.DataFrame, +): + sql = model_manipulation_sql_generator.ml_generate_text_embedding( + source_df=mock_df, + struct_options={"option_key1": 1, "option_key2": 2.2}, + ) + assert ( + sql + == """SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `my_project_id.my_dataset_id.my_model_id`, + (input_X_sql), STRUCT( + 1 AS option_key1, + 2.2 AS option_key2))""" ) From 7ab65e88deb0080e9c36c2709f8a5385ccaf8cf2 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 27 Sep 2023 21:00:26 -0700 Subject: [PATCH 43/80] fix: generate unique ids on join to avoid id collisions (#65) * fix: generate unique ids on join to avoid id collisions --- bigframes/core/joins/single_column.py | 80 +++++++++++---------------- 1 file changed, 31 insertions(+), 49 deletions(-) diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 434cc2cd794..8a9825cf0b5 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -16,6 +16,7 @@ from __future__ import annotations +import itertools import typing from typing import Callable, Literal, Tuple @@ -25,7 +26,7 @@ import bigframes.constants as constants import bigframes.core as core -import bigframes.core.guid +import bigframes.core.guid as guid import bigframes.core.joins.row_identity import bigframes.core.ordering @@ -122,17 +123,38 @@ def join_by_column( ), ) else: + lmapping = { + col_id: guid.generate_guid() + for col_id in itertools.chain( + left.column_names, left._hidden_ordering_column_names + ) + } + rmapping = { + col_id: guid.generate_guid() + for col_id in itertools.chain( + right.column_names, right._hidden_ordering_column_names + ) + } + + def get_column_left(col_id): + return lmapping[col_id] + + def get_column_right(col_id): + return rmapping[col_id] + left_table = left._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True, + col_id_overrides=lmapping, ) right_table = right._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True, + col_id_overrides=rmapping, ) join_conditions = [ - value_to_join_key(left_table[left_index]) - == value_to_join_key(right_table[right_index]) + value_to_join_key(left_table[lmapping[left_index]]) + == value_to_join_key(right_table[rmapping[right_index]]) for left_index, right_index in zip(left_column_ids, right_column_ids) ] @@ -145,38 +167,6 @@ def join_by_column( rname="{name}_y", ) - def get_column_left(key: str) -> str: - if ( - how == "inner" - and key in left_column_ids - and key in combined_table.columns - ): - # Ibis doesn't rename the column if the values are guaranteed - # to be equal on left and right (because they're part of an - # inner join condition). See: - # https://github.com/ibis-project/ibis/pull/4651 - pass - elif key in right_table.columns: - key = f"{key}_x" - - return key - - def get_column_right(key: str) -> str: - if ( - how == "inner" - and key in right_column_ids - and key in combined_table.columns - ): - # Ibis doesn't rename the column if the values are guaranteed - # to be equal on left and right (because they're part of an - # inner join condition). See: - # https://github.com/ibis-project/ibis/pull/4651 - pass - elif key in left_table.columns: - key = f"{key}_y" - - return key - # Preserve ordering accross joins. ordering = join_orderings( left._ordering, @@ -245,20 +235,14 @@ def get_join_cols( join_key_cols: list[ibis_types.Value] = [] for left_col, right_col in zip(left_join_cols, right_join_cols): if not coalesce_join_keys: - join_key_cols.append( - left_col.name(bigframes.core.guid.generate_guid(prefix="index_")) - ) - join_key_cols.append( - right_col.name(bigframes.core.guid.generate_guid(prefix="index_")) - ) + join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) + join_key_cols.append(right_col.name(guid.generate_guid(prefix="index_"))) else: if how == "left" or how == "inner": - join_key_cols.append( - left_col.name(bigframes.core.guid.generate_guid(prefix="index_")) - ) + join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) elif how == "right": join_key_cols.append( - right_col.name(bigframes.core.guid.generate_guid(prefix="index_")) + right_col.name(guid.generate_guid(prefix="index_")) ) elif how == "outer": # The left index and the right index might contain null values, for @@ -269,16 +253,14 @@ def get_join_cols( # Don't need to coalesce if they are exactly the same column. if left_col.name("index").equals(right_col.name("index")): join_key_cols.append( - left_col.name( - bigframes.core.guid.generate_guid(prefix="index_") - ) + left_col.name(guid.generate_guid(prefix="index_")) ) else: join_key_cols.append( ibis.coalesce( left_col, right_col, - ).name(bigframes.core.guid.generate_guid(prefix="index_")) + ).name(guid.generate_guid(prefix="index_")) ) else: raise ValueError( From 0e0493f6099e4084d1978afb3f60fb3e0a872379 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Thu, 28 Sep 2023 14:42:06 -0500 Subject: [PATCH 44/80] chore(main): release 0.5.0 (#35) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 48 ++++++++++++++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de2edcf31e9..e4b2bff3c71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,54 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.5.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.4.0...v0.5.0) (2023-09-28) + + +### Features + +* Add `DataFrame.kurtosis` / `DF.kurt` method ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Add `DataFrame.rolling` and `DataFrame.expanding` methods ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Add `items`, `apply` methods to `DataFrame`. ([#43](https://github.com/googleapis/python-bigquery-dataframes/issues/43)) ([3adc1b3](https://github.com/googleapis/python-bigquery-dataframes/commit/3adc1b3aa3e2b218d4fa5debdaa4298276bdf801)) +* Add axis param to simple df aggregations ([#52](https://github.com/googleapis/python-bigquery-dataframes/issues/52)) ([9cf9972](https://github.com/googleapis/python-bigquery-dataframes/commit/9cf99721ed83704e6ee28b15c699326c431eb252)) +* Add index `dtype`, `astype`, `drop`, `fillna`, aggregate attributes. ([#38](https://github.com/googleapis/python-bigquery-dataframes/issues/38)) ([1a254a4](https://github.com/googleapis/python-bigquery-dataframes/commit/1a254a496633957b9506dd8392dcc6fd10762201)) +* Add ml.preprocessing.LabelEncoder ([#50](https://github.com/googleapis/python-bigquery-dataframes/issues/50)) ([2510461](https://github.com/googleapis/python-bigquery-dataframes/commit/25104610e5ffe526315923946533a66713c1d155)) +* Add ml.preprocessing.MaxAbsScaler ([#56](https://github.com/googleapis/python-bigquery-dataframes/issues/56)) ([14b262b](https://github.com/googleapis/python-bigquery-dataframes/commit/14b262bde2bb86093bf4df63862e369c5a84b0ad)) +* Add ml.preprocessing.MinMaxScaler ([#64](https://github.com/googleapis/python-bigquery-dataframes/issues/64)) ([392113b](https://github.com/googleapis/python-bigquery-dataframes/commit/392113b70d6a8c407accbb6684d75b31261e3741)) +* Add more index methods ([#54](https://github.com/googleapis/python-bigquery-dataframes/issues/54)) ([a6e32aa](https://github.com/googleapis/python-bigquery-dataframes/commit/a6e32aa875370063c48ce7922c2aa369a770bd30)) +* Support `calculate_p_values` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `class_weights="balanced"` in `LogisticRegression` model ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `df[column_name] = df_only_one_column` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `early_stop` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `enable_global_explain` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `l2_reg` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `learn_rate_strategy` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `ls_init_learn_rate` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `max_iterations` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `min_rel_progress` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support `optimize_strategy` parameter in `bigframes.ml.linear_model.LinearRegression` ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) +* Support casting string to integer or float ([#59](https://github.com/googleapis/python-bigquery-dataframes/issues/59)) ([3502f83](https://github.com/googleapis/python-bigquery-dataframes/commit/3502f835b35c437933430698e7a1c9badaddcb99)) + + +### Bug Fixes + +* Fix header skipping logic in `read_csv` ([#49](https://github.com/googleapis/python-bigquery-dataframes/issues/49)) ([d56258c](https://github.com/googleapis/python-bigquery-dataframes/commit/d56258cbfcda168cb9e437a021e282818d622d6a)) +* Generate unique ids on join to avoid id collisions ([#65](https://github.com/googleapis/python-bigquery-dataframes/issues/65)) ([7ab65e8](https://github.com/googleapis/python-bigquery-dataframes/commit/7ab65e88deb0080e9c36c2709f8a5385ccaf8cf2)) +* LabelEncoder params consistent with Sklearn ([#60](https://github.com/googleapis/python-bigquery-dataframes/issues/60)) ([632caec](https://github.com/googleapis/python-bigquery-dataframes/commit/632caec420a7e23188f01b96a00c354d205da74e)) +* Loosen filter items tests to accomodate shifting pandas impl ([#41](https://github.com/googleapis/python-bigquery-dataframes/issues/41)) ([edabdbb](https://github.com/googleapis/python-bigquery-dataframes/commit/edabdbb131150707ea9211292cacbb60b8d076dd)) + + +### Performance Improvements + +* Add ability to cache dataframe and series to session table ([#51](https://github.com/googleapis/python-bigquery-dataframes/issues/51)) ([416d7cb](https://github.com/googleapis/python-bigquery-dataframes/commit/416d7cb9b560d7e33dcc0227f03a00d43f55ba0d)) +* Inline small `Series` and `DataFrames` in query text ([#45](https://github.com/googleapis/python-bigquery-dataframes/issues/45)) ([5e199ec](https://github.com/googleapis/python-bigquery-dataframes/commit/5e199ecf1ecf13a68a2ed0dd4464afd9db977ab1)) +* Reimplement unpivot to use cross join rather than union ([#47](https://github.com/googleapis/python-bigquery-dataframes/issues/47)) ([f9a93ce](https://github.com/googleapis/python-bigquery-dataframes/commit/f9a93ce71d053aa17b1e3a2946c90e0227076184)) +* Simplify join order to use multiple order keys instead of string. ([#36](https://github.com/googleapis/python-bigquery-dataframes/issues/36)) ([5056da6](https://github.com/googleapis/python-bigquery-dataframes/commit/5056da6b385dbcfc179d2bcbb6549fa539428cda)) + + +### Documentation + +* Link to Remote Functions code samples from README and API reference ([c1900c2](https://github.com/googleapis/python-bigquery-dataframes/commit/c1900c29a44199d5d8d036d6d842b4f00448fa79)) + ## [0.4.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.3.2...v0.4.0) (2023-09-16) diff --git a/bigframes/version.py b/bigframes/version.py index 65b984a0d73..ad3c3082c52 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.4.0" +__version__ = "0.5.0" From bc7be7f3529d39e647913824ca8bb86446da941c Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 28 Sep 2023 15:34:21 -0500 Subject: [PATCH 45/80] chore: remove Git-on-Borg check from config (#69) * chore: remove Git-on-Borg check from config Continuous jobs are incorrectly identifying themselves as Git-on-Borg jobs and using the wrong job path. * update job configs too --- .kokoro/build.sh | 6 +----- .kokoro/continuous/common.cfg | 2 +- .kokoro/continuous/nightly.cfg | 2 +- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/.kokoro/build.sh b/.kokoro/build.sh index f91c541c6c9..a0fa4bc787b 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -15,11 +15,7 @@ set -eo pipefail -if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then - PROJECT_SCM="github/python-bigquery-dataframes" -else - PROJECT_SCM="git/bigframes" -fi +PROJECT_SCM="github/python-bigquery-dataframes" if [[ -z "${PROJECT_ROOT:-}" ]]; then PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}" diff --git a/.kokoro/continuous/common.cfg b/.kokoro/continuous/common.cfg index 5d40578ac79..97e0651aa92 100644 --- a/.kokoro/continuous/common.cfg +++ b/.kokoro/continuous/common.cfg @@ -7,4 +7,4 @@ action { } } -build_file: "bigframes/.kokoro/build.sh" +build_file: "python-bigquery-dataframes/.kokoro/build.sh" diff --git a/.kokoro/continuous/nightly.cfg b/.kokoro/continuous/nightly.cfg index 63c3f51d053..2b7111664f7 100644 --- a/.kokoro/continuous/nightly.cfg +++ b/.kokoro/continuous/nightly.cfg @@ -1,3 +1,3 @@ # Format: //devtools/kokoro/config/proto/build.proto -build_file: "bigframes/.kokoro/release-nightly.sh" +build_file: "python-bigquery-dataframes/.kokoro/release-nightly.sh" From bf050cf475ad8a9e3e0ca3f896ddaf96dbe13ae3 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 28 Sep 2023 15:51:25 -0700 Subject: [PATCH 46/80] feat: add update and align methods to dataframe (#57) * feat: add update and align methods to dataframe --- bigframes/core/block_transforms.py | 72 ++++++++++++ bigframes/dataframe.py | 108 +++++++++++------- tests/system/small/test_dataframe.py | 76 +++++++++++- .../bigframes_vendored/pandas/core/frame.py | 62 ++++++++++ 4 files changed, 271 insertions(+), 47 deletions(-) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index d22112417cc..6e99a7c7740 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -504,3 +504,75 @@ def _kurt_from_moments_and_count( kurt_id, na_cond_id, ops.partial_arg3(ops.where_op, None) ) return block, kurt_id + + +def align( + left_block: blocks.Block, + right_block: blocks.Block, + join: str = "outer", + axis: typing.Union[str, int, None] = None, +) -> typing.Tuple[blocks.Block, blocks.Block]: + axis_n = core.utils.get_axis_number(axis) if axis is not None else None + # Must align columns first as other way will likely create extra joins + if (axis_n is None) or axis_n == 1: + left_block, right_block = align_columns(left_block, right_block, join=join) + if (axis_n is None) or axis_n == 0: + left_block, right_block = align_rows(left_block, right_block, join=join) + return left_block, right_block + + +def align_rows( + left_block: blocks.Block, + right_block: blocks.Block, + join: str = "outer", +): + joined_index, (get_column_left, get_column_right) = left_block.index.join( + right_block.index, how=join + ) + left_columns = [get_column_left(col) for col in left_block.value_columns] + right_columns = [get_column_right(col) for col in right_block.value_columns] + + left_block = joined_index._block.select_columns(left_columns) + right_block = joined_index._block.select_columns(right_columns) + return left_block, right_block + + +def align_columns( + left_block: blocks.Block, + right_block: blocks.Block, + join: str = "outer", +): + columns, lcol_indexer, rcol_indexer = left_block.column_labels.join( + right_block.column_labels, how=join, return_indexers=True + ) + column_indices = zip( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)), + rcol_indexer if (rcol_indexer is not None) else range(len(columns)), + ) + left_column_ids = [] + right_column_ids = [] + + original_left_block = left_block + original_right_block = right_block + + for left_index, right_index in column_indices: + if left_index >= 0: + left_col_id = original_left_block.value_columns[left_index] + else: + dtype = right_block.dtypes[right_index] + left_block, left_col_id = left_block.create_constant( + None, dtype=dtype, label=original_right_block.column_labels[right_index] + ) + left_column_ids.append(left_col_id) + + if right_index >= 0: + right_col_id = original_right_block.value_columns[right_index] + else: + dtype = original_left_block.dtypes[left_index] + right_block, right_col_id = right_block.create_constant( + None, dtype=dtype, label=left_block.column_labels[left_index] + ) + right_column_ids.append(right_col_id) + left_final = left_block.select_columns(left_column_ids) + right_final = right_block.select_columns(right_column_ids) + return left_final, right_final diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 113355589b1..828d2df64ef 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -745,6 +745,55 @@ def rpow( __rpow__ = rpow + def align( + self, + other: typing.Union[DataFrame, bigframes.series.Series], + join: str = "outer", + axis: typing.Union[str, int, None] = None, + ) -> typing.Tuple[ + typing.Union[DataFrame, bigframes.series.Series], + typing.Union[DataFrame, bigframes.series.Series], + ]: + axis_n = utils.get_axis_number(axis) if axis else None + if axis_n == 1 and isinstance(other, bigframes.series.Series): + raise NotImplementedError( + f"align with series and axis=1 not supported. {constants.FEEDBACK_LINK}" + ) + left_block, right_block = block_ops.align( + self._block, other._block, join=join, axis=axis + ) + return DataFrame(left_block), other.__class__(right_block) + + def update(self, other, join: str = "left", overwrite=True, filter_func=None): + other = other if isinstance(other, DataFrame) else DataFrame(other) + if join != "left": + raise ValueError("Only 'left' join supported for update") + + if filter_func is not None: # Will always take other if possible + + def update_func( + left: bigframes.series.Series, right: bigframes.series.Series + ) -> bigframes.series.Series: + return left.mask(right.notna() & filter_func(left), right) + + elif overwrite: + + def update_func( + left: bigframes.series.Series, right: bigframes.series.Series + ) -> bigframes.series.Series: + return left.mask(right.notna(), right) + + else: + + def update_func( + left: bigframes.series.Series, right: bigframes.series.Series + ) -> bigframes.series.Series: + return left.mask(left.isna(), right) + + result = self.combine(other, update_func, how=join) + + self._set_block(result._block) + def combine( self, other: DataFrame, @@ -753,56 +802,31 @@ def combine( ], fill_value=None, overwrite: bool = True, + *, + how: str = "outer", ) -> DataFrame: - # Join rows - joined_index, (get_column_left, get_column_right) = self._block.index.join( - other._block.index, how="outer" - ) - columns, lcol_indexer, rcol_indexer = self.columns.join( - other.columns, how="outer", return_indexers=True - ) + l_aligned, r_aligned = block_ops.align(self._block, other._block, join=how) - column_indices = zip( - lcol_indexer if (lcol_indexer is not None) else range(len(columns)), - rcol_indexer if (lcol_indexer is not None) else range(len(columns)), + other_missing_labels = self._block.column_labels.difference( + other._block.column_labels ) - block = joined_index._block + l_frame = DataFrame(l_aligned) + r_frame = DataFrame(r_aligned) results = [] - for left_index, right_index in column_indices: - if left_index >= 0 and right_index >= 0: # -1 indices indicate missing - left_col_id = get_column_left(self._block.value_columns[left_index]) - right_col_id = get_column_right(other._block.value_columns[right_index]) - left_series = bigframes.series.Series(block.select_column(left_col_id)) - right_series = bigframes.series.Series( - block.select_column(right_col_id) - ) + for (label, lseries), (_, rseries) in zip(l_frame.items(), r_frame.items()): + if not ((label in other_missing_labels) and not overwrite): if fill_value is not None: - left_series = left_series.fillna(fill_value) - right_series = right_series.fillna(fill_value) - results.append(func(left_series, right_series)) - elif left_index >= 0: - # Does not exist in other - if overwrite: - dtype = self.dtypes[left_index] - block, null_col_id = block.create_constant(None, dtype=dtype) - result = bigframes.series.Series(block.select_column(null_col_id)) - results.append(result) + result = func( + lseries.fillna(fill_value), rseries.fillna(fill_value) + ) else: - left_col_id = get_column_left(self._block.value_columns[left_index]) - result = bigframes.series.Series(block.select_column(left_col_id)) - if fill_value is not None: - result = result.fillna(fill_value) - results.append(result) - elif right_index >= 0: - right_col_id = get_column_right(other._block.value_columns[right_index]) - result = bigframes.series.Series(block.select_column(right_col_id)) - if fill_value is not None: - result = result.fillna(fill_value) - results.append(result) + result = func(lseries, rseries) else: - # Should not be possible - raise ValueError("No right or left index.") + result = ( + lseries.fillna(fill_value) if fill_value is not None else lseries + ) + results.append(result) if all([isinstance(val, bigframes.series.Series) for val in results]): import bigframes.core.reshape as rs diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index adf17848ee6..ba76c4b0d3e 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1211,6 +1211,77 @@ def test_combine( pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) +@pytest.mark.parametrize( + ("overwrite", "filter_func"), + [ + (True, None), + (False, None), + (True, lambda x: x.isna() | (x % 2 == 0)), + ], + ids=[ + "default", + "overwritefalse", + "customfilter", + ], +) +def test_df_update(overwrite, filter_func): + if pd.__version__.startswith("1."): + pytest.skip("dtype handled differently in pandas 1.x.") + index1 = pandas.Index([1, 2, 3, 4], dtype="Int64") + index2 = pandas.Index([1, 2, 4, 5], dtype="Int64") + pd_df1 = pandas.DataFrame( + {"a": [1, None, 3, 4], "b": [5, 6, None, 8]}, dtype="Int64", index=index1 + ) + pd_df2 = pandas.DataFrame( + {"a": [None, 20, 30, 40], "c": [90, None, 110, 120]}, + dtype="Int64", + index=index2, + ) + + bf_df1 = dataframe.DataFrame(pd_df1) + bf_df2 = dataframe.DataFrame(pd_df2) + + bf_df1.update(bf_df2, overwrite=overwrite, filter_func=filter_func) + pd_df1.update(pd_df2, overwrite=overwrite, filter_func=filter_func) + + pd.testing.assert_frame_equal(bf_df1.to_pandas(), pd_df1) + + +@pytest.mark.parametrize( + ("join", "axis"), + [ + ("outer", None), + ("outer", 0), + ("outer", 1), + ("left", 0), + ("right", 1), + ("inner", None), + ("inner", 1), + ], +) +def test_df_align(join, axis): + index1 = pandas.Index([1, 2, 3, 4], dtype="Int64") + index2 = pandas.Index([1, 2, 4, 5], dtype="Int64") + pd_df1 = pandas.DataFrame( + {"a": [1, None, 3, 4], "b": [5, 6, None, 8]}, dtype="Int64", index=index1 + ) + pd_df2 = pandas.DataFrame( + {"a": [None, 20, 30, 40], "c": [90, None, 110, 120]}, + dtype="Int64", + index=index2, + ) + + bf_df1 = dataframe.DataFrame(pd_df1) + bf_df2 = dataframe.DataFrame(pd_df2) + + bf_result1, bf_result2 = bf_df1.align(bf_df2, join=join, axis=axis) + pd_result1, pd_result2 = pd_df1.align(pd_df2, join=join, axis=axis) + + # Don't check dtype as pandas does unnecessary float conversion + pd.testing.assert_frame_equal(bf_result1.to_pandas(), pd_result1, check_dtype=False) + pd.testing.assert_frame_equal(bf_result2.to_pandas(), pd_result2, check_dtype=False) + + def test_combine_first( scalars_df_index, scalars_df_2_index, @@ -1232,11 +1303,6 @@ def test_combine_first( pd_df_b.columns = ["b", "a", "d"] pd_result = pd_df_a.combine_first(pd_df_b) - print("pandas") - print(pd_result.to_string()) - print("bigframes") - print(bf_result.to_string()) - # Some dtype inconsistency for all-NULL columns pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 6ce11cd7e98..5cd9fe51637 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -503,6 +503,35 @@ def drop( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def align( + self, + other, + join="outer", + axis=None, + ) -> tuple: + """ + Align two objects on their axes with the specified join method. + + Join method is specified for each axis Index. + + Args: + other (DataFrame or Series): + join ({{'outer', 'inner', 'left', 'right'}}, default 'outer'): + Type of alignment to be performed. + left: use only keys from left frame, preserve key order. + right: use only keys from right frame, preserve key order. + outer: use union of keys from both frames, sort keys lexicographically. + inner: use intersection of keys from both frames, + preserve the order of the left keys. + + axis (allowed axis of the other object, default None): + Align on index (0), columns (1), or both (None). + + Returns: + tuple of (DataFrame, type of other): Aligned objects. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def rename( self, *, @@ -1265,6 +1294,39 @@ def combine_first(self, other) -> DataFrame: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def update( + self, other, join: str = "left", overwrite: bool = True, filter_func=None + ) -> DataFrame: + """ + Modify in place using non-NA values from another DataFrame. + + Aligns on indices. There is no return value. + + Args: + other (DataFrame, or object coercible into a DataFrame): + Should have at least one matching index/column label + with the original DataFrame. If a Series is passed, + its name attribute must be set, and that will be + used as the column name to align with the original DataFrame. + join ({'left'}, default 'left'): + Only left join is implemented, keeping the index and columns of the + original object. + overwrite (bool, default True): + How to handle non-NA values for overlapping keys: + True: overwrite original DataFrame's values + with values from `other`. + False: only update values that are NA in + the original DataFrame. + + filter_func (callable(1d-array) -> bool 1d-array, optional): + Can choose to replace values other than NA. Return True for values + that should be updated. + + Returns: + None: This method directly changes calling object. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Data reshaping From d8910d4d9335626d25891fb150d646bdb3527564 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 29 Sep 2023 10:03:21 -0700 Subject: [PATCH 47/80] chore: add vertex sdkhand bigframes integration notebooks (#70) --- .../vertex_sdk/sdk2_bigframes_pytorch.ipynb | 723 +++++++++++++++++ .../vertex_sdk/sdk2_bigframes_sklearn.ipynb | 727 ++++++++++++++++++ .../sdk2_bigframes_tensorflow.ipynb | 646 ++++++++++++++++ noxfile.py | 3 + 4 files changed, 2099 insertions(+) create mode 100644 notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb create mode 100644 notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb create mode 100644 notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb diff --git a/notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb b/notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb new file mode 100644 index 00000000000..598d958f0c3 --- /dev/null +++ b/notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb @@ -0,0 +1,723 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "# Train a pytorch model with Vertex AI SDK 2.0 and Bigframes\n", + "\n", + "\n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"VertexOpen in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "This tutorial demonstrates how to train a pytorch model using Vertex AI local-to-remote training with Vertex AI SDK 2.0 and BigQuery Bigframes as the data source.\n", + "\n", + "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn to use `Vertex AI SDK 2.0` with Bigframes as input data source.\n", + "\n", + "\n", + "This tutorial uses the following Google Cloud ML services:\n", + "\n", + "- `Vertex AI Training`\n", + "- `Vertex AI Remote Training`\n", + "\n", + "\n", + "The steps performed include:\n", + "\n", + "- Initialize a dataframe from a BigQuery table and split the dataset\n", + "- Perform transformations as a Vertex AI remote training.\n", + "- Train the model remotely and evaluate the model locally\n", + "\n", + "**Local-to-remote training**\n", + "\n", + "```\n", + "import vertexai\n", + "from my_module import MyModelClass\n", + "\n", + "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n", + "\n", + "# Wrap the model class with `vertex_ai.preview.remote`\n", + "MyModelClass = vertexai.preview.remote(MyModelClass)\n", + "\n", + "# Instantiate the class\n", + "model = MyModelClass(...)\n", + "\n", + "# Optional set remote config\n", + "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n", + "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n", + "\n", + "# This `fit` call will be executed remotely\n", + "model.fit(...)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses the IRIS dataset, which predicts the iris species." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* Vertex AI\n", + "* BigQuery\n", + "* Cloud Storage\n", + "\n", + "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n", + "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n", + "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages required to execute this notebook. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "# Install the packages\n", + "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n", + "! pip3 install --upgrade --quiet bigframes\n", + "! pip3 install --upgrade --quiet torch" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", + "\n", + "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Region\n", + "\n", + "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "region" + }, + "outputs": [], + "source": [ + "REGION = \"us-central1\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**1. Vertex AI Workbench**\n", + "* Do nothing as you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**2. Local JupyterLab instance, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**3. Colab, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f6b2ccc891ed" + }, + "source": [ + "**4. Service account or other**\n", + "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zgPO1eR3CYjk" + }, + "source": [ + "### Create a Cloud Storage bucket\n", + "\n", + "Create a storage bucket to store intermediate artifacts such as datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MzGDU7TWdts_" + }, + "outputs": [], + "source": [ + "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-EcIXiGsCePi" + }, + "source": [ + "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NIq7R4HZCfIc" + }, + "outputs": [], + "source": [ + "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries and define constants" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "import torch\n", + "import vertexai\n", + "from vertexai.preview import VertexModel\n", + "\n", + "bf.options.bigquery.location = \"us\" # Dataset is in 'us' not 'us-central1'\n", + "bf.options.bigquery.project = PROJECT_ID\n", + "\n", + "from bigframes.ml.model_selection import \\\n", + " train_test_split as bf_train_test_split" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "## Initialize Vertex AI SDK for Python\n", + "\n", + "Initialize the Vertex AI SDK for Python for your project and corresponding bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "outputs": [], + "source": [ + "vertexai.init(\n", + " project=PROJECT_ID,\n", + " location=REGION,\n", + " staging_bucket=BUCKET_URI,\n", + ")\n", + "\n", + "REMOTE_JOB_NAME = \"sdk2-bigframes-pytorch\"\n", + "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "105334524e96" + }, + "source": [ + "## Prepare the dataset\n", + "\n", + "Now load the Iris dataset and split the data into train and test sets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b44cdc4e03f1" + }, + "outputs": [], + "source": [ + "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n", + "\n", + "species_categories = {\n", + " \"versicolor\": 0,\n", + " \"virginica\": 1,\n", + " \"setosa\": 2,\n", + "}\n", + "df[\"species\"] = df[\"species\"].map(species_categories)\n", + "\n", + "# Assign an index column name\n", + "index_col = \"index\"\n", + "df.index.name = index_col" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9cb8616b1997" + }, + "outputs": [], + "source": [ + "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n", + "label_columns = df[[\"species\"]]\n", + "train_X, test_X, train_y, test_y = bf_train_test_split(\n", + " feature_columns, label_columns, test_size=0.2\n", + ")\n", + "\n", + "print(\"X_train size: \", train_X.size)\n", + "print(\"X_test size: \", test_X.size)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "23fe7b734b08" + }, + "outputs": [], + "source": [ + "# Switch to remote mode for training\n", + "vertexai.preview.init(remote=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5904a0f1bb03" + }, + "source": [ + "## PyTorch remote training with CPU (Custom PyTorch model)\n", + "\n", + "First, train a PyTorch model as a remote training job:\n", + "\n", + "- Reinitialize Vertex AI for remote training.\n", + "- Set TorchLogisticRegression for the remote training job.\n", + "- Invoke TorchLogisticRegression locally which will launch the remote training job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2a1b85195a17" + }, + "outputs": [], + "source": [ + "# define the custom model\n", + "class TorchLogisticRegression(VertexModel, torch.nn.Module):\n", + " def __init__(self, input_size: int, output_size: int):\n", + " torch.nn.Module.__init__(self)\n", + " VertexModel.__init__(self)\n", + " self.linear = torch.nn.Linear(input_size, output_size)\n", + " self.softmax = torch.nn.Softmax(dim=1)\n", + "\n", + " def forward(self, x):\n", + " return self.softmax(self.linear(x))\n", + "\n", + " @vertexai.preview.developer.mark.train()\n", + " def train(self, X, y, num_epochs, lr):\n", + " X = X.to(torch.float32)\n", + " y = torch.flatten(y) # necessary to get 1D tensor\n", + " dataloader = torch.utils.data.DataLoader(\n", + " torch.utils.data.TensorDataset(X, y),\n", + " batch_size=10,\n", + " shuffle=True,\n", + " generator=torch.Generator(device=X.device),\n", + " )\n", + "\n", + " criterion = torch.nn.CrossEntropyLoss()\n", + " optimizer = torch.optim.SGD(self.parameters(), lr=lr)\n", + "\n", + " for t in range(num_epochs):\n", + " for batch, (X, y) in enumerate(dataloader):\n", + " optimizer.zero_grad()\n", + " pred = self(X)\n", + " loss = criterion(pred, y)\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " @vertexai.preview.developer.mark.predict()\n", + " def predict(self, X):\n", + " X = torch.tensor(X).to(torch.float32)\n", + " with torch.no_grad():\n", + " pred = torch.argmax(self(X), dim=1)\n", + " return pred" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4e35593f520a" + }, + "outputs": [], + "source": [ + "# Switch to remote mode for training\n", + "vertexai.preview.init(remote=True)\n", + "\n", + "# Instantiate model\n", + "model = TorchLogisticRegression(4, 3)\n", + "\n", + "# Set training config\n", + "model.train.vertex.remote_config.custom_commands = [\n", + " \"pip install torchdata\",\n", + " \"pip install torcharrow\",\n", + "]\n", + "model.train.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-torch-model\"\n", + "model.train.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "\n", + "# Train model on Vertex\n", + "model.train(train_X, train_y, num_epochs=200, lr=0.05)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "edf4d0708f02" + }, + "source": [ + "## Remote prediction\n", + "\n", + "Obtain predictions from the trained model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "42dfbff0ca15" + }, + "outputs": [], + "source": [ + "vertexai.preview.init(remote=True)\n", + "\n", + "# Set remote config\n", + "model.predict.vertex.remote_config.custom_commands = [\n", + " \"pip install torchdata\",\n", + " \"pip install torcharrow\",\n", + "]\n", + "model.predict.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-torch-predict\"\n", + "model.predict.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "\n", + "predictions = model.predict(test_X)\n", + "\n", + "print(f\"Remote predictions: {predictions}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "4340ed8316cd" + }, + "source": [ + "## Local evaluation\n", + "\n", + "Evaluate model results locally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "eb27a31cec6f" + }, + "outputs": [], + "source": [ + "# User must convert bigframes to torch tensor for local evaluation\n", + "train_X_tensor = torch.from_numpy(\n", + " train_X.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n", + ")\n", + "train_y_tensor = torch.from_numpy(\n", + " train_y.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n", + ")\n", + "\n", + "test_X_tensor = torch.from_numpy(\n", + " test_X.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n", + ")\n", + "test_y_tensor = torch.from_numpy(\n", + " test_y.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "7db44ad81389" + }, + "outputs": [], + "source": [ + "from sklearn.metrics import accuracy_score\n", + "\n", + "# Switch to local mode for evaluation\n", + "vertexai.preview.init(remote=False)\n", + "\n", + "# Evaluate model's accuracy score\n", + "print(\n", + " f\"Train accuracy: {accuracy_score(train_y_tensor, model.predict(train_X_tensor))}\"\n", + ")\n", + "\n", + "print(f\"Test accuracy: {accuracy_score(test_y_tensor, model.predict(test_X_tensor))}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Delete Cloud Storage objects that were created\n", + "delete_bucket = False\n", + "if delete_bucket or os.getenv(\"IS_TESTING\"):\n", + " ! gsutil -m rm -r $BUCKET_URI" + ] + } + ], + "metadata": { + "colab": { + "collapsed_sections": [], + "name": "sdk2_bigframes_pytorch.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb b/notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb new file mode 100644 index 00000000000..021c0707535 --- /dev/null +++ b/notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb @@ -0,0 +1,727 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "# Train a scikit-learn model with Vertex AI SDK 2.0 and Bigframes\n", + "\n", + "\n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"VertexOpen in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "This tutorial demonstrates how to train a scikit-learn model using Vertex AI local-to-remote training with Vertex AI SDK 2.0 and BigQuery Bigframes as the data source.\n", + "\n", + "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn to use `Vertex AI SDK 2.0` with Bigframes as input data source.\n", + "\n", + "\n", + "This tutorial uses the following Google Cloud ML services:\n", + "\n", + "- `Vertex AI Training`\n", + "- `Vertex AI Remote Training`\n", + "\n", + "\n", + "The steps performed include:\n", + "\n", + "- Initialize a dataframe from a BigQuery table and split the dataset\n", + "- Perform transformations as a Vertex AI remote training.\n", + "- Train the model remotely and evaluate the model locally\n", + "\n", + "**Local-to-remote training**\n", + "\n", + "```\n", + "import vertexai\n", + "from my_module import MyModelClass\n", + "\n", + "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n", + "\n", + "# Wrap the model class with `vertex_ai.preview.remote`\n", + "MyModelClass = vertexai.preview.remote(MyModelClass)\n", + "\n", + "# Instantiate the class\n", + "model = MyModelClass(...)\n", + "\n", + "# Optional set remote config\n", + "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n", + "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n", + "\n", + "# This `fit` call will be executed remotely\n", + "model.fit(...)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses the IRIS dataset, which predicts the iris species." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* Vertex AI\n", + "* BigQuery\n", + "* Cloud Storage\n", + "\n", + "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n", + "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n", + "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages required to execute this notebook. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "# Install the packages\n", + "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n", + "! pip3 install --upgrade --quiet bigframes" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", + "\n", + "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Region\n", + "\n", + "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "region" + }, + "outputs": [], + "source": [ + "REGION = \"us-central1\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**1. Vertex AI Workbench**\n", + "* Do nothing as you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**2. Local JupyterLab instance, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**3. Colab, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f6b2ccc891ed" + }, + "source": [ + "**4. Service account or other**\n", + "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zgPO1eR3CYjk" + }, + "source": [ + "### Create a Cloud Storage bucket\n", + "\n", + "Create a storage bucket to store intermediate artifacts such as datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MzGDU7TWdts_" + }, + "outputs": [], + "source": [ + "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-EcIXiGsCePi" + }, + "source": [ + "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NIq7R4HZCfIc" + }, + "outputs": [], + "source": [ + "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries and define constants" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "import vertexai\n", + "\n", + "bf.options.bigquery.location = \"us\" # Dataset is in 'us' not 'us-central1'\n", + "bf.options.bigquery.project = PROJECT_ID\n", + "\n", + "from bigframes.ml.model_selection import \\\n", + " train_test_split as bf_train_test_split\n", + "\n", + "REMOTE_JOB_NAME = \"sdk2-bigframes-sklearn\"\n", + "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "## Initialize Vertex AI SDK for Python\n", + "\n", + "Initialize the Vertex AI SDK for Python for your project and corresponding bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "outputs": [], + "source": [ + "vertexai.init(\n", + " project=PROJECT_ID,\n", + " location=REGION,\n", + " staging_bucket=BUCKET_URI,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "105334524e96" + }, + "source": [ + "## Prepare the dataset\n", + "\n", + "Now load the Iris dataset and split the data into train and test sets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b44cdc4e03f1" + }, + "outputs": [], + "source": [ + "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n", + "\n", + "species_categories = {\n", + " \"versicolor\": 0,\n", + " \"virginica\": 1,\n", + " \"setosa\": 2,\n", + "}\n", + "df[\"species\"] = df[\"species\"].map(species_categories)\n", + "\n", + "# Assign an index column name\n", + "index_col = \"index\"\n", + "df.index.name = index_col" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9cb8616b1997" + }, + "outputs": [], + "source": [ + "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n", + "label_columns = df[[\"species\"]]\n", + "train_X, test_X, train_y, test_y = bf_train_test_split(\n", + " feature_columns, label_columns, test_size=0.2\n", + ")\n", + "\n", + "print(\"X_train size: \", train_X.size)\n", + "print(\"X_test size: \", test_X.size)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8306545fcc57" + }, + "source": [ + "## Feature transformation\n", + "\n", + "Next, you do feature transformations on the data using the Vertex AI remote training service.\n", + "\n", + "First, you re-initialize Vertex AI to enable remote training." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "55e701c31036" + }, + "outputs": [], + "source": [ + "# Switch to remote mode for training\n", + "vertexai.preview.init(remote=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "4a0e9d59b273" + }, + "source": [ + "### Execute remote job for fit_transform() on training data\n", + "\n", + "Next, indicate that the `StandardScalar` class is to be executed remotely. Then set up the data transform and call the `fit_transform()` method is executed remotely." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "90333089d362" + }, + "outputs": [], + "source": [ + "from sklearn.preprocessing import StandardScaler\n", + "\n", + "# Wrap classes to enable Vertex remote execution\n", + "StandardScaler = vertexai.preview.remote(StandardScaler)\n", + "\n", + "# Instantiate transformer\n", + "transformer = StandardScaler()\n", + "\n", + "# Set training config\n", + "transformer.fit_transform.vertex.remote_config.display_name = (\n", + " f\"{REMOTE_JOB_NAME}-fit-transformer-bigframes\"\n", + ")\n", + "transformer.fit_transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "\n", + "# Execute transformer on Vertex (train_X is bigframes.dataframe.DataFrame, X_train is np.array)\n", + "X_train = transformer.fit_transform(train_X)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6bf95574c907" + }, + "source": [ + "### Remote transform on test data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "da6eea22a89a" + }, + "outputs": [], + "source": [ + "# Transform test dataset before calculate test score\n", + "transformer.transform.vertex.remote_config.display_name = (\n", + " REMOTE_JOB_NAME + \"-transformer\"\n", + ")\n", + "transformer.transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "\n", + "# Execute transformer on Vertex (test_X is bigframes.dataframe.DataFrame, X_test is np.array)\n", + "X_test = transformer.transform(test_X)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ddf906c886e4" + }, + "source": [ + "## Remote training\n", + "\n", + "First, train the scikit-learn model as a remote training job:\n", + "\n", + "- Set LogisticRegression for the remote training job.\n", + "- Invoke LogisticRegression locally which will launch the remote training job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "c7b0116fa60c" + }, + "outputs": [], + "source": [ + "from sklearn.linear_model import LogisticRegression\n", + "\n", + "# Wrap classes to enable Vertex remote execution\n", + "LogisticRegression = vertexai.preview.remote(LogisticRegression)\n", + "\n", + "# Instantiate model, warm_start=True for uptraining\n", + "model = LogisticRegression(warm_start=True)\n", + "\n", + "# Set training config\n", + "model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-sklearn-model\"\n", + "model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "\n", + "# Train model on Vertex\n", + "model.fit(train_X, train_y)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ffe1d5903bcb" + }, + "source": [ + "## Remote prediction\n", + "\n", + "Obtain predictions from the trained model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d00ce35920fa" + }, + "outputs": [], + "source": [ + "# Remote evaluation\n", + "vertexai.preview.init(remote=True)\n", + "\n", + "# Evaluate model's accuracy score\n", + "predictions = model.predict(test_X)\n", + "\n", + "print(f\"Remote predictions: {predictions}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a8cd6cbd4403" + }, + "source": [ + "## Local evaluation\n", + "\n", + "Score model results locally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dc105dafdfb9" + }, + "outputs": [], + "source": [ + "# User must convert bigframes to pandas dataframe for local evaluation\n", + "train_X_pd = train_X.to_pandas().reset_index(drop=True)\n", + "train_y_pd = train_y.to_pandas().reset_index(drop=True)\n", + "\n", + "test_X_pd = test_X.to_pandas().reset_index(drop=True)\n", + "test_y_pd = test_y.to_pandas().reset_index(drop=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "25fec549de69" + }, + "outputs": [], + "source": [ + "# Switch to local mode for testing\n", + "vertexai.preview.init(remote=False)\n", + "\n", + "# Evaluate model's accuracy score\n", + "print(f\"Train accuracy: {model.score(train_X_pd, train_y_pd)}\")\n", + "\n", + "print(f\"Test accuracy: {model.score(test_X_pd, test_y_pd)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Delete Cloud Storage objects that were created\n", + "delete_bucket = False\n", + "if delete_bucket or os.getenv(\"IS_TESTING\"):\n", + " ! gsutil -m rm -r $BUCKET_URI" + ] + } + ], + "metadata": { + "colab": { + "collapsed_sections": [], + "name": "sdk2_bigframes_sklearn.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb b/notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb new file mode 100644 index 00000000000..e6843b66b57 --- /dev/null +++ b/notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb @@ -0,0 +1,646 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "# Train a Tensorflow Keras model with Vertex AI SDK 2.0 and Bigframes \n", + "\n", + "\n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"VertexOpen in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "This tutorial demonstrates how to train a tensorflow keras model using Vertex AI local-to-remote training with Vertex AI SDK 2.0 and BigQuery Bigframes as the data source.\n", + "\n", + "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn to use `Vertex AI SDK 2.0` with Bigframes as input data source.\n", + "\n", + "\n", + "This tutorial uses the following Google Cloud ML services:\n", + "\n", + "- `Vertex AI Training`\n", + "- `Vertex AI Remote Training`\n", + "\n", + "\n", + "The steps performed include:\n", + "\n", + "- Initialize a dataframe from a BigQuery table and split the dataset\n", + "- Perform transformations as a Vertex AI remote training.\n", + "- Train the model remotely and evaluate the model locally\n", + "\n", + "**Local-to-remote training**\n", + "\n", + "```\n", + "import vertexai\n", + "from my_module import MyModelClass\n", + "\n", + "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n", + "\n", + "# Wrap the model class with `vertex_ai.preview.remote`\n", + "MyModelClass = vertexai.preview.remote(MyModelClass)\n", + "\n", + "# Instantiate the class\n", + "model = MyModelClass(...)\n", + "\n", + "# Optional set remote config\n", + "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n", + "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n", + "\n", + "# This `fit` call will be executed remotely\n", + "model.fit(...)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset\n", + "\n", + "This tutorial uses the IRIS dataset, which predicts the iris species." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aed92deeb4a0" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* Vertex AI\n", + "* BigQuery\n", + "* Cloud Storage\n", + "\n", + "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n", + "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n", + "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n", + "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", + "to generate a cost estimate based on your projected usage." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i7EUnXsZhAGF" + }, + "source": [ + "## Installation\n", + "\n", + "Install the following packages required to execute this notebook. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "# Install the packages\n", + "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n", + "! pip3 install --upgrade --quiet bigframes\n", + "! pip3 install --upgrade --quiet tensorflow==2.12.0" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "58707a750154" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f200f10a1da3" + }, + "outputs": [], + "source": [ + "# Automatically restart kernel after installs so that your environment can access the new packages\n", + "# import IPython\n", + "\n", + "# app = IPython.Application.instance()\n", + "# app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BF1j6f9HApxa" + }, + "source": [ + "## Before you begin\n", + "\n", + "### Set up your Google Cloud project\n", + "\n", + "**The following steps are required, regardless of your notebook environment.**\n", + "\n", + "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n", + "\n", + "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", + "\n", + "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WReHDGG5g0XY" + }, + "source": [ + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oM1iC_MfAts1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "region" + }, + "source": [ + "#### Region\n", + "\n", + "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "region" + }, + "outputs": [], + "source": [ + "REGION = \"us-central1\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sBCra4QMA2wR" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "74ccc9e52986" + }, + "source": [ + "**1. Vertex AI Workbench**\n", + "* Do nothing as you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "de775a3773ba" + }, + "source": [ + "**2. Local JupyterLab instance, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "254614fa0c46" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ef21552ccea8" + }, + "source": [ + "**3. Colab, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "603adbbf0532" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f6b2ccc891ed" + }, + "source": [ + "**4. Service account or other**\n", + "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zgPO1eR3CYjk" + }, + "source": [ + "### Create a Cloud Storage bucket\n", + "\n", + "Create a storage bucket to store intermediate artifacts such as datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MzGDU7TWdts_" + }, + "outputs": [], + "source": [ + "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-EcIXiGsCePi" + }, + "source": [ + "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NIq7R4HZCfIc" + }, + "outputs": [], + "source": [ + "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "960505627ddf" + }, + "source": [ + "### Import libraries and define constants" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PyQmSRbKA8r-" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "import tensorflow as tf\n", + "import vertexai\n", + "from tensorflow import keras\n", + "\n", + "bf.options.bigquery.location = \"us\" # Dataset is in 'us' not 'us-central1'\n", + "bf.options.bigquery.project = PROJECT_ID\n", + "\n", + "from bigframes.ml.model_selection import \\\n", + " train_test_split as bf_train_test_split" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "source": [ + "## Initialize Vertex AI SDK for Python\n", + "\n", + "Initialize the Vertex AI SDK for Python for your project and corresponding bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "init_aip:mbsdk,all" + }, + "outputs": [], + "source": [ + "vertexai.init(\n", + " project=PROJECT_ID,\n", + " location=REGION,\n", + " staging_bucket=BUCKET_URI,\n", + ")\n", + "\n", + "REMOTE_JOB_NAME = \"sdk2-bigframes-tensorflow\"\n", + "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "105334524e96" + }, + "source": [ + "## Prepare the dataset\n", + "\n", + "Now load the Iris dataset and split the data into train and test sets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "94576deccd8c" + }, + "outputs": [], + "source": [ + "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n", + "\n", + "species_categories = {\n", + " \"versicolor\": 0,\n", + " \"virginica\": 1,\n", + " \"setosa\": 2,\n", + "}\n", + "df[\"target\"] = df[\"species\"].map(species_categories)\n", + "df = df.drop(columns=[\"species\"])\n", + "\n", + "train, test = bf_train_test_split(df, test_size=0.2)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cfcbce726efa" + }, + "source": [ + "## Remote training with GPU\n", + "\n", + "First, train a TensorFlow model as a remote training job:\n", + "\n", + "- Reinitialize Vertex AI for remote training.\n", + "- Instantiate the tensorflow keras model for the remote training job.\n", + "- Invoke the tensorflow keras model.fit() locally which will launch the remote training job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fd865b0c4e8b" + }, + "outputs": [], + "source": [ + "# Switch to remote mode for training\n", + "vertexai.preview.init(remote=True)\n", + "\n", + "keras.Sequential = vertexai.preview.remote(keras.Sequential)\n", + "\n", + "# Instantiate model\n", + "model = keras.Sequential(\n", + " [keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]\n", + ")\n", + "\n", + "# Specify optimizer and loss function\n", + "model.compile(optimizer=\"adam\", loss=\"mean_squared_error\")\n", + "\n", + "# Set training config\n", + "model.fit.vertex.remote_config.enable_cuda = True\n", + "model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-keras-model-gpu\"\n", + "model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "model.fit.vertex.remote_config.custom_commands = [\"pip install tensorflow-io==0.32.0\"]\n", + "\n", + "# Manually set compute resources this time\n", + "model.fit.vertex.remote_config.machine_type = \"n1-highmem-4\"\n", + "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_K80\"\n", + "model.fit.vertex.remote_config.accelerator_count = 4\n", + "\n", + "# Train model on Vertex\n", + "model.fit(train, epochs=10)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f1af94ac1477" + }, + "source": [ + "## Remote prediction\n", + "\n", + "Obtain predictions from the trained model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1d75879948b5" + }, + "outputs": [], + "source": [ + "vertexai.preview.init(remote=True)\n", + "\n", + "# Set remote config\n", + "model.predict.vertex.remote_config.enable_cuda = False\n", + "model.predict.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-keras-predict-cpu\"\n", + "model.predict.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n", + "model.predict.vertex.remote_config.custom_commands = [\n", + " \"pip install tensorflow-io==0.32.0\"\n", + "]\n", + "\n", + "predictions = model.predict(train)\n", + "\n", + "print(f\"Remote predictions: {predictions}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "798b77c95067" + }, + "source": [ + "## Local evaluation\n", + "\n", + "Evaluate model results locally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "88e734e30791" + }, + "outputs": [], + "source": [ + "# User must convert bigframes to pandas dataframe for local evaluation\n", + "feature_columns = [\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]\n", + "label_columns = [\"target\"]\n", + "\n", + "train_X_np = train[feature_columns].to_pandas().values.astype(float)\n", + "train_y_np = train[label_columns].to_pandas().values.astype(float)\n", + "train_ds = tf.data.Dataset.from_tensor_slices((train_X_np, train_y_np))\n", + "\n", + "test_X_np = test[feature_columns].to_pandas().values.astype(float)\n", + "test_y_np = test[label_columns].to_pandas().values.astype(float)\n", + "test_ds = tf.data.Dataset.from_tensor_slices((test_X_np, test_y_np))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cb8637f783ad" + }, + "outputs": [], + "source": [ + "# Switch to local mode for evaluation\n", + "vertexai.preview.init(remote=False)\n", + "\n", + "# Evaluate model's mean square errors\n", + "print(f\"Train loss: {model.evaluate(train_ds.batch(32))}\")\n", + "\n", + "print(f\"Test loss: {model.evaluate(test_ds.batch(32))}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TpV-iwP9qw9c" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, you can delete the individual resources you created in this tutorial:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sx_vKniMq9ZX" + }, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Delete Cloud Storage objects that were created\n", + "delete_bucket = False\n", + "if delete_bucket or os.getenv(\"IS_TESTING\"):\n", + " ! gsutil -m rm -r $BUCKET_URI" + ] + } + ], + "metadata": { + "colab": { + "collapsed_sections": [], + "name": "sdk2_bigframes_tensorflow.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/noxfile.py b/noxfile.py index 033bbfefe40..81ed9c2b2c4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -610,6 +610,9 @@ def notebook(session): "notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb", "notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb", "notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb", + "notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb", + "notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb", + "notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb", # The experimental notebooks imagine features that don't yet # exist or only exist as temporary prototypes. "notebooks/experimental/longer_ml_demo.ipynb", From 4a84714e2fb07f70c70c79f8b8da9fcb41096e33 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 29 Sep 2023 14:16:37 -0700 Subject: [PATCH 48/80] feat: add df.unstack (#63) --- bigframes/core/__init__.py | 113 +++++++++----- bigframes/core/blocks.py | 143 +++++++----------- bigframes/core/utils.py | 20 +++ bigframes/dataframe.py | 21 +++ tests/system/small/test_dataframe.py | 20 +++ tests/system/small/test_multiindex.py | 31 ++++ .../bigframes_vendored/pandas/core/frame.py | 15 ++ 7 files changed, 241 insertions(+), 122 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index c529f833510..5e0675fd134 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -963,10 +963,11 @@ def unpivot( ], *, passthrough_columns: typing.Sequence[str] = (), - index_col_id: str = "index", + index_col_ids: typing.Sequence[str] = ["index"], dtype: typing.Union[ bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] ] = pandas.Float64Dtype(), + how="left", ) -> ArrayValue: """ Unpivot ArrayValue columns. @@ -981,8 +982,11 @@ def unpivot( Returns: ArrayValue: The unpivoted ArrayValue """ - table = self._to_ibis_expr(ordering_mode="offset_col") + if how not in ("left", "right"): + raise ValueError("'how' must be 'left' or 'right'") + table = self._to_ibis_expr(ordering_mode="unordered", expose_hidden_cols=True) row_n = len(row_labels) + hidden_col_ids = self._hidden_ordering_column_names.keys() if not all( len(source_columns) == row_n for _, source_columns in unpivot_columns ): @@ -992,33 +996,44 @@ def unpivot( unpivot_table = table.cross_join( ibis.memtable({unpivot_offset_id: range(row_n)}) ) - unpivot_offsets_value = ( - ( - (unpivot_table[ORDER_ID_COLUMN] * row_n) - + unpivot_table[unpivot_offset_id] - ) - .cast(ibis_dtypes.int64) - .name(ORDER_ID_COLUMN), - ) - # Use ibis memtable to infer type of rowlabels (if possible) # TODO: Allow caller to specify dtype - labels_ibis_type = ibis.memtable({"col": row_labels})["col"].type() - labels_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(labels_ibis_type) - cases = [ - ( - i, - bigframes.dtypes.literal_to_ibis_scalar( - row_labels[i], force_dtype=labels_dtype # type:ignore - ), - ) - for i in range(len(row_labels)) + if isinstance(row_labels[0], tuple): + labels_table = ibis.memtable(row_labels) + labels_ibis_types = [ + labels_table[col].type() for col in labels_table.columns + ] + else: + labels_ibis_types = [ibis.memtable({"col": row_labels})["col"].type()] + labels_dtypes = [ + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type) + for ibis_type in labels_ibis_types ] - labels_value = ( - typing.cast(ibis_types.IntegerColumn, unpivot_table[unpivot_offset_id]) - .cases(cases, default=None) # type:ignore - .name(index_col_id) - ) + + label_columns = [] + for label_part, (col_id, label_dtype) in enumerate( + zip(index_col_ids, labels_dtypes) + ): + # interpret as tuples even if it wasn't originally so can apply same logic for multi-column labels + labels_as_tuples = [ + label if isinstance(label, tuple) else (label,) for label in row_labels + ] + cases = [ + ( + i, + bigframes.dtypes.literal_to_ibis_scalar( + label_tuple[label_part], # type:ignore + force_dtype=label_dtype, # type:ignore + ), + ) + for i, label_tuple in enumerate(labels_as_tuples) + ] + labels_value = ( + typing.cast(ibis_types.IntegerColumn, unpivot_table[unpivot_offset_id]) + .cases(cases, default=None) # type:ignore + .name(col_id) + ) + label_columns.append(labels_value) unpivot_values = [] for j in range(len(unpivot_columns)): @@ -1042,23 +1057,53 @@ def unpivot( unpivot_values.append(unpivot_value.name(result_col)) unpivot_table = unpivot_table.select( - passthrough_columns, labels_value, *unpivot_values, unpivot_offsets_value + passthrough_columns, + *label_columns, + *unpivot_values, + *hidden_col_ids, + unpivot_offset_id, ) + # Extend the original ordering using unpivot_offset_id + old_ordering = self._ordering + if how == "left": + new_ordering = ExpressionOrdering( + ordering_value_columns=[ + *old_ordering.ordering_value_columns, + OrderingColumnReference(unpivot_offset_id), + ], + total_ordering_columns=frozenset( + [*old_ordering.total_ordering_columns, unpivot_offset_id] + ), + ) + else: # how=="right" + new_ordering = ExpressionOrdering( + ordering_value_columns=[ + OrderingColumnReference(unpivot_offset_id), + *old_ordering.ordering_value_columns, + ], + total_ordering_columns=frozenset( + [*old_ordering.total_ordering_columns, unpivot_offset_id] + ), + ) value_columns = [ unpivot_table[value_col_id] for value_col_id, _ in unpivot_columns ] passthrough_values = [unpivot_table[col] for col in passthrough_columns] + hidden_ordering_columns = [ + unpivot_table[unpivot_offset_id], + *[unpivot_table[hidden_col] for hidden_col in hidden_col_ids], + ] return ArrayValue( session=self._session, table=unpivot_table, - columns=[unpivot_table[index_col_id], *value_columns, *passthrough_values], - hidden_ordering_columns=[unpivot_table[ORDER_ID_COLUMN]], - ordering=ExpressionOrdering( - ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], - integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - ), + columns=[ + *[unpivot_table[col_id] for col_id in index_col_ids], + *value_columns, + *passthrough_values, + ], + hidden_ordering_columns=hidden_ordering_columns, + ordering=new_ordering, ) def assign(self, source_id: str, destination_id: str) -> ArrayValue: diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index b53c2212c10..863852c684c 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -838,7 +838,7 @@ def aggregate_all_and_stack( ] result_expr = self.expr.aggregate(aggregations, dropna=dropna).unpivot( row_labels=self.column_labels.to_list(), - index_col_id="index", + index_col_ids=["index"], unpivot_columns=[(value_col_id, self.value_columns)], dtype=dtype, ) @@ -849,7 +849,7 @@ def aggregate_all_and_stack( expr_with_offsets, offset_col = self.expr.promote_offsets() stacked_expr = expr_with_offsets.unpivot( row_labels=self.column_labels.to_list(), - index_col_id=guid.generate_guid(), + index_col_ids=[guid.generate_guid()], unpivot_columns=[(value_col_id, self.value_columns)], passthrough_columns=[*self.index_columns, offset_col], dtype=dtype, @@ -1041,7 +1041,7 @@ def summarize( expr = self.expr.aggregate(aggregations).unpivot( labels, unpivot_columns=columns, - index_col_id=label_col_id, + index_col_ids=[label_col_id], ) labels = self._get_labels_for_columns(column_ids) return Block(expr, column_labels=labels, index_columns=[label_col_id]) @@ -1225,116 +1225,83 @@ def pivot( return result_block.with_column_labels(column_index) - def stack(self): + def stack(self, how="left", dropna=True, sort=True, levels: int = 1): """Unpivot last column axis level into row axis""" - if isinstance(self.column_labels, pd.MultiIndex): - return self._stack_multi() - else: - return self._stack_mono() - - def _stack_mono(self): - if isinstance(self.column_labels, pd.MultiIndex): - raise ValueError("Expected single level index") - # These are the values that will be turned into rows - stack_values = self.column_labels.drop_duplicates().sort_values() - # Get matching columns - unpivot_columns: List[Tuple[str, List[str]]] = [] - dtypes: List[bigframes.dtypes.Dtype] = [] - col_id = guid.generate_guid("unpivot_") - dtype = None - input_columns: Sequence[Optional[str]] = [] - for uvalue in stack_values: - matching_ids = self.label_to_col_id.get(uvalue, []) - input_id = matching_ids[0] if len(matching_ids) > 0 else None - if input_id: - if dtype and dtype != self._column_type(input_id): - raise NotImplementedError( - "Cannot stack columns with non-matching dtypes." - ) - else: - dtype = self._column_type(input_id) - input_columns.append(input_id) - unpivot_columns.append((col_id, input_columns)) - if dtype: - dtypes.append(dtype or pd.Float64Dtype()) + col_labels, row_labels = utils.split_index(self.column_labels, levels=levels) + if dropna: + row_labels = row_labels.drop_duplicates() + if sort: + row_labels = row_labels.sort_values() - added_index_column = col_id = guid.generate_guid() - unpivot_expr = self._expr.unpivot( - row_labels=stack_values, - passthrough_columns=self.index_columns, - unpivot_columns=unpivot_columns, - index_col_id=added_index_column, - dtype=dtypes, - ) - block = Block( - unpivot_expr, - index_columns=[*self.index_columns, added_index_column], - column_labels=[None], - index_labels=[*self._index_labels, self.column_labels.names[-1]], - ) - return block - - def _stack_multi(self): - if not isinstance(self.column_labels, pd.MultiIndex): - raise ValueError("Expected multi-index") - - # These are the values that will be turned into rows - stack_values = ( - self.column_labels.get_level_values(-1).drop_duplicates().sort_values() - ) + row_label_tuples = utils.index_as_tuples(row_labels) - result_col_labels = ( - self.column_labels.droplevel(-1) - .drop_duplicates() - .sort_values() - .dropna(how="all") - ) + if col_labels is not None: + result_index = col_labels.drop_duplicates().sort_values().dropna(how="all") + result_col_labels = utils.index_as_tuples(result_index) + else: + result_index = pd.Index([None]) + result_col_labels = list([()]) # Get matching columns unpivot_columns: List[Tuple[str, List[str]]] = [] dtypes = [] for val in result_col_labels: col_id = guid.generate_guid("unpivot_") - dtype = None - input_columns: Sequence[Optional[str]] = [] - for uvalue in stack_values: - # Need to unpack if still a multi-index after dropping 1 level - label_to_match = ( - (val, uvalue) if result_col_labels.nlevels == 1 else (*val, uvalue) - ) - matching_ids = self.label_to_col_id.get(label_to_match, []) - input_id = matching_ids[0] if len(matching_ids) > 0 else None - if input_id: - if dtype and dtype != self._column_type(input_id): - raise NotImplementedError( - "Cannot stack columns with non-matching dtypes." - ) - else: - dtype = self._column_type(input_id) - input_columns.append(input_id) - # Input column i is the first one that + input_columns, dtype = self._create_stack_column(val, row_label_tuples) unpivot_columns.append((col_id, input_columns)) if dtype: dtypes.append(dtype or pd.Float64Dtype()) - added_index_column = col_id = guid.generate_guid() + added_index_columns = [guid.generate_guid() for _ in range(row_labels.nlevels)] unpivot_expr = self._expr.unpivot( - row_labels=stack_values, + row_labels=row_label_tuples, passthrough_columns=self.index_columns, unpivot_columns=unpivot_columns, - index_col_id=added_index_column, + index_col_ids=added_index_columns, dtype=dtypes, + how=how, ) + new_index_level_names = self.column_labels.names[-levels:] + if how == "left": + index_columns = [*self.index_columns, *added_index_columns] + index_labels = [*self._index_labels, *new_index_level_names] + else: + index_columns = [*added_index_columns, *self.index_columns] + index_labels = [*new_index_level_names, *self._index_labels] + block = Block( unpivot_expr, - index_columns=[*self.index_columns, added_index_column], - column_labels=result_col_labels, - index_labels=[*self._index_labels, self.column_labels.names[-1]], + index_columns=index_columns, + column_labels=result_index, + index_labels=index_labels, ) return block + def _create_stack_column( + self, col_label: typing.Tuple, stack_labels: typing.Sequence[typing.Tuple] + ): + dtype = None + input_columns: list[Optional[str]] = [] + for uvalue in stack_labels: + label_to_match = (*col_label, *uvalue) + label_to_match = ( + label_to_match[0] if len(label_to_match) == 1 else label_to_match + ) + matching_ids = self.label_to_col_id.get(label_to_match, []) + input_id = matching_ids[0] if len(matching_ids) > 0 else None + if input_id: + if dtype and dtype != self._column_type(input_id): + raise NotImplementedError( + "Cannot stack columns with non-matching dtypes." + ) + else: + dtype = self._column_type(input_id) + input_columns.append(input_id) + # Input column i is the first one that + return input_columns, dtype or pd.Float64Dtype() + def _column_type(self, col_id: str) -> bigframes.dtypes.Dtype: col_offset = self.value_columns.index(col_id) dtype = self.dtypes[col_offset] diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py index 75175690ce0..1a5d3ab37ca 100644 --- a/bigframes/core/utils.py +++ b/bigframes/core/utils.py @@ -49,6 +49,26 @@ def combine_indices(index1: pd.Index, index2: pd.Index) -> pd.MultiIndex: return multi_index +def index_as_tuples(index: pd.Index) -> typing.Sequence[typing.Tuple]: + if isinstance(index, pd.MultiIndex): + return [label for label in index] + else: + return [(label,) for label in index] + + +def split_index( + index: pd.Index, levels: int = 1 +) -> typing.Tuple[typing.Optional[pd.Index], pd.Index]: + nlevels = index.nlevels + remaining = nlevels - levels + if remaining > 0: + return index.droplevel(list(range(remaining, nlevels))), index.droplevel( + list(range(0, remaining)) + ) + else: + return (None, index) + + def get_standardized_ids( col_labels: Iterable[Hashable], idx_labels: Iterable[Hashable] = () ) -> tuple[list[str], list[str]]: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 828d2df64ef..ea06e28cdfc 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1706,6 +1706,27 @@ def stack(self): return bigframes.series.Series(result_block) return DataFrame(result_block) + def unstack(self): + block = self._block + # Special case, unstack with mono-index transpose into a series + if self.index.nlevels == 1: + block = block.stack( + how="right", dropna=False, sort=False, levels=self.columns.nlevels + ) + return bigframes.series.Series(block) + + # Pivot by last level of index + index_ids = block.index_columns + block = block.reset_index(drop=False) + block = block.set_index(index_ids[:-1]) + + pivot_block = block.pivot( + columns=[index_ids[-1]], + values=self._block.value_columns, + values_in_index=True, + ) + return DataFrame(pivot_block) + def _drop_non_numeric(self, keep_bool=True) -> DataFrame: types_to_keep = set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES) if not keep_bool: diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index ba76c4b0d3e..e71b1430e61 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1771,6 +1771,26 @@ def test_df_stack(scalars_dfs): pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) +def test_df_unstack(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + # To match bigquery dataframes + scalars_pandas_df = scalars_pandas_df.copy() + scalars_pandas_df.columns = scalars_pandas_df.columns.astype("string[pyarrow]") + # Can only stack identically-typed columns + columns = [ + "rowindex_2", + "int64_col", + "int64_too", + ] + + # unstack on mono-index produces series + bf_result = scalars_df[columns].unstack().to_pandas() + pd_result = scalars_pandas_df[columns].unstack() + + # Pandas produces NaN, where bq dataframes produces pd.NA + pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + @pytest.mark.parametrize( ("values", "index", "columns"), [ diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 1e38b47b4cc..10f2a74b21b 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -729,6 +729,26 @@ def test_column_multi_index_stack(scalars_df_index, scalars_pandas_df_index): ) +def test_column_multi_index_unstack(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_too", "int64_col", "rowindex_2"] + level1 = pandas.Index(["b", "a", "b"], dtype="string[pyarrow]") + # Need resulting column to be pyarrow string rather than object dtype + level2 = pandas.Index(["a", "b", "b"], dtype="string[pyarrow]") + multi_columns = pandas.MultiIndex.from_arrays([level1, level2]) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.unstack().to_pandas() + # Shifting sort behavior in stack + pd_result = pd_df.unstack() + + # Pandas produces NaN, where bq dataframes produces pd.NA + # Column ordering seems to depend on pandas version + pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) + + @pytest.mark.skip(reason="Pandas fails in newer versions.") def test_column_multi_index_w_na_stack(scalars_df_index, scalars_pandas_df_index): columns = ["int64_too", "int64_col", "rowindex_2"] @@ -866,6 +886,17 @@ def test_column_multi_index_reorder_levels(scalars_df_index, scalars_pandas_df_i pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_multi_index_unstack(hockey_df, hockey_pandas_df): + bf_result = ( + hockey_df.set_index(["team_name", "season", "position"]).unstack().to_pandas() + ) + pd_result = hockey_pandas_df.set_index( + ["team_name", "season", "position"] + ).unstack() + + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + def test_column_multi_index_swaplevel(scalars_df_index, scalars_pandas_df_index): columns = ["int64_too", "string_col", "bool_col"] multi_columns = pandas.MultiIndex.from_tuples( diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 5cd9fe51637..348145a4d65 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -1972,6 +1972,21 @@ def stack(self): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def unstack(self): + """ + Pivot a level of the (necessarily hierarchical) index labels. + + Returns a DataFrame having a new level of column labels whose inner-most level + consists of the pivoted index labels. + + If the index is not a MultiIndex, the output will be a Series + (the analogue of stack when the columns are not a MultiIndex). + + Returns: + DataFrame or Series + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # Add index and columns From 971d091cac9ad662145a3d43d8f9a785eb0ccc23 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Mon, 2 Oct 2023 20:24:53 -0700 Subject: [PATCH 49/80] fix: fix df/series.iloc by list with multiindex (#79) * fix: fix df/series.iloc by list with multiindex * fix git error --------- Co-authored-by: Henry J Solberg --- bigframes/core/indexers.py | 15 ++++++++------- tests/system/small/test_dataframe.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index a538c80711b..e22a5aed857 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -332,8 +332,6 @@ def _iloc_getitem_series_or_dataframe( elif isinstance(key, slice): return series_or_dataframe._slice(key.start, key.stop, key.step) elif pd.api.types.is_list_like(key): - # TODO(henryjsolberg): support MultiIndex - if len(key) == 0: return typing.cast( typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], @@ -346,15 +344,18 @@ def _iloc_getitem_series_or_dataframe( original_series_name if original_series_name is not None else "0" ) df = series_or_dataframe.to_frame() - original_index_name = df.index.name - temporary_index_name = guid.generate_guid(prefix="temp_iloc_index_") - df = df.rename_axis(temporary_index_name) + original_index_names = df.index.names + temporary_index_names = [ + guid.generate_guid(prefix="temp_iloc_index_") + for _ in range(len(df.index.names)) + ] + df = df.rename_axis(temporary_index_names) # set to offset index and use regular loc, then restore index df = df.reset_index(drop=False) result = df.loc[key] - result = result.set_index(temporary_index_name) - result = result.rename_axis(original_index_name) + result = result.set_index(temporary_index_names) + result = result.rename_axis(original_index_names) if isinstance(series_or_dataframe, bigframes.series.Series): result = result[series_name] diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index e71b1430e61..aad6faa4b6b 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2525,6 +2525,24 @@ def test_iloc_list(scalars_df_index, scalars_pandas_df_index): ) +def test_iloc_list_multiindex(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + scalars_df = scalars_df.copy() + scalars_pandas_df = scalars_pandas_df.copy() + scalars_df = scalars_df.set_index(["bytes_col", "numeric_col"]) + scalars_pandas_df = scalars_pandas_df.set_index(["bytes_col", "numeric_col"]) + + index_list = [0, 0, 0, 5, 4, 7] + + bf_result = scalars_df.iloc[index_list] + pd_result = scalars_pandas_df.iloc[index_list] + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + def test_iloc_empty_list(scalars_df_index, scalars_pandas_df_index): index_list = [] From c9fa85cc338be5e9a8dde59b255690aedbbc1127 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 3 Oct 2023 00:59:44 -0700 Subject: [PATCH 50/80] feat: add multi-column dataframe merge (#73) --- bigframes/core/blocks.py | 73 ++++++++++ bigframes/core/joins/single_column.py | 60 +++------ bigframes/core/utils.py | 33 +++++ bigframes/dataframe.py | 127 +++++------------- tests/system/small/test_dataframe.py | 35 ++++- .../bigframes_vendored/pandas/core/frame.py | 12 +- .../pandas/core/reshape/merge.py | 12 +- 7 files changed, 206 insertions(+), 146 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 863852c684c..0161d173618 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -38,6 +38,7 @@ import bigframes.core as core import bigframes.core.guid as guid import bigframes.core.indexes as indexes +import bigframes.core.joins as joins import bigframes.core.ordering as ordering import bigframes.core.utils import bigframes.core.utils as utils @@ -1403,6 +1404,78 @@ def concat( result_block = result_block.reset_index() return result_block + def merge( + self, + other: Block, + how: typing.Literal[ + "inner", + "left", + "outer", + "right", + ], + left_col_ids: typing.Sequence[str], + right_col_ids: typing.Sequence[str], + sort: bool, + suffixes: tuple[str, str] = ("_x", "_y"), + ) -> Block: + ( + joined_expr, + coalesced_join_cols, + (get_column_left, get_column_right), + ) = joins.join_by_column( + self.expr, + left_col_ids, + other.expr, + right_col_ids, + how=how, + sort=sort, + ) + + # which join key parts should be coalesced + merge_join_key_mask = [ + str(self.col_id_to_label[left_id]) == str(other.col_id_to_label[right_id]) + for left_id, right_id in zip(left_col_ids, right_col_ids) + ] + labels_to_coalesce = [ + self.col_id_to_label[col_id] + for i, col_id in enumerate(left_col_ids) + if merge_join_key_mask[i] + ] + + def left_col_mapping(col_id: str) -> str: + if col_id in left_col_ids: + join_key_part = left_col_ids.index(col_id) + if merge_join_key_mask[join_key_part]: + return coalesced_join_cols[join_key_part] + return get_column_left(col_id) + + def right_col_mapping(col_id: str) -> typing.Optional[str]: + if col_id in right_col_ids: + join_key_part = right_col_ids.index(col_id) + if merge_join_key_mask[join_key_part]: + return None + return get_column_right(col_id) + + left_columns = [left_col_mapping(col_id) for col_id in self.value_columns] + + right_columns = [ + typing.cast(str, right_col_mapping(col_id)) + for col_id in other.value_columns + if right_col_mapping(col_id) + ] + + expr = joined_expr.select_columns([*left_columns, *right_columns]) + labels = utils.merge_column_labels( + self.column_labels, + other.column_labels, + coalesce_labels=labels_to_coalesce, + suffixes=suffixes, + ) + + # Constructs default index + expr, offset_index_id = expr.promote_offsets() + return Block(expr, index_columns=[offset_index_id], column_labels=labels) + def _force_reproject(self) -> Block: """Forces a reprojection of the underlying tables expression. Used to force predicate/order application before subsequent operations.""" return Block( diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 8a9825cf0b5..2d616fc3f0c 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -44,7 +44,6 @@ def join_by_column( "right", ], sort: bool = False, - coalesce_join_keys: bool = True, allow_row_identity_join: bool = True, ) -> Tuple[ core.ArrayValue, @@ -59,8 +58,6 @@ def join_by_column( right: Expression for right table to join. right_column_ids: Column IDs (not label) to join by. how: The type of join to perform. - coalesce_join_keys: if set to False, returned column ids will contain - both left and right join key columns. allow_row_identity_join (bool): If True, allow matching by row identity. Set to False to always perform a true JOIN in generated SQL. @@ -71,8 +68,6 @@ def join_by_column( * Sequence[str]: Column IDs of the coalesced join columns. Sometimes either the left/right table will have missing rows. This column pulls the non-NULL value from either left/right. - If coalesce_join_keys is False, will return uncombined left and - right key columns. * Tuple[Callable, Callable]: For a given column ID from left or right, respectively, return the new column id from the combined expression. """ @@ -100,9 +95,7 @@ def join_by_column( right_join_keys = [ combined_expr.get_column(get_column_right(col)) for col in right_column_ids ] - join_key_cols = get_join_cols( - left_join_keys, right_join_keys, how, coalesce_join_keys - ) + join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how) join_key_ids = [col.get_name() for col in join_key_cols] combined_expr = combined_expr.projection( [*join_key_cols, *combined_expr.columns] @@ -182,9 +175,7 @@ def get_column_right(col_id): right_join_keys = [ combined_table[get_column_right(col)] for col in right_column_ids ] - join_key_cols = get_join_cols( - left_join_keys, right_join_keys, how, coalesce_join_keys - ) + join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how) # We could filter out the original join columns, but predicates/ordering # might still reference them in implicit joins. columns = ( @@ -226,46 +217,35 @@ def get_column_right(col_id): ) -def get_join_cols( +def get_coalesced_join_cols( left_join_cols: typing.Iterable[ibis_types.Value], right_join_cols: typing.Iterable[ibis_types.Value], how: str, - coalesce_join_keys: bool = True, ) -> typing.List[ibis_types.Value]: join_key_cols: list[ibis_types.Value] = [] for left_col, right_col in zip(left_join_cols, right_join_cols): - if not coalesce_join_keys: + if how == "left" or how == "inner": join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) + elif how == "right": join_key_cols.append(right_col.name(guid.generate_guid(prefix="index_"))) - else: - if how == "left" or how == "inner": + elif how == "outer": + # The left index and the right index might contain null values, for + # example due to an outer join with different numbers of rows. Coalesce + # these to take the index value from either column. + # Use a random name in case the left index and the right index have the + # same name. In such a case, _x and _y suffixes will already be used. + # Don't need to coalesce if they are exactly the same column. + if left_col.name("index").equals(right_col.name("index")): join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) - elif how == "right": - join_key_cols.append( - right_col.name(guid.generate_guid(prefix="index_")) - ) - elif how == "outer": - # The left index and the right index might contain null values, for - # example due to an outer join with different numbers of rows. Coalesce - # these to take the index value from either column. - # Use a random name in case the left index and the right index have the - # same name. In such a case, _x and _y suffixes will already be used. - # Don't need to coalesce if they are exactly the same column. - if left_col.name("index").equals(right_col.name("index")): - join_key_cols.append( - left_col.name(guid.generate_guid(prefix="index_")) - ) - else: - join_key_cols.append( - ibis.coalesce( - left_col, - right_col, - ).name(guid.generate_guid(prefix="index_")) - ) else: - raise ValueError( - f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}" + join_key_cols.append( + ibis.coalesce( + left_col, + right_col, + ).name(guid.generate_guid(prefix="index_")) ) + else: + raise ValueError(f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}") return join_key_cols diff --git a/bigframes/core/utils.py b/bigframes/core/utils.py index 1a5d3ab37ca..dc7c7090117 100644 --- a/bigframes/core/utils.py +++ b/bigframes/core/utils.py @@ -104,3 +104,36 @@ def get_standardized_ids( idx_ids, col_ids = ids[: len(idx_ids)], ids[len(idx_ids) :] return col_ids, idx_ids + + +def merge_column_labels( + left_labels: pd.Index, + right_labels: pd.Index, + coalesce_labels: typing.Sequence, + suffixes: tuple[str, str] = ("_x", "_y"), +) -> pd.Index: + result_labels = [] + + for col_label in left_labels: + if col_label in right_labels: + if col_label in coalesce_labels: + # Merging on the same column only returns 1 key column from coalesce both. + # Take the left key column. + result_labels.append(col_label) + else: + result_labels.append(str(col_label) + suffixes[0]) + else: + result_labels.append(col_label) + + for col_label in right_labels: + if col_label in left_labels: + if col_label in coalesce_labels: + # Merging on the same column only returns 1 key column from coalesce both. + # Pass the right key column. + pass + else: + result_labels.append(str(col_label) + suffixes[1]) + else: + result_labels.append(col_label) + + return pd.Index(result_labels) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index ea06e28cdfc..b8ae37ee6c0 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -46,7 +46,6 @@ import bigframes.core.indexers as indexers import bigframes.core.indexes as indexes import bigframes.core.io -import bigframes.core.joins as joins import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.core.window @@ -1779,12 +1778,10 @@ def merge( ] = "inner", # TODO(garrettwu): Currently can take inner, outer, left and right. To support # cross joins - # TODO(garrettwu): Support "on" list of columns and None. Currently a single - # column must be provided - on: Optional[str] = None, + on: Union[blocks.Label, Sequence[blocks.Label], None] = None, *, - left_on: Optional[str] = None, - right_on: Optional[str] = None, + left_on: Union[blocks.Label, Sequence[blocks.Label], None] = None, + right_on: Union[blocks.Label, Sequence[blocks.Label], None] = None, sort: bool = False, suffixes: tuple[str, str] = ("_x", "_y"), ) -> DataFrame: @@ -1798,97 +1795,41 @@ def merge( ) left_on, right_on = on, on - left = self - left_on_sql = self._sql_names(left_on) - # 0 elements already throws an exception - if len(left_on_sql) > 1: - raise ValueError(f"The column label {left_on} is not unique.") - left_on_sql = left_on_sql[0] - - right_on_sql = right._sql_names(right_on) - if len(right_on_sql) > 1: - raise ValueError(f"The column label {right_on} is not unique.") - right_on_sql = right_on_sql[0] - - ( - joined_expr, - join_key_ids, - (get_column_left, get_column_right), - ) = joins.join_by_column( - left._block.expr, - [left_on_sql], - right._block.expr, - [right_on_sql], - how=how, - sort=sort, - # In merging on the same column, it only returns 1 key column from coalesced both. - # While if 2 different columns, both will be presented in the result. - coalesce_join_keys=(left_on == right_on), - ) - # TODO(swast): Add suffixes to the column labels instead of reusing the - # column IDs as the new labels. - # Drop the index column(s) to be consistent with pandas. - left_columns = [ - join_key_ids[0] if (col_id == left_on_sql) else get_column_left(col_id) - for col_id in left._block.value_columns - ] - - right_columns = [] - for col_id in right._block.value_columns: - if col_id == right_on_sql: - # When left_on == right_on - if len(join_key_ids) > 1: - right_columns.append(join_key_ids[1]) - else: - right_columns.append(get_column_right(col_id)) - - expr = joined_expr.select_columns([*left_columns, *right_columns]) - labels = self._get_merged_col_labels( - right, left_on=left_on, right_on=right_on, suffixes=suffixes - ) + if utils.is_list_like(left_on): + left_on = list(left_on) # type: ignore + else: + left_on = [left_on] - # Constructs default index - expr, offset_index_id = expr.promote_offsets() - block = blocks.Block( - expr, index_columns=[offset_index_id], column_labels=labels + if utils.is_list_like(right_on): + right_on = list(right_on) # type: ignore + else: + right_on = [right_on] + + left_join_ids = [] + for label in left_on: # type: ignore + left_col_id = self._resolve_label_exact(label) + # 0 elements already throws an exception + if not left_col_id: + raise ValueError(f"No column {label} found in self.") + left_join_ids.append(left_col_id) + + right_join_ids = [] + for label in right_on: # type: ignore + right_col_id = right._resolve_label_exact(label) + if not right_col_id: + raise ValueError(f"No column {label} found in other.") + right_join_ids.append(right_col_id) + + block = self._block.merge( + right._block, + how, + left_join_ids, + right_join_ids, + sort=sort, + suffixes=suffixes, ) return DataFrame(block) - def _get_merged_col_labels( - self, - right: DataFrame, - left_on: str, - right_on: str, - suffixes: tuple[str, str] = ("_x", "_y"), - ) -> List[blocks.Label]: - on_col_equal = left_on == right_on - - left_col_labels: list[blocks.Label] = [] - for col_label in self._block.column_labels: - if col_label in right._block.column_labels: - if on_col_equal and col_label == left_on: - # Merging on the same column only returns 1 key column from coalesce both. - # Take the left key column. - left_col_labels.append(col_label) - else: - left_col_labels.append(str(col_label) + suffixes[0]) - else: - left_col_labels.append(col_label) - - right_col_labels: list[blocks.Label] = [] - for col_label in right._block.column_labels: - if col_label in self._block.column_labels: - if on_col_equal and col_label == left_on: - # Merging on the same column only returns 1 key column from coalesce both. - # Pass the right key column. - pass - else: - right_col_labels.append(str(col_label) + suffixes[1]) - else: - right_col_labels.append(col_label) - - return left_col_labels + right_col_labels - def join( self, other: DataFrame, *, on: Optional[str] = None, how: str = "left" ) -> DataFrame: diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index aad6faa4b6b..a27d8fd487c 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -757,7 +757,7 @@ def test_df_isin_dict(scalars_dfs): ("right",), ], ) -def test_merge(scalars_dfs, merge_how): +def test_df_merge(scalars_dfs, merge_how): scalars_df, scalars_pandas_df = scalars_dfs on = "rowindex_2" left_columns = ["int64_col", "float64_col", "rowindex_2"] @@ -782,6 +782,39 @@ def test_merge(scalars_dfs, merge_how): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.parametrize( + ("left_on", "right_on"), + [ + (["int64_col", "rowindex_2"], ["int64_col", "rowindex_2"]), + (["rowindex_2", "int64_col"], ["int64_col", "rowindex_2"]), + (["rowindex_2", "float64_col"], ["int64_col", "rowindex_2"]), + ], +) +def test_df_merge_multi_key(scalars_dfs, left_on, right_on): + scalars_df, scalars_pandas_df = scalars_dfs + left_columns = ["int64_col", "float64_col", "rowindex_2"] + right_columns = ["int64_col", "bool_col", "string_col", "rowindex_2"] + + left = scalars_df[left_columns] + # Offset the rows somewhat so that outer join can have an effect. + right = scalars_df[right_columns].assign(rowindex_2=scalars_df["rowindex_2"] + 2) + + df = left.merge(right, "outer", left_on=left_on, right_on=right_on, sort=True) + bf_result = df.to_pandas() + + pd_result = scalars_pandas_df[left_columns].merge( + scalars_pandas_df[right_columns].assign( + rowindex_2=scalars_pandas_df["rowindex_2"] + 2 + ), + "outer", + left_on=left_on, + right_on=right_on, + sort=True, + ) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + @pytest.mark.parametrize( ("merge_how",), [ diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 348145a4d65..11aaf84b14e 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -1468,14 +1468,14 @@ def merge( ``inner``: use intersection of keys from both frames, similar to a SQL inner join; preserve the order of the left keys. - on: - Column join on. It must be found in both DataFrames. Either on or left_on + right_on + on (label or list of labels): + Columns to join on. It must be found in both DataFrames. Either on or left_on + right_on must be passed in. - left_on: - Column join on in the left DataFrame. Either on or left_on + right_on + left_on (label or list of labels): + Columns to join on in the left DataFrame. Either on or left_on + right_on must be passed in. - right_on: - Column join on in the right DataFrame. Either on or left_on + right_on + right_on (label or list of labels): + Columns to join on in the right DataFrame. Either on or left_on + right_on must be passed in. sort: Default False. Sort the join keys lexicographically in the diff --git a/third_party/bigframes_vendored/pandas/core/reshape/merge.py b/third_party/bigframes_vendored/pandas/core/reshape/merge.py index ee02d698da3..cc81de405b9 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/merge.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/merge.py @@ -51,14 +51,14 @@ def merge( ``inner``: use intersection of keys from both frames, similar to a SQL inner join; preserve the order of the left keys. - on: - Column join on. It must be found in both DataFrames. Either on or left_on + right_on + on (label or list of labels): + Columns to join on. It must be found in both DataFrames. Either on or left_on + right_on must be passed in. - left_on: - Column join on in the left DataFrame. Either on or left_on + right_on + left_on (label or list of labels): + Columns to join on in the left DataFrame. Either on or left_on + right_on must be passed in. - right_on: - Column join on in the right DataFrame. Either on or left_on + right_on + right_on (label or list of labels): + Columns to join on in the right DataFrame. Either on or left_on + right_on must be passed in. sort: Default False. Sort the join keys lexicographically in the From c506200dfc770b981d81013c63d822466ea6350b Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Tue, 3 Oct 2023 10:39:27 -0500 Subject: [PATCH 51/80] chore: [autoapprove] bump cryptography from 41.0.3 to 41.0.4 (#75) Source-Link: https://github.com/googleapis/synthtool/commit/dede53ff326079b457cfb1aae5bbdc82cbb51dc3 Post-Processor: gcr.io/cloud-devrel-public-resources/owlbot-python:latest@sha256:fac304457974bb530cc5396abd4ab25d26a469cd3bc97cbfb18c8d4324c584eb Co-authored-by: Owl Bot Co-authored-by: Tim Swast --- .github/.OwlBot.lock.yaml | 4 ++-- .gitignore | 1 + .kokoro/requirements.txt | 49 ++++++++++++++++++++------------------- 3 files changed, 28 insertions(+), 26 deletions(-) diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index a3da1b0d4cd..a9bdb1b7ac0 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,5 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:3e3800bb100af5d7f9e810d48212b37812c1856d20ffeafb99ebe66461b61fc7 -# created: 2023-08-02T10:53:29.114535628Z + digest: sha256:fac304457974bb530cc5396abd4ab25d26a469cd3bc97cbfb18c8d4324c584eb +# created: 2023-10-02T21:31:03.517640371Z diff --git a/.gitignore b/.gitignore index b4243ced74e..d083ea1ddc3 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ docs.metadata # Virtual environment env/ +venv/ # Test logs coverage.xml diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 029bd342de9..96d593c8c82 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -113,30 +113,30 @@ commonmark==0.9.1 \ --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 # via rich -cryptography==41.0.3 \ - --hash=sha256:0d09fb5356f975974dbcb595ad2d178305e5050656affb7890a1583f5e02a306 \ - --hash=sha256:23c2d778cf829f7d0ae180600b17e9fceea3c2ef8b31a99e3c694cbbf3a24b84 \ - --hash=sha256:3fb248989b6363906827284cd20cca63bb1a757e0a2864d4c1682a985e3dca47 \ - --hash=sha256:41d7aa7cdfded09b3d73a47f429c298e80796c8e825ddfadc84c8a7f12df212d \ - --hash=sha256:42cb413e01a5d36da9929baa9d70ca90d90b969269e5a12d39c1e0d475010116 \ - --hash=sha256:4c2f0d35703d61002a2bbdcf15548ebb701cfdd83cdc12471d2bae80878a4207 \ - --hash=sha256:4fd871184321100fb400d759ad0cddddf284c4b696568204d281c902fc7b0d81 \ - --hash=sha256:5259cb659aa43005eb55a0e4ff2c825ca111a0da1814202c64d28a985d33b087 \ - --hash=sha256:57a51b89f954f216a81c9d057bf1a24e2f36e764a1ca9a501a6964eb4a6800dd \ - --hash=sha256:652627a055cb52a84f8c448185922241dd5217443ca194d5739b44612c5e6507 \ - --hash=sha256:67e120e9a577c64fe1f611e53b30b3e69744e5910ff3b6e97e935aeb96005858 \ - --hash=sha256:6af1c6387c531cd364b72c28daa29232162010d952ceb7e5ca8e2827526aceae \ - --hash=sha256:6d192741113ef5e30d89dcb5b956ef4e1578f304708701b8b73d38e3e1461f34 \ - --hash=sha256:7efe8041897fe7a50863e51b77789b657a133c75c3b094e51b5e4b5cec7bf906 \ - --hash=sha256:84537453d57f55a50a5b6835622ee405816999a7113267739a1b4581f83535bd \ - --hash=sha256:8f09daa483aedea50d249ef98ed500569841d6498aa9c9f4b0531b9964658922 \ - --hash=sha256:95dd7f261bb76948b52a5330ba5202b91a26fbac13ad0e9fc8a3ac04752058c7 \ - --hash=sha256:a74fbcdb2a0d46fe00504f571a2a540532f4c188e6ccf26f1f178480117b33c4 \ - --hash=sha256:a983e441a00a9d57a4d7c91b3116a37ae602907a7618b882c8013b5762e80574 \ - --hash=sha256:ab8de0d091acbf778f74286f4989cf3d1528336af1b59f3e5d2ebca8b5fe49e1 \ - --hash=sha256:aeb57c421b34af8f9fe830e1955bf493a86a7996cc1338fe41b30047d16e962c \ - --hash=sha256:ce785cf81a7bdade534297ef9e490ddff800d956625020ab2ec2780a556c313e \ - --hash=sha256:d0d651aa754ef58d75cec6edfbd21259d93810b73f6ec246436a21b7841908de +cryptography==41.0.4 \ + --hash=sha256:004b6ccc95943f6a9ad3142cfabcc769d7ee38a3f60fb0dddbfb431f818c3a67 \ + --hash=sha256:047c4603aeb4bbd8db2756e38f5b8bd7e94318c047cfe4efeb5d715e08b49311 \ + --hash=sha256:0d9409894f495d465fe6fda92cb70e8323e9648af912d5b9141d616df40a87b8 \ + --hash=sha256:23a25c09dfd0d9f28da2352503b23e086f8e78096b9fd585d1d14eca01613e13 \ + --hash=sha256:2ed09183922d66c4ec5fdaa59b4d14e105c084dd0febd27452de8f6f74704143 \ + --hash=sha256:35c00f637cd0b9d5b6c6bd11b6c3359194a8eba9c46d4e875a3660e3b400005f \ + --hash=sha256:37480760ae08065437e6573d14be973112c9e6dcaf5f11d00147ee74f37a3829 \ + --hash=sha256:3b224890962a2d7b57cf5eeb16ccaafba6083f7b811829f00476309bce2fe0fd \ + --hash=sha256:5a0f09cefded00e648a127048119f77bc2b2ec61e736660b5789e638f43cc397 \ + --hash=sha256:5b72205a360f3b6176485a333256b9bcd48700fc755fef51c8e7e67c4b63e3ac \ + --hash=sha256:7e53db173370dea832190870e975a1e09c86a879b613948f09eb49324218c14d \ + --hash=sha256:7febc3094125fc126a7f6fb1f420d0da639f3f32cb15c8ff0dc3997c4549f51a \ + --hash=sha256:80907d3faa55dc5434a16579952ac6da800935cd98d14dbd62f6f042c7f5e839 \ + --hash=sha256:86defa8d248c3fa029da68ce61fe735432b047e32179883bdb1e79ed9bb8195e \ + --hash=sha256:8ac4f9ead4bbd0bc8ab2d318f97d85147167a488be0e08814a37eb2f439d5cf6 \ + --hash=sha256:93530900d14c37a46ce3d6c9e6fd35dbe5f5601bf6b3a5c325c7bffc030344d9 \ + --hash=sha256:9eeb77214afae972a00dee47382d2591abe77bdae166bda672fb1e24702a3860 \ + --hash=sha256:b5f4dfe950ff0479f1f00eda09c18798d4f49b98f4e2006d644b3301682ebdca \ + --hash=sha256:c3391bd8e6de35f6f1140e50aaeb3e2b3d6a9012536ca23ab0d9c35ec18c8a91 \ + --hash=sha256:c880eba5175f4307129784eca96f4e70b88e57aa3f680aeba3bab0e980b0f37d \ + --hash=sha256:cecfefa17042941f94ab54f769c8ce0fe14beff2694e9ac684176a2535bf9714 \ + --hash=sha256:e40211b4923ba5a6dc9769eab704bdb3fbb58d56c5b336d30996c24fcf12aadb \ + --hash=sha256:efc8ad4e6fc4f1752ebfb58aefece8b4e3c4cae940b0994d43649bdfce8d0d4f # via # gcp-releasetool # secretstorage @@ -382,6 +382,7 @@ protobuf==3.20.3 \ # gcp-docuploader # gcp-releasetool # google-api-core + # googleapis-common-protos pyasn1==0.4.8 \ --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \ --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba From 158c00ca1e48f0d8175ed7d2a9b47ab1dc8f288f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 3 Oct 2023 12:50:14 -0500 Subject: [PATCH 52/80] chore: fix flaky system tests (#84) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Remove a mutation that could affect subsequent tests. * Make error message from API more flexible. 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://togithub.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 # 🦕 --- tests/system/small/test_pandas_options.py | 10 +++++----- tests/system/small/test_session.py | 7 +++---- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index 6510c4fa271..956b29ae12f 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -75,7 +75,7 @@ def test_read_gbq_start_sets_session_location( # Now read_gbq* from another location should fail with pytest.raises( google.api_core.exceptions.NotFound, - match=f"404 Not found: Dataset {dataset_id_permanent} was not found in location {tokyo_location}", + match=dataset_id_permanent, ): read_method(query) @@ -100,7 +100,7 @@ def test_read_gbq_start_sets_session_location( # Now read_gbq* from another location should fail with pytest.raises( google.api_core.exceptions.NotFound, - match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -146,7 +146,7 @@ def test_read_gbq_after_session_start_must_comply_with_default_location( # Doing read_gbq* from a table in another location should fail with pytest.raises( google.api_core.exceptions.NotFound, - match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -194,7 +194,7 @@ def test_read_gbq_must_comply_with_set_location_US( # Starting user journey with read_gbq* from another location should fail with pytest.raises( google.api_core.exceptions.NotFound, - match=f"404 Not found: Dataset {dataset_id_permanent_tokyo} was not found in location US", + match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -244,7 +244,7 @@ def test_read_gbq_must_comply_with_set_location_non_US( # Starting user journey with read_gbq* from another location should fail with pytest.raises( google.api_core.exceptions.NotFound, - match=f"404 Not found: Dataset {dataset_id_permanent} was not found in location {tokyo_location}", + match=dataset_id_permanent, ): read_method(query) diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 614c9537649..d2560174e43 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -321,11 +321,10 @@ def test_read_pandas_multi_index(session, scalars_pandas_df_multi_index): def test_read_pandas_rowid_exists_adds_suffix(session, scalars_pandas_df_default_index): - scalars_pandas_df_default_index["rowid"] = np.arange( - scalars_pandas_df_default_index.shape[0] - ) + pandas_df = scalars_pandas_df_default_index.copy() + pandas_df["rowid"] = np.arange(pandas_df.shape[0]) - df = session.read_pandas(scalars_pandas_df_default_index) + df = session.read_pandas(pandas_df) total_order_col = df._block._expr._ordering.total_order_col assert total_order_col and total_order_col.column_id == "rowid_2" From 8f3b5b240f0f28fef92465abc53504e875d7335a Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 3 Oct 2023 15:16:15 -0500 Subject: [PATCH 53/80] fix: avoid `403 response too large to return` error with `read_gbq` and large query results (#77) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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://togithub.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) Towards internal issue 303057336 🦕 --- bigframes/core/__init__.py | 4 +- bigframes/core/io.py | 45 +++++- bigframes/session.py | 155 +++++++++++--------- tests/system/small/ml/test_core.py | 3 +- tests/system/small/ml/test_decomposition.py | 5 +- tests/system/small/test_session.py | 1 + tests/unit/core/test_io.py | 55 +++++++ 7 files changed, 193 insertions(+), 75 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 5e0675fd134..8008c1189a1 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -1198,8 +1198,8 @@ def cached(self, cluster_cols: typing.Sequence[str]) -> ArrayValue: destination = self._session._ibis_to_session_table( ibis_expr, cluster_cols=cluster_cols, api_name="cache" ) - table_expression = self._session.ibis_client.sql( - f"SELECT * FROM `_SESSION`.`{destination.table_id}`" + table_expression = self._session.ibis_client.table( + f"{destination.project}.{destination.dataset_id}.{destination.table_id}" ) new_columns = [table_expression[column] for column in self.column_names] new_hidden_columns = [ diff --git a/bigframes/core/io.py b/bigframes/core/io.py index 3c2e5a25f5a..d47efbdddcf 100644 --- a/bigframes/core/io.py +++ b/bigframes/core/io.py @@ -16,7 +16,8 @@ import datetime import textwrap -from typing import Dict, Union +import types +from typing import Dict, Iterable, Union import google.cloud.bigquery as bigquery @@ -89,6 +90,48 @@ def create_snapshot_sql( ) +# BigQuery REST API returns types in Legacy SQL format +# https://cloud.google.com/bigquery/docs/data-types but we use Standard SQL +# names +# https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types +BQ_STANDARD_TYPES = types.MappingProxyType( + { + "BOOLEAN": "BOOL", + "INTEGER": "INT64", + "FLOAT": "FLOAT64", + } +) + + +def bq_field_to_type_sql(field: bigquery.SchemaField): + if field.mode == "REPEATED": + nested_type = bq_field_to_type_sql( + bigquery.SchemaField( + field.name, field.field_type, mode="NULLABLE", fields=field.fields + ) + ) + return f"ARRAY<{nested_type}>" + + if field.field_type == "RECORD": + nested_fields_sql = ", ".join( + bq_field_to_sql(child_field) for child_field in field.fields + ) + return f"STRUCT<{nested_fields_sql}>" + + type_ = field.field_type + return BQ_STANDARD_TYPES.get(type_, type_) + + +def bq_field_to_sql(field: bigquery.SchemaField): + name = field.name + type_ = bq_field_to_type_sql(field) + return f"`{name}` {type_}" + + +def bq_schema_to_sql(schema: Iterable[bigquery.SchemaField]): + return ", ".join(bq_field_to_sql(field) for field in schema) + + def format_option(key: str, value: Union[bool, str]) -> str: if isinstance(value, bool): return f"{key}=true" if value else f"{key}=false" diff --git a/bigframes/session.py b/bigframes/session.py index 7b827c7dcfc..ac48c977cb0 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -449,13 +449,6 @@ def _query_to_destination( index_cols: List[str], api_name: str, ) -> Tuple[Optional[bigquery.TableReference], Optional[bigquery.QueryJob]]: - # If there are no index columns, then there's no reason to cache to a - # (clustered) session table, as we'll just have to query it again to - # create a default index & ordering. - if not index_cols: - _, query_job = self._start_query(query) - return query_job.destination, query_job - # If a dry_run indicates this is not a query type job, then don't # bother trying to do a CREATE TEMP TABLE ... AS SELECT ... statement. dry_run_config = bigquery.QueryJobConfig() @@ -465,15 +458,24 @@ def _query_to_destination( _, query_job = self._start_query(query) return query_job.destination, query_job - # Make sure we cluster by the index column(s) so that subsequent - # operations are as speedy as they can be. + # Create a table to workaround BigQuery 10 GB query results limit. See: + # internal issue 303057336. + # Since we have a `statement_type == 'SELECT'`, schema should be populated. + schema = typing.cast(Iterable[bigquery.SchemaField], dry_run_job.schema) + temp_table = self._create_session_table_empty(api_name, schema, index_cols) + + job_config = bigquery.QueryJobConfig() + job_config.destination = temp_table + try: - ibis_expr = self.ibis_client.sql(query) - return self._ibis_to_session_table(ibis_expr, index_cols, api_name), None + # Write to temp table to workaround BigQuery 10 GB query results + # limit. See: internal issue 303057336. + _, query_job = self._start_query(query, job_config=job_config) + return query_job.destination, query_job except google.api_core.exceptions.BadRequest: - # Some SELECT statements still aren't compatible with CREATE TEMP - # TABLE ... AS SELECT ... statements. For example, if the query has - # a top-level ORDER BY, this conflicts with our ability to cluster + # Some SELECT statements still aren't compatible with cluster + # tables as the destination. For example, if the query has a + # top-level ORDER BY, this conflicts with our ability to cluster # the table by the index column(s). _, query_job = self._start_query(query) return query_job.destination, query_job @@ -1231,6 +1233,54 @@ def _create_session_table(self) -> bigquery.TableReference: ) return dataset.table(table_name) + def _create_session_table_empty( + self, + api_name: str, + schema: Iterable[bigquery.SchemaField], + cluster_cols: List[str], + ) -> bigquery.TableReference: + # Can't set a table in _SESSION as destination via query job API, so we + # run DDL, instead. + table = self._create_session_table() + schema_sql = bigframes_io.bq_schema_to_sql(schema) + + clusterable_cols = [ + col.name + for col in schema + if col.name in cluster_cols and _can_cluster_bq(col) + ][:_MAX_CLUSTER_COLUMNS] + + if clusterable_cols: + cluster_cols_sql = ", ".join( + f"`{cluster_col}`" for cluster_col in clusterable_cols + ) + cluster_sql = f"CLUSTER BY {cluster_cols_sql}" + else: + cluster_sql = "" + + ddl_text = f""" + CREATE TEMP TABLE + `_SESSION`.`{table.table_id}` + ({schema_sql}) + {cluster_sql} + """ + + job_config = bigquery.QueryJobConfig() + + # Include a label so that Dataplex Lineage can identify temporary + # tables that BigQuery DataFrames creates. Googlers: See internal issue + # 296779699. We're labeling the job instead of the table because + # otherwise we get `BadRequest: 400 OPTIONS on temporary tables are not + # supported`. + job_config.labels = {"source": "bigquery-dataframes-temp"} + job_config.labels["bigframes-api"] = api_name + + _, query_job = self._start_query(ddl_text, job_config=job_config) + + # Use fully-qualified name instead of `_SESSION` name so that the + # created table can be used as the destination table. + return query_job.destination + def _create_sequential_ordering( self, table: ibis_types.Table, @@ -1249,7 +1299,9 @@ def _create_sequential_ordering( cluster_cols=list(index_cols) + [default_ordering_name], api_name=api_name, ) - table = self.ibis_client.sql(f"SELECT * FROM `{table_ref.table_id}`") + table = self.ibis_client.table( + f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" + ) ordering_reference = core.OrderingColumnReference(default_ordering_name) ordering = core.ExpressionOrdering( ordering_value_columns=[ordering_reference], @@ -1264,55 +1316,13 @@ def _ibis_to_session_table( cluster_cols: Iterable[str], api_name: str, ) -> bigquery.TableReference: - clusterable_cols = [ - col for col in cluster_cols if _can_cluster(table[col].type()) - ][:_MAX_CLUSTER_COLUMNS] - return self._query_to_session_table( + desination, _ = self._query_to_destination( self.ibis_client.compile(table), - cluster_cols=clusterable_cols, + index_cols=list(cluster_cols), api_name=api_name, ) - - def _query_to_session_table( - self, - query_text: str, - cluster_cols: Iterable[str], - api_name: str, - ) -> bigquery.TableReference: - if len(list(cluster_cols)) > _MAX_CLUSTER_COLUMNS: - raise ValueError( - f"Too many cluster columns: {list(cluster_cols)}, max {_MAX_CLUSTER_COLUMNS} allowed." - ) - # Can't set a table in _SESSION as destination via query job API, so we - # run DDL, instead. - table = self._create_session_table() - cluster_cols_sql = ", ".join(f"`{cluster_col}`" for cluster_col in cluster_cols) - - # TODO(swast): This might not support multi-statement SQL queries (scripts). - ddl_text = f""" - CREATE TEMP TABLE `_SESSION`.`{table.table_id}` - CLUSTER BY {cluster_cols_sql} - AS {query_text} - """ - - job_config = bigquery.QueryJobConfig() - - # Include a label so that Dataplex Lineage can identify temporary - # tables that BigQuery DataFrames creates. Googlers: See internal issue - # 296779699. We're labeling the job instead of the table because - # otherwise we get `BadRequest: 400 OPTIONS on temporary tables are not - # supported`. - job_config.labels = {"source": "bigquery-dataframes-temp"} - job_config.labels["bigframes-api"] = api_name - - try: - self._start_query( - ddl_text, job_config=job_config - ) # Wait for the job to complete - except google.api_core.exceptions.Conflict: - # Allow query retry to succeed. - pass - return table + # There should always be a destination table for this query type. + return typing.cast(bigquery.TableReference, desination) def remote_function( self, @@ -1494,14 +1504,21 @@ def connect(context: Optional[bigquery_options.BigQueryOptions] = None) -> Sessi return Session(context) -def _can_cluster(ibis_type: ibis_dtypes.DataType): +def _can_cluster_bq(field: bigquery.SchemaField): # https://cloud.google.com/bigquery/docs/clustered-tables # Notably, float is excluded - return ( - ibis_type.is_integer() - or ibis_type.is_string() - or ibis_type.is_decimal() - or ibis_type.is_date() - or ibis_type.is_timestamp() - or ibis_type.is_boolean() + type_ = field.field_type + return type_ in ( + "INTEGER", + "INT64", + "STRING", + "NUMERIC", + "DECIMAL", + "BIGNUMERIC", + "BIGDECIMAL", + "DATE", + "DATETIME", + "TIMESTAMP", + "BOOL", + "BOOLEAN", ) diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index ace943956fe..f911dd7eebc 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -23,6 +23,7 @@ import bigframes from bigframes.ml import core +import tests.system.utils def test_model_eval( @@ -224,7 +225,7 @@ def test_pca_model_principal_component_info(penguins_bqml_pca_model: core.BqmlMo "cumulative_explained_variance_ratio": [0.469357, 0.651283, 0.812383], }, ) - pd.testing.assert_frame_equal( + tests.system.utils.assert_pandas_df_equal_ignore_ordering( result, expected, check_exact=False, diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index c71bbbe3b0c..e31681f4a09 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -15,6 +15,7 @@ import pandas as pd from bigframes.ml import decomposition +import tests.system.utils def test_pca_predict(penguins_pca_model, new_penguins_df): @@ -129,7 +130,7 @@ def test_pca_explained_variance_(penguins_pca_model: decomposition.PCA): "explained_variance": [3.278657, 1.270829, 1.125354], }, ) - pd.testing.assert_frame_equal( + tests.system.utils.assert_pandas_df_equal_ignore_ordering( result, expected, check_exact=False, @@ -148,7 +149,7 @@ def test_pca_explained_variance_ratio_(penguins_pca_model: decomposition.PCA): "explained_variance_ratio": [0.469357, 0.181926, 0.1611], }, ) - pd.testing.assert_frame_equal( + tests.system.utils.assert_pandas_df_equal_ignore_ordering( result, expected, check_exact=False, diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index d2560174e43..53ddfa3c495 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -57,6 +57,7 @@ def test_read_gbq_tokyo( ), pytest.param( """SELECT + t.int64_col + 1 as my_ints, t.float64_col * 2 AS my_floats, CONCAT(t.string_col, "_2") AS my_strings, t.int64_col > 0 AS my_bools, diff --git a/tests/unit/core/test_io.py b/tests/unit/core/test_io.py index c5074f80c2c..afb38a5f753 100644 --- a/tests/unit/core/test_io.py +++ b/tests/unit/core/test_io.py @@ -13,8 +13,10 @@ # limitations under the License. import datetime +from typing import Iterable import google.cloud.bigquery as bigquery +import pytest import bigframes.core.io @@ -47,3 +49,56 @@ def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): # Don't need the project ID for _SESSION tables. assert "my-test-project" not in sql + + +@pytest.mark.parametrize( + ("schema", "expected"), + ( + ( + [bigquery.SchemaField("My Column", "INTEGER")], + "`My Column` INT64", + ), + ( + [ + bigquery.SchemaField("My Column", "INTEGER"), + bigquery.SchemaField("Float Column", "FLOAT"), + bigquery.SchemaField("Bool Column", "BOOLEAN"), + ], + "`My Column` INT64, `Float Column` FLOAT64, `Bool Column` BOOL", + ), + ( + [ + bigquery.SchemaField("My Column", "INTEGER", mode="REPEATED"), + bigquery.SchemaField("Float Column", "FLOAT", mode="REPEATED"), + bigquery.SchemaField("Bool Column", "BOOLEAN", mode="REPEATED"), + ], + "`My Column` ARRAY, `Float Column` ARRAY, `Bool Column` ARRAY", + ), + ( + [ + bigquery.SchemaField( + "My Column", + "RECORD", + mode="REPEATED", + fields=( + bigquery.SchemaField("Float Column", "FLOAT", mode="REPEATED"), + bigquery.SchemaField("Bool Column", "BOOLEAN", mode="REPEATED"), + bigquery.SchemaField( + "Nested Column", + "RECORD", + fields=(bigquery.SchemaField("Int Column", "INTEGER"),), + ), + ), + ), + ], + ( + "`My Column` ARRAY," + + " `Bool Column` ARRAY," + + " `Nested Column` STRUCT<`Int Column` INT64>>>" + ), + ), + ), +) +def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str): + pass From 17afac9ff70a2b93ed70dc7bcce7beb9a53c2ece Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 3 Oct 2023 16:04:14 -0500 Subject: [PATCH 54/80] feat: support STRUCT data type with `Series.struct.field` to extract child fields (#71) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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://togithub.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 # 🦕 --- .kokoro/build.sh | 3 + bigframes/dataframe.py | 10 +- bigframes/dtypes.py | 63 ++++++++++++- bigframes/operations/base.py | 10 +- bigframes/operations/structs.py | 61 ++++++++++++ bigframes/series.py | 5 + noxfile.py | 2 +- tests/system/small/test_dataframe.py | 14 ++- tests/unit/test_dtypes.py | 64 +++++++++++++ .../pandas/core/arrays/__init__.py | 0 .../pandas/core/arrays/arrow/__init__.py | 0 .../pandas/core/arrays/arrow/accessors.py | 94 +++++++++++++++++++ .../bigframes_vendored/sklearn/__init__.py | 0 .../sklearn/ensemble/__init__.py | 0 .../bigframes_vendored/xgboost/__init__.py | 0 15 files changed, 318 insertions(+), 8 deletions(-) create mode 100644 bigframes/operations/structs.py create mode 100644 third_party/bigframes_vendored/pandas/core/arrays/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/arrays/arrow/__init__.py create mode 100644 third_party/bigframes_vendored/pandas/core/arrays/arrow/accessors.py create mode 100644 third_party/bigframes_vendored/sklearn/__init__.py create mode 100644 third_party/bigframes_vendored/sklearn/ensemble/__init__.py create mode 100644 third_party/bigframes_vendored/xgboost/__init__.py diff --git a/.kokoro/build.sh b/.kokoro/build.sh index a0fa4bc787b..58eaa7fedf9 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -26,6 +26,9 @@ 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 diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index b8ae37ee6c0..345afbe6e6b 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -160,7 +160,15 @@ def __init__( columns=columns, # type:ignore dtype=dtype, # type:ignore ) - if pd_dataframe.size < MAX_INLINE_DF_SIZE: + if ( + pd_dataframe.size < MAX_INLINE_DF_SIZE + # TODO(swast): Workaround data types limitation in inline data. + and not any( + dt.pyarrow_dtype + for dt in pd_dataframe.dtypes + if isinstance(dt, pandas.ArrowDtype) + ) + ): self._block = blocks.block_from_local( pd_dataframe, session or bigframes.pandas.get_global_session() ) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 59d3007fab0..46a7a1cb503 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -84,10 +84,10 @@ BIDIRECTIONAL_MAPPINGS: Iterable[Tuple[IbisDtype, Dtype]] = ( (ibis_dtypes.boolean, pd.BooleanDtype()), + (ibis_dtypes.date, pd.ArrowDtype(pa.date32())), (ibis_dtypes.float64, pd.Float64Dtype()), (ibis_dtypes.int64, pd.Int64Dtype()), (ibis_dtypes.string, pd.StringDtype(storage="pyarrow")), - (ibis_dtypes.date, pd.ArrowDtype(pa.date32())), (ibis_dtypes.time, pd.ArrowDtype(pa.time64("us"))), (ibis_dtypes.Timestamp(timezone=None), pd.ArrowDtype(pa.timestamp("us"))), ( @@ -100,6 +100,19 @@ pandas: ibis for ibis, pandas in BIDIRECTIONAL_MAPPINGS } +IBIS_TO_ARROW: Dict[ibis_dtypes.DataType, pa.DataType] = { + ibis_dtypes.boolean: pa.bool_(), + ibis_dtypes.date: pa.date32(), + ibis_dtypes.float64: pa.float64(), + ibis_dtypes.int64: pa.int64(), + ibis_dtypes.string: pa.string(), + ibis_dtypes.time: pa.time64("us"), + ibis_dtypes.Timestamp(timezone=None): pa.timestamp("us"), + ibis_dtypes.Timestamp(timezone="UTC"): pa.timestamp("us", tz="UTC"), +} + +ARROW_TO_IBIS = {arrow: ibis for ibis, arrow in IBIS_TO_ARROW.items()} + IBIS_TO_BIGFRAMES: Dict[ibis_dtypes.DataType, Union[Dtype, np.dtype[Any]]] = { ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS } @@ -148,11 +161,12 @@ def ibis_dtype_to_bigframes_dtype( # Special cases: Ibis supports variations on these types, but currently # our IO returns them as objects. Eventually, we should support them as # ArrowDType (and update the IO accordingly) - if isinstance(ibis_dtype, ibis_dtypes.Array) or isinstance( - ibis_dtype, ibis_dtypes.Struct - ): + if isinstance(ibis_dtype, ibis_dtypes.Array): return np.dtype("O") + if isinstance(ibis_dtype, ibis_dtypes.Struct): + return pd.ArrowDtype(ibis_dtype_to_arrow_dtype(ibis_dtype)) + if ibis_dtype in IBIS_TO_BIGFRAMES: return IBIS_TO_BIGFRAMES[ibis_dtype] elif isinstance(ibis_dtype, ibis_dtypes.Null): @@ -164,6 +178,26 @@ def ibis_dtype_to_bigframes_dtype( ) +def ibis_dtype_to_arrow_dtype(ibis_dtype: ibis_dtypes.DataType) -> pa.DataType: + if isinstance(ibis_dtype, ibis_dtypes.Array): + return pa.list_(ibis_dtype_to_arrow_dtype(ibis_dtype.value_type)) + + if isinstance(ibis_dtype, ibis_dtypes.Struct): + return pa.struct( + [ + (name, ibis_dtype_to_arrow_dtype(dtype)) + for name, dtype in ibis_dtype.fields.items() + ] + ) + + if ibis_dtype in IBIS_TO_ARROW: + return IBIS_TO_ARROW[ibis_dtype] + else: + raise ValueError( + f"Unexpected Ibis data type {ibis_dtype}. {constants.FEEDBACK_LINK}" + ) + + def ibis_value_to_canonical_type(value: ibis_types.Value) -> ibis_types.Value: """Converts an Ibis expression to canonical type. @@ -187,6 +221,24 @@ def ibis_table_to_canonical_types(table: ibis_types.Table) -> ibis_types.Table: return table.select(*casted_columns) +def arrow_dtype_to_ibis_dtype(arrow_dtype: pa.DataType) -> ibis_dtypes.DataType: + if pa.types.is_struct(arrow_dtype): + struct_dtype = typing.cast(pa.StructType, arrow_dtype) + return ibis_dtypes.Struct.from_tuples( + [ + (field.name, arrow_dtype_to_ibis_dtype(field.type)) + for field in struct_dtype + ] + ) + + if arrow_dtype in ARROW_TO_IBIS: + return ARROW_TO_IBIS[arrow_dtype] + else: + raise ValueError( + f"Unexpected Arrow data type {arrow_dtype}. {constants.FEEDBACK_LINK}" + ) + + def bigframes_dtype_to_ibis_dtype( bigframes_dtype: Union[DtypeString, Dtype, np.dtype[Any]] ) -> ibis_dtypes.DataType: @@ -202,6 +254,9 @@ def bigframes_dtype_to_ibis_dtype( Raises: ValueError: If passed a dtype not supported by BigQuery DataFrames. """ + if isinstance(bigframes_dtype, pd.ArrowDtype): + return arrow_dtype_to_ibis_dtype(bigframes_dtype.pyarrow_dtype) + type_string = str(bigframes_dtype) if type_string in BIGFRAMES_STRING_TO_BIGFRAMES: bigframes_dtype = BIGFRAMES_STRING_TO_BIGFRAMES[ diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index add6af57f41..51eaad18b9b 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -86,7 +86,15 @@ def __init__( if pd_series.name is None: # to_frame will set default numeric column label if unnamed, but we do not support int column label, so must rename pd_dataframe = pd_dataframe.set_axis(["unnamed_col"], axis=1) - if pd_dataframe.size < MAX_INLINE_SERIES_SIZE: + if ( + pd_dataframe.size < MAX_INLINE_SERIES_SIZE + # TODO(swast): Workaround data types limitation in inline data. + and not any( + dt.pyarrow_dtype + for dt in pd_dataframe.dtypes + if isinstance(dt, pd.ArrowDtype) + ) + ): self._block = blocks.block_from_local( pd_dataframe, session or bigframes.pandas.get_global_session() ) diff --git a/bigframes/operations/structs.py b/bigframes/operations/structs.py new file mode 100644 index 00000000000..80d51115d01 --- /dev/null +++ b/bigframes/operations/structs.py @@ -0,0 +1,61 @@ +# 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. + +from __future__ import annotations + +import typing + +import ibis.expr.types as ibis_types + +import bigframes.dataframe +import bigframes.operations +import bigframes.operations.base +import bigframes.series +import third_party.bigframes_vendored.pandas.core.arrays.arrow.accessors as vendoracessors + + +class StructField(bigframes.operations.UnaryOp): + def __init__(self, name_or_index: str | int): + self._name_or_index = name_or_index + + def _as_ibis(self, x: ibis_types.Value): + struct_value = typing.cast(ibis_types.StructValue, x) + if isinstance(self._name_or_index, str): + name = self._name_or_index + else: + name = struct_value.names[self._name_or_index] + return struct_value[name].name(name) + + +class StructAccessor( + bigframes.operations.base.SeriesMethods, vendoracessors.StructAccessor +): + __doc__ = vendoracessors.StructAccessor.__doc__ + + def field(self, name_or_index: str | int) -> bigframes.series.Series: + series = self._apply_unary_op(StructField(name_or_index)) + if isinstance(name_or_index, str): + name = name_or_index + else: + struct_field = self._dtype.pyarrow_dtype[name_or_index] + name = struct_field.name + return series.rename(name) + + def explode(self) -> bigframes.dataframe.DataFrame: + import bigframes.pandas + + pa_type = self._dtype.pyarrow_dtype + return bigframes.pandas.concat( + [self.field(i) for i in range(pa_type.num_fields)], axis="columns" + ) diff --git a/bigframes/series.py b/bigframes/series.py index 47298d59f50..84d737210ef 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -51,6 +51,7 @@ import bigframes.operations.base import bigframes.operations.datetimes as dt import bigframes.operations.strings as strings +import bigframes.operations.structs as structs import third_party.bigframes_vendored.pandas.core.series as vendored_pandas_series LevelType = typing.Union[str, int] @@ -118,6 +119,10 @@ def query_job(self) -> Optional[bigquery.QueryJob]: self._set_internal_query_job(self._compute_dry_run()) return self._query_job + @property + def struct(self) -> structs.StructAccessor: + return structs.StructAccessor(self._block) + def _set_internal_query_job(self, query_job: bigquery.QueryJob): self._query_job = query_job diff --git a/noxfile.py b/noxfile.py index 81ed9c2b2c4..a113e1fcde4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -362,7 +362,7 @@ def doctest(session: nox.sessions.Session): run_system( session=session, prefix_name="doctest", - extra_pytest_options=("--doctest-modules",), + extra_pytest_options=("--doctest-modules", "third_party"), test_folder="bigframes", check_cov=True, ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a27d8fd487c..2880932d103 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -917,7 +917,19 @@ def test_get_dtypes_array_struct(session): dtypes = df.dtypes pd.testing.assert_series_equal( dtypes, - pd.Series({"array_column": np.dtype("O"), "struct_column": np.dtype("O")}), + pd.Series( + { + "array_column": np.dtype("O"), + "struct_column": pd.ArrowDtype( + pa.struct( + [ + ("string_field", pa.string()), + ("float_field", pa.float64()), + ] + ) + ), + } + ), ) diff --git a/tests/unit/test_dtypes.py b/tests/unit/test_dtypes.py index bb8ae570dca..3baff2e1f58 100644 --- a/tests/unit/test_dtypes.py +++ b/tests/unit/test_dtypes.py @@ -85,6 +85,70 @@ def test_ibis_float32_raises_unexpected_datatype(): bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_dtypes.float32) +IBIS_ARROW_DTYPES = ( + (ibis_dtypes.boolean, pa.bool_()), + (ibis_dtypes.date, pa.date32()), + (ibis_dtypes.Timestamp(), pa.timestamp("us")), + (ibis_dtypes.float64, pa.float64()), + ( + ibis_dtypes.Timestamp(timezone="UTC"), + pa.timestamp("us", tz="UTC"), + ), + ( + ibis_dtypes.Struct.from_tuples( + [ + ("name", ibis_dtypes.string()), + ("version", ibis_dtypes.int64()), + ] + ), + pa.struct( + [ + ("name", pa.string()), + ("version", pa.int64()), + ] + ), + ), + ( + ibis_dtypes.Struct.from_tuples( + [ + ( + "nested", + ibis_dtypes.Struct.from_tuples( + [ + ("field", ibis_dtypes.string()), + ] + ), + ), + ] + ), + pa.struct( + [ + ( + "nested", + pa.struct( + [ + ("field", pa.string()), + ] + ), + ), + ] + ), + ), +) + + +@pytest.mark.parametrize(("ibis_dtype", "arrow_dtype"), IBIS_ARROW_DTYPES) +def test_arrow_dtype_to_ibis_dtype(ibis_dtype, arrow_dtype): + result = bigframes.dtypes.arrow_dtype_to_ibis_dtype(arrow_dtype) + assert result == ibis_dtype + + +@pytest.mark.parametrize(("ibis_dtype", "arrow_dtype"), IBIS_ARROW_DTYPES) +def test_ibis_dtype_to_arrow_dtype(ibis_dtype, arrow_dtype): + result = bigframes.dtypes.ibis_dtype_to_arrow_dtype(ibis_dtype) + assert result == arrow_dtype + + @pytest.mark.parametrize( ["bigframes_dtype", "ibis_dtype"], [ diff --git a/third_party/bigframes_vendored/pandas/core/arrays/__init__.py b/third_party/bigframes_vendored/pandas/core/arrays/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/arrays/arrow/__init__.py b/third_party/bigframes_vendored/pandas/core/arrays/arrow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/pandas/core/arrays/arrow/accessors.py b/third_party/bigframes_vendored/pandas/core/arrays/arrow/accessors.py new file mode 100644 index 00000000000..8e3ea06a3d4 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/arrays/arrow/accessors.py @@ -0,0 +1,94 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/arrays/arrow/accessors.py +"""Accessors for arrow-backed data.""" + +from __future__ import annotations + +from bigframes import constants + + +class StructAccessor: + """ + Accessor object for structured data properties of the Series values. + """ + + def field(self, name_or_index: str | int): + """ + Extract a child field of a struct as a Series. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import pyarrow as pa + >>> bpd.options.display.progress_bar = None + >>> s = bpd.Series( + ... [ + ... {"version": 1, "project": "pandas"}, + ... {"version": 2, "project": "pandas"}, + ... {"version": 1, "project": "numpy"}, + ... ], + ... dtype=bpd.ArrowDtype(pa.struct( + ... [("version", pa.int64()), ("project", pa.string())] + ... )) + ... ) + + Extract by field name. + + >>> s.struct.field("project") + 0 pandas + 1 pandas + 2 numpy + Name: project, dtype: string + + Extract by field index. + + >>> s.struct.field(0) + 0 1 + 1 2 + 2 1 + Name: version, dtype: Int64 + + Args: + name_or_index: + Name (str) or index (int) of the child field to extract. + + Returns: + Series: + The data corresponding to the selected child field. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def explode(self): + """ + Extract all child fields of a struct as a DataFrame. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import pyarrow as pa + >>> bpd.options.display.progress_bar = None + >>> s = bpd.Series( + ... [ + ... {"version": 1, "project": "pandas"}, + ... {"version": 2, "project": "pandas"}, + ... {"version": 1, "project": "numpy"}, + ... ], + ... dtype=bpd.ArrowDtype(pa.struct( + ... [("version", pa.int64()), ("project", pa.string())] + ... )) + ... ) + + Extract all child fields. + + >>> s.struct.explode() + version project + 0 1 pandas + 1 2 pandas + 2 1 numpy + + [3 rows x 2 columns] + + Returns: + DataFrame: + The data corresponding to all child fields. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/__init__.py b/third_party/bigframes_vendored/sklearn/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/sklearn/ensemble/__init__.py b/third_party/bigframes_vendored/sklearn/ensemble/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/third_party/bigframes_vendored/xgboost/__init__.py b/third_party/bigframes_vendored/xgboost/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 781307ec22d31a7657f8ee5c6eedc0e419450ccd Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 3 Oct 2023 16:12:52 -0700 Subject: [PATCH 55/80] feat: add idxmin, idxmax to series, dataframe (#74) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add idxmin, idxmax to series, dataframe * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot --- bigframes/core/block_transforms.py | 51 +++++++++++++++++++ bigframes/dataframe.py | 6 +++ bigframes/series.py | 28 ++++++++++ tests/system/small/test_dataframe.py | 28 ++++++++++ tests/system/small/test_multiindex.py | 11 ++++ tests/system/small/test_series.py | 12 +++++ .../bigframes_vendored/pandas/core/frame.py | 22 ++++++++ .../bigframes_vendored/pandas/core/series.py | 40 +++++++++------ 8 files changed, 182 insertions(+), 16 deletions(-) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 6e99a7c7740..30c79029817 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -17,6 +17,7 @@ import pandas as pd +import bigframes.constants as constants import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.ordering as ordering @@ -576,3 +577,53 @@ def align_columns( left_final = left_block.select_columns(left_column_ids) right_final = right_block.select_columns(right_column_ids) return left_final, right_final + + +def idxmin(block: blocks.Block) -> blocks.Block: + return _idx_extrema(block, "min") + + +def idxmax(block: blocks.Block) -> blocks.Block: + return _idx_extrema(block, "max") + + +def _idx_extrema( + block: blocks.Block, min_or_max: typing.Literal["min", "max"] +) -> blocks.Block: + if len(block.index_columns) != 1: + # TODO: Need support for tuple dtype + raise NotImplementedError( + f"idxmin not support for multi-index. {constants.FEEDBACK_LINK}" + ) + + original_block = block + result_cols = [] + for value_col in original_block.value_columns: + direction = ( + ordering.OrderingDirection.ASC + if min_or_max == "min" + else ordering.OrderingDirection.DESC + ) + # Have to find the min for each + order_refs = [ + ordering.OrderingColumnReference(value_col, direction), + *[ + ordering.OrderingColumnReference(idx_col) + for idx_col in original_block.index_columns + ], + ] + window_spec = core.WindowSpec(ordering=order_refs) + idx_col = original_block.index_columns[0] + block, result_col = block.apply_window_op( + idx_col, agg_ops.first_op, window_spec + ) + result_cols.append(result_col) + + block = block.select_columns(result_cols).with_column_labels( + original_block.column_labels + ) + # Stack the entire column axis to produce single-column result + # Assumption: uniform dtype for stackability + return block.aggregate_all_and_stack( + agg_ops.AnyValueOp(), dtype=block.dtypes[0] + ).with_column_labels([original_block.index.name]) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 345afbe6e6b..eea8beb130a 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1642,6 +1642,12 @@ def agg( aggregate = agg + def idxmin(self) -> bigframes.series.Series: + return bigframes.series.Series(block_ops.idxmin(self._block)) + + def idxmax(self) -> bigframes.series.Series: + return bigframes.series.Series(block_ops.idxmax(self._block)) + def describe(self) -> DataFrame: df_numeric = self._drop_non_numeric(keep_bool=False) if len(df_numeric.columns) == 0: diff --git a/bigframes/series.py b/bigframes/series.py index 84d737210ef..8815a6abdee 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -887,6 +887,34 @@ def argmin(self) -> int: scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] ) + def idxmax(self) -> blocks.Label: + block = self._block.order_by( + [ + OrderingColumnReference( + self._value_column, direction=OrderingDirection.DESC + ), + *[ + OrderingColumnReference(idx_col) + for idx_col in self._block.index_columns + ], + ] + ) + block = block.slice(0, 1) + return indexes.Index._from_block(block).to_pandas()[0] + + def idxmin(self) -> blocks.Label: + block = self._block.order_by( + [ + OrderingColumnReference(self._value_column), + *[ + OrderingColumnReference(idx_col) + for idx_col in self._block.index_columns + ], + ] + ) + block = block.slice(0, 1) + return indexes.Index._from_block(block).to_pandas()[0] + @property def is_monotonic_increasing(self) -> bool: return typing.cast( diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 2880932d103..f3e7a0c0e9d 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1292,6 +1292,34 @@ def test_df_update(overwrite, filter_func): pd.testing.assert_frame_equal(bf_df1.to_pandas(), pd_df1) +def test_df_idxmin(): + pd_df = pd.DataFrame( + {"a": [1, 2, 3], "b": [7, None, 3], "c": [4, 4, 4]}, index=["x", "y", "z"] + ) + bf_df = dataframe.DataFrame(pd_df) + + bf_result = bf_df.idxmin().to_pandas() + pd_result = pd_df.idxmin() + + pd.testing.assert_series_equal( + bf_result, pd_result, check_index_type=False, check_dtype=False + ) + + +def test_df_idxmax(): + pd_df = pd.DataFrame( + {"a": [1, 2, 3], "b": [7, None, 3], "c": [4, 4, 4]}, index=["x", "y", "z"] + ) + bf_df = dataframe.DataFrame(pd_df) + + bf_result = bf_df.idxmax().to_pandas() + pd_result = pd_df.idxmax() + + pd.testing.assert_series_equal( + bf_result, pd_result, check_index_type=False, check_dtype=False + ) + + @pytest.mark.parametrize( ("join", "axis"), [ diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 10f2a74b21b..19f1c557ef2 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -41,6 +41,17 @@ def test_reset_multi_index(scalars_df_index, scalars_pandas_df_index): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_series_multi_index_idxmin(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.set_index(["bool_col", "int64_too"])[ + "float64_col" + ].idxmin() + pd_result = scalars_pandas_df_index.set_index(["bool_col", "int64_too"])[ + "float64_col" + ].idxmin() + + assert bf_result == pd_result + + def test_binop_series_series_matching_multi_indices( scalars_df_index, scalars_pandas_df_index ): diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 588dcc2c83a..645638fb3aa 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2468,6 +2468,18 @@ def test_argmax(scalars_df_index, scalars_pandas_df_index): assert bf_result == pd_result +def test_series_idxmin(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.string_col.idxmin() + pd_result = scalars_pandas_df_index.string_col.idxmin() + assert bf_result == pd_result + + +def test_series_idxmax(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.int64_too.idxmax() + pd_result = scalars_pandas_df_index.int64_too.idxmax() + assert bf_result == pd_result + + def test_getattr_attribute_error_when_pandas_has(scalars_df_index): # asof is implemented in pandas but not in bigframes with pytest.raises(AttributeError): diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 11aaf84b14e..17d941fbdd0 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -1805,6 +1805,28 @@ def nsmallest(self, n: int, columns, keep: str = "first"): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def idxmin(self): + """ + Return index of first occurrence of minimum over requested axis. + + NA/null values are excluded. + + Returns: + Series: Indexes of minima along the specified axis. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def idxmax(self): + """ + Return index of first occurrence of maximum over requested axis. + + NA/null values are excluded. + + Returns: + Series: Indexes of maxima along the specified axis. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def nunique(self): """ Count number of distinct elements in specified axis. diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index d58c1ccc3b4..a41a3454ca4 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -8,7 +8,6 @@ import numpy as np from pandas._libs import lib from pandas._typing import Axis, FilePath, NaPosition, WriteBuffer -import pandas.io.formats.format as fmt from bigframes import constants from third_party.bigframes_vendored.pandas.core.generic import NDFrame @@ -151,21 +150,6 @@ def to_string( str or None: String representation of Series if ``buf=None``, otherwise None. """ - formatter = fmt.SeriesFormatter( - self, - name=name, - length=length, - header=header, - index=index, - dtype=dtype, - na_rep=na_rep, - float_format=float_format, - min_rows=min_rows, - max_rows=max_rows, - ) - result = formatter.to_string() - - # catch contract violations raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def to_markdown( @@ -475,6 +459,30 @@ def duplicated(self, keep="first") -> Series: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def idxmin(self) -> Hashable: + """ + Return the row label of the minimum value. + + If multiple values equal the minimum, the first row label with that + value is returned. + + Returns: + Index: Label of the minimum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def idxmax(self) -> Hashable: + """ + Return the row label of the maximum value. + + If multiple values equal the maximum, the first row label with that + value is returned. + + Returns: + Index: Label of the maximum value. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def round(self, decimals: int = 0) -> Series: """ Round each value in a Series to the given number of decimals. From fff3d45f03ffbc7bb23143a1572e3dd157463ca9 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 3 Oct 2023 18:57:06 -0700 Subject: [PATCH 56/80] fix: change return type of `Series.loc[scalar]` (#40) * bug: change return type of `Series.loc[scalar]` Change-Id: Id60a7da3021972da5c8a28fb8f3620e10643c0ed * add scalar case and update return types * remove unneeded iloc in series getitem test * fix test_series_get_with_default_index * Run query manual for clarity/redundance --------- Co-authored-by: Henry J Solberg --- bigframes/core/indexers.py | 81 ++++++++++++++++++---------- bigframes/ml/model_selection.py | 8 +-- tests/system/small/test_dataframe.py | 2 +- tests/system/small/test_series.py | 6 +-- 4 files changed, 62 insertions(+), 35 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index e22a5aed857..1a88b2abd60 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -15,7 +15,7 @@ from __future__ import annotations import typing -from typing import Tuple +from typing import Tuple, Union import ibis import pandas as pd @@ -29,20 +29,19 @@ import bigframes.series if typing.TYPE_CHECKING: - LocSingleKey = typing.Union[bigframes.series.Series, indexes.Index, slice] + LocSingleKey = Union[ + bigframes.series.Series, indexes.Index, slice, bigframes.core.scalar.Scalar + ] class LocSeriesIndexer: def __init__(self, series: bigframes.series.Series): self._series = series - def __getitem__(self, key) -> bigframes.series.Series: - """ - Only indexing by a boolean bigframes.series.Series or list of index entries is currently supported - """ - return typing.cast( - bigframes.series.Series, _loc_getitem_series_or_dataframe(self._series, key) - ) + def __getitem__( + self, key + ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: + return _loc_getitem_series_or_dataframe(self._series, key) def __setitem__(self, key, value) -> None: # TODO(swast): support MultiIndex @@ -84,7 +83,7 @@ def __init__(self, series: bigframes.series.Series): def __getitem__( self, key - ) -> bigframes.core.scalar.Scalar | bigframes.series.Series: + ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: """ Index series using integer offsets. Currently supports index by key type: @@ -103,13 +102,17 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame): self._dataframe = dataframe @typing.overload - def __getitem__(self, key: LocSingleKey) -> bigframes.dataframe.DataFrame: + def __getitem__( + self, key: LocSingleKey + ) -> Union[bigframes.dataframe.DataFrame, pd.Series]: ... # Technically this is wrong since we can have duplicate column labels, but # this is expected to be rare. @typing.overload - def __getitem__(self, key: Tuple[LocSingleKey, str]) -> bigframes.series.Series: + def __getitem__( + self, key: Tuple[LocSingleKey, str] + ) -> Union[bigframes.series.Series, bigframes.core.scalar.Scalar]: ... def __getitem__(self, key): @@ -173,7 +176,7 @@ class ILocDataFrameIndexer: def __init__(self, dataframe: bigframes.dataframe.DataFrame): self._dataframe = dataframe - def __getitem__(self, key) -> bigframes.dataframe.DataFrame | pd.Series: + def __getitem__(self, key) -> Union[bigframes.dataframe.DataFrame, pd.Series]: """ Index dataframe using integer offsets. Currently supports index by key type: @@ -188,21 +191,26 @@ def __getitem__(self, key) -> bigframes.dataframe.DataFrame | pd.Series: @typing.overload def _loc_getitem_series_or_dataframe( series_or_dataframe: bigframes.series.Series, key -) -> bigframes.series.Series: +) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: ... @typing.overload def _loc_getitem_series_or_dataframe( series_or_dataframe: bigframes.dataframe.DataFrame, key -) -> bigframes.dataframe.DataFrame: +) -> Union[bigframes.dataframe.DataFrame, pd.Series]: ... def _loc_getitem_series_or_dataframe( - series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, + series_or_dataframe: Union[bigframes.dataframe.DataFrame, bigframes.series.Series], key: LocSingleKey, -) -> bigframes.dataframe.DataFrame | bigframes.series.Series: +) -> Union[ + bigframes.dataframe.DataFrame, + bigframes.series.Series, + pd.Series, + bigframes.core.scalar.Scalar, +]: if isinstance(key, bigframes.series.Series) and key.dtype == "boolean": return series_or_dataframe[key] elif isinstance(key, bigframes.series.Series): @@ -222,7 +230,7 @@ def _loc_getitem_series_or_dataframe( # TODO(henryjsolberg): support MultiIndex if len(key) == 0: # type: ignore return typing.cast( - typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], + Union[bigframes.dataframe.DataFrame, bigframes.series.Series], series_or_dataframe.iloc[0:0], ) @@ -258,11 +266,22 @@ def _loc_getitem_series_or_dataframe( ) keys_df = keys_df.set_index(index_name, drop=True) keys_df.index.name = None - return _perform_loc_list_join(series_or_dataframe, keys_df) + result = _perform_loc_list_join(series_or_dataframe, keys_df) + pandas_result = result.to_pandas() + # although loc[scalar_key] returns multiple results when scalar_key + # is not unique, we download the results here and return the computed + # individual result (as a scalar or pandas series) when the key is unique, + # since we expect unique index keys to be more common. loc[[scalar_key]] + # can be used to retrieve one-item DataFrames or Series. + if len(pandas_result) == 1: + return pandas_result.iloc[0] + # when the key is not unique, we return a bigframes data type + # as usual for methods that return dataframes/series + return result else: raise TypeError( - "Invalid argument type. loc currently only supports indexing with a " - "boolean bigframes Series, a list of index entries or a single index entry. " + "Invalid argument type. Expected bigframes.Series, bigframes.Index, " + "list, : (empty slice), or scalar. " f"{constants.FEEDBACK_LINK}" ) @@ -284,9 +303,9 @@ def _perform_loc_list_join( def _perform_loc_list_join( - series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, + series_or_dataframe: Union[bigframes.dataframe.DataFrame, bigframes.series.Series], keys_df: bigframes.dataframe.DataFrame, -) -> bigframes.series.Series | bigframes.dataframe.DataFrame: +) -> Union[bigframes.series.Series, bigframes.dataframe.DataFrame]: # right join based on the old index so that the matching rows from the user's # original dataframe will be duplicated and reordered appropriately original_index_names = series_or_dataframe.index.names @@ -309,20 +328,26 @@ def _perform_loc_list_join( @typing.overload def _iloc_getitem_series_or_dataframe( series_or_dataframe: bigframes.series.Series, key -) -> bigframes.series.Series | bigframes.core.scalar.Scalar: +) -> Union[bigframes.series.Series, bigframes.core.scalar.Scalar]: ... @typing.overload def _iloc_getitem_series_or_dataframe( series_or_dataframe: bigframes.dataframe.DataFrame, key -) -> bigframes.dataframe.DataFrame | pd.Series: +) -> Union[bigframes.dataframe.DataFrame, pd.Series]: ... def _iloc_getitem_series_or_dataframe( - series_or_dataframe: bigframes.dataframe.DataFrame | bigframes.series.Series, key -) -> bigframes.dataframe.DataFrame | bigframes.series.Series | bigframes.core.scalar.Scalar | pd.Series: + series_or_dataframe: Union[bigframes.dataframe.DataFrame, bigframes.series.Series], + key, +) -> Union[ + bigframes.dataframe.DataFrame, + bigframes.series.Series, + bigframes.core.scalar.Scalar, + pd.Series, +]: if isinstance(key, int): internal_slice_result = series_or_dataframe._slice(key, key + 1, 1) result_pd_df = internal_slice_result.to_pandas() @@ -334,7 +359,7 @@ def _iloc_getitem_series_or_dataframe( elif pd.api.types.is_list_like(key): if len(key) == 0: return typing.cast( - typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series], + Union[bigframes.dataframe.DataFrame, bigframes.series.Series], series_or_dataframe.iloc[0:0], ) df = series_or_dataframe diff --git a/bigframes/ml/model_selection.py b/bigframes/ml/model_selection.py index 110cbcf4934..443b9e7be6e 100644 --- a/bigframes/ml/model_selection.py +++ b/bigframes/ml/model_selection.py @@ -17,6 +17,7 @@ https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection.""" +import typing from typing import List, Union from bigframes.ml import utils @@ -79,9 +80,10 @@ def train_test_split( train_index = split_dfs[0].index test_index = split_dfs[1].index - split_dfs += [ - df.loc[index] for df in dfs[1:] for index in (train_index, test_index) - ] + split_dfs += typing.cast( + List[bpd.DataFrame], + [df.loc[index] for df in dfs[1:] for index in (train_index, test_index)], + ) # convert back to Series. results: List[Union[bpd.DataFrame, bpd.Series]] = [] diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index f3e7a0c0e9d..b8616a54d65 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2081,7 +2081,7 @@ def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index bf_result = scalars_df_index.loc[index] pd_result = scalars_pandas_df_index.loc[index] pd.testing.assert_series_equal( - bf_result.to_pandas().iloc[0, :], + bf_result, pd_result, ) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 645638fb3aa..8c1c36720bf 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -118,7 +118,7 @@ def test_series_get_with_default_index(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = scalars_df[col_name].get(key) pd_result = scalars_pandas_df[col_name].get(key) - assert bf_result.to_pandas().iloc[0] == pd_result + assert bf_result == pd_result @pytest.mark.parametrize( @@ -157,7 +157,7 @@ def test_series___getitem___with_default_index(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_result = scalars_df[col_name][key] pd_result = scalars_pandas_df[col_name][key] - assert bf_result.to_pandas().iloc[0] == pd_result + assert bf_result == pd_result @pytest.mark.parametrize( @@ -2652,7 +2652,7 @@ def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index index = -2345 bf_result = scalars_df_index.date_col.loc[index] pd_result = scalars_pandas_df_index.date_col.loc[index] - assert bf_result.to_pandas().iloc[0] == pd_result + assert bf_result == pd_result def test_series_bool_interpretation_error(scalars_df_index): From 24c625638984f6a84191c7a4c8ac9fb6c3cf1dca Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Wed, 4 Oct 2023 07:20:45 -0700 Subject: [PATCH 57/80] feat: add ml.preprocessing.KBinsDiscretizer (#81) * feat: add ml.preprocessing.KBinsDiscretizer * fix: address all the comments * fix: address additional comments * fix: fix the failed test * Empty commit * Trigger Kokoro --- bigframes/ml/compose.py | 13 +- bigframes/ml/pipeline.py | 12 +- bigframes/ml/preprocessing.py | 152 ++++++++++++++++-- bigframes/ml/sql.py | 9 ++ tests/system/large/ml/test_pipeline.py | 46 ++++++ tests/system/small/ml/test_preprocessing.py | 131 ++++++++++++++- tests/unit/ml/test_compose.py | 26 +++ tests/unit/ml/test_sql.py | 7 + .../sklearn/preprocessing/_discretization.py | 47 ++++++ 9 files changed, 426 insertions(+), 17 deletions(-) create mode 100644 third_party/bigframes_vendored/sklearn/preprocessing/_discretization.py diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index 9effbf1968c..bf046ff6914 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -31,6 +31,7 @@ preprocessing.StandardScaler, preprocessing.MaxAbsScaler, preprocessing.MinMaxScaler, + preprocessing.KBinsDiscretizer, preprocessing.LabelEncoder, ] @@ -91,18 +92,24 @@ def transformers_( return result - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql( + self, + columns: List[str], + X: bpd.DataFrame, + ) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: columns (List[str]): a list of column names to transform + X (bpd.DataFrame): + The Dataframe with training data. Returns: a list of tuples of (sql_expression, output_name)""" return [ - transformer._compile_to_sql([column])[0] + transformer._compile_to_sql([column], X=X)[0] for column in columns for _, transformer, target_column in self.transformers_ if column == target_column @@ -115,7 +122,7 @@ def fit( ) -> ColumnTransformer: (X,) = utils.convert_to_dataframe(X) - compiled_transforms = self._compile_to_sql(X.columns.tolist()) + compiled_transforms = self._compile_to_sql(X.columns.tolist(), X) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] self._bqml_model = self._bqml_model_factory.create_model( diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index ac02c391129..ad0b3fae111 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -52,6 +52,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): preprocessing.OneHotEncoder, preprocessing.MaxAbsScaler, preprocessing.MinMaxScaler, + preprocessing.KBinsDiscretizer, preprocessing.LabelEncoder, ), ): @@ -93,7 +94,7 @@ def fit( ) -> Pipeline: (X,) = utils.convert_to_dataframe(X) - compiled_transforms = self._transform._compile_to_sql(X.columns.tolist()) + compiled_transforms = self._transform._compile_to_sql(X.columns.tolist(), X=X) transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] if y is not None: @@ -151,6 +152,7 @@ def _extract_as_column_transformer( preprocessing.StandardScaler, preprocessing.MaxAbsScaler, preprocessing.MinMaxScaler, + preprocessing.KBinsDiscretizer, preprocessing.LabelEncoder, ], Union[str, List[str]], @@ -190,6 +192,13 @@ def _extract_as_column_transformer( *preprocessing.MinMaxScaler._parse_from_sql(transform_sql), ) ) + elif transform_sql.startswith("ML.BUCKETIZE"): + transformers.append( + ( + "k_bins_discretizer", + *preprocessing.KBinsDiscretizer._parse_from_sql(transform_sql), + ) + ) elif transform_sql.startswith("ML.LABEL_ENCODER"): transformers.append( ( @@ -213,6 +222,7 @@ def _merge_column_transformer( preprocessing.OneHotEncoder, preprocessing.MaxAbsScaler, preprocessing.MinMaxScaler, + preprocessing.KBinsDiscretizer, preprocessing.LabelEncoder, ]: """Try to merge the column transformer to a simple transformer.""" diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index caf4657a633..5f44d402184 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -23,6 +23,7 @@ from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.preprocessing._data +import third_party.bigframes_vendored.sklearn.preprocessing._discretization import third_party.bigframes_vendored.sklearn.preprocessing._encoder import third_party.bigframes_vendored.sklearn.preprocessing._label @@ -44,12 +45,15 @@ def __init__(self): def __eq__(self, other: Any) -> bool: return type(other) is StandardScaler and self._bqml_model == other._bqml_model - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: - columns: a list of column names to transform + columns: + a list of column names to transform. + X (default None): + Ignored. Returns: a list of tuples of (sql_expression, output_name)""" return [ @@ -124,12 +128,15 @@ def __init__(self): def __eq__(self, other: Any) -> bool: return type(other) is MaxAbsScaler and self._bqml_model == other._bqml_model - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: - columns: a list of column names to transform + columns: + a list of column names to transform. + X (default None): + Ignored. Returns: a list of tuples of (sql_expression, output_name)""" return [ @@ -204,12 +211,15 @@ def __init__(self): def __eq__(self, other: Any) -> bool: return type(other) is MinMaxScaler and self._bqml_model == other._bqml_model - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: - columns: a list of column names to transform + columns: + a list of column names to transform. + X (default None): + Ignored. Returns: a list of tuples of (sql_expression, output_name)""" return [ @@ -267,6 +277,124 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +class KBinsDiscretizer( + base.Transformer, + third_party.bigframes_vendored.sklearn.preprocessing._discretization.KBinsDiscretizer, +): + __doc__ = ( + third_party.bigframes_vendored.sklearn.preprocessing._discretization.KBinsDiscretizer.__doc__ + ) + + def __init__( + self, + n_bins: int = 5, + strategy: Literal["uniform", "quantile"] = "quantile", + ): + if strategy != "uniform": + raise NotImplementedError( + f"Only strategy = 'uniform' is supported now, input is {strategy}." + ) + if n_bins < 2: + raise ValueError( + f"n_bins has to be larger than or equal to 2, input is {n_bins}." + ) + self.n_bins = n_bins + self.strategy = strategy + self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return ( + type(other) is KBinsDiscretizer + and self.n_bins == other.n_bins + and self._bqml_model == other._bqml_model + ) + + def _compile_to_sql( + self, + columns: List[str], + X: bpd.DataFrame, + ) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: + a list of column names to transform + X: + The Dataframe with training data. + + Returns: a list of tuples of (sql_expression, output_name)""" + array_split_points = {} + if self.strategy == "uniform": + for column in columns: + min_value = X[column].min() + max_value = X[column].max() + bin_size = (max_value - min_value) / self.n_bins + array_split_points[column] = [ + min_value + i * bin_size for i in range(self.n_bins - 1) + ] + + return [ + ( + self._base_sql_generator.ml_bucketize( + column, array_split_points[column], f"kbinsdiscretizer_{column}" + ), + f"kbinsdiscretizer_{column}", + ) + for column in columns + ] + + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[KBinsDiscretizer, str]: + """Parse SQL to tuple(KBinsDiscretizer, column_label). + + Args: + sql: SQL string of format "ML.BUCKETIZE({col_label}, array_split_points, FALSE) OVER()" + + Returns: + tuple(KBinsDiscretizer, column_label)""" + s = sql[sql.find("(") + 1 : sql.find(")")] + array_split_points = s[s.find("[") + 1 : s.find("]")] + col_label = s[: s.find(",")] + n_bins = array_split_points.count(",") + 2 + return cls(n_bins, "uniform"), col_label + + def fit( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> KBinsDiscretizer: + (X,) = utils.convert_to_dataframe(X) + + compiled_transforms = self._compile_to_sql(X.columns.tolist(), X) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = self._bqml_model_factory.create_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + return self + + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.transform(X) + return typing.cast( + bpd.DataFrame, + df[self._output_names], + ) + + class OneHotEncoder( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, @@ -308,13 +436,15 @@ def __eq__(self, other: Any) -> bool: and self.max_categories == other.max_categories ) - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: columns: - a list of column names to transform + a list of column names to transform. + X (default None): + Ignored. Returns: a list of tuples of (sql_expression, output_name)""" @@ -432,13 +562,15 @@ def __eq__(self, other: Any) -> bool: and self.max_categories == other.max_categories ) - def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]: + def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]: """Compile this transformer to a list of SQL expressions that can be included in a BQML TRANSFORM clause Args: columns: - a list of column names to transform + a list of column names to transform. + X (default None): + Ignored. Returns: a list of tuples of (sql_expression, output_name)""" diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 57c8ba672af..601b2710999 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -85,6 +85,15 @@ def ml_min_max_scaler(self, numeric_expr_sql: str, name: str) -> str: """Encode ML.MIN_MAX_SCALER for BQML""" return f"""ML.MIN_MAX_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_bucketize( + self, + numeric_expr_sql: str, + array_split_points: Iterable[Union[int, float]], + name: str, + ) -> str: + """Encode ML.MIN_MAX_SCALER for BQML""" + return f"""ML.BUCKETIZE({numeric_expr_sql}, {array_split_points}, FALSE) AS {name}""" + def ml_one_hot_encoder( self, numeric_expr_sql: str, diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 34a2ca01011..9294740dd67 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -580,6 +580,11 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind preprocessing.MinMaxScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -657,6 +662,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.MinMaxScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -696,9 +706,19 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id ("standard_scaler", preprocessing.StandardScaler(), "culmen_length_mm"), ("max_abs_scaler", preprocessing.MaxAbsScaler(), "culmen_length_mm"), ("min_max_scaler", preprocessing.MinMaxScaler(), "culmen_length_mm"), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + "culmen_length_mm", + ), ("standard_scaler", preprocessing.StandardScaler(), "flipper_length_mm"), ("max_abs_scaler", preprocessing.MaxAbsScaler(), "flipper_length_mm"), ("min_max_scaler", preprocessing.MinMaxScaler(), "flipper_length_mm"), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + "flipper_length_mm", + ), ] assert transformers == expected @@ -791,6 +811,32 @@ def test_pipeline_min_max_scaler_to_gbq(penguins_df_default_index, dataset_id): assert pl_loaded._estimator.fit_intercept is False +def test_pipeline_k_bins_discretizer_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ("transform", preprocessing.KBinsDiscretizer(strategy="uniform")), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "culmen_length_mm", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_k_bins_discretizer", replace=True + ) + assert isinstance(pl_loaded._transform, preprocessing.KBinsDiscretizer) + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False + + def test_pipeline_one_hot_encoder_to_gbq(penguins_df_default_index, dataset_id): pl = pipeline.Pipeline( [ diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py index fc8f3251bda..45548acca32 100644 --- a/tests/system/small/ml/test_preprocessing.py +++ b/tests/system/small/ml/test_preprocessing.py @@ -121,7 +121,7 @@ def test_standard_scaler_series_normalizes(penguins_df_default_index, new_pengui def test_max_abs_scaler_normalizes(penguins_df_default_index, new_penguins_df): - # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.MaxAbsScaler, when BQML's change is in prod. scaler = bigframes.ml.preprocessing.MaxAbsScaler() scaler.fit( penguins_df_default_index[ @@ -211,7 +211,7 @@ def test_max_abs_scaler_series_normalizes(penguins_df_default_index, new_penguin pd.testing.assert_frame_equal(result, expected, rtol=1e-3) -def test_min_max_scaler_normalizeds_fit_transform(new_penguins_df): +def test_min_max_scaler_normalized_fit_transform(new_penguins_df): scaler = bigframes.ml.preprocessing.MinMaxScaler() result = scaler.fit_transform( new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] @@ -265,7 +265,7 @@ def test_min_max_scaler_series_normalizes(penguins_df_default_index, new_penguin def test_min_max_scaler_normalizes(penguins_df_default_index, new_penguins_df): - # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.StandardScaler, when BQML's change is in prod. + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.MinMaxScaler, when BQML's change is in prod. scaler = bigframes.ml.preprocessing.MinMaxScaler() scaler.fit( penguins_df_default_index[ @@ -304,6 +304,131 @@ def test_min_max_scaler_normalizes(penguins_df_default_index, new_penguins_df): pd.testing.assert_frame_equal(result, expected, rtol=1e-3) +def test_k_bins_discretizer_normalized_fit_transform_default_params(new_penguins_df): + discretizer = bigframes.ml.preprocessing.KBinsDiscretizer(strategy="uniform") + result = discretizer.fit_transform( + new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] + ).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "kbinsdiscretizer_culmen_depth_mm": ["bin_5", "bin_2", "bin_4"], + "kbinsdiscretizer_culmen_length_mm": ["bin_5", "bin_3", "bin_2"], + "kbinsdiscretizer_flipper_length_mm": ["bin_5", "bin_2", "bin_4"], + }, + dtype="string[pyarrow]", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_k_bins_discretizer_series_normalizes( + penguins_df_default_index, new_penguins_df +): + discretizer = bigframes.ml.preprocessing.KBinsDiscretizer(strategy="uniform") + discretizer.fit(penguins_df_default_index["culmen_length_mm"]) + + result = discretizer.transform( + penguins_df_default_index["culmen_length_mm"] + ).to_pandas() + result = discretizer.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "kbinsdiscretizer_culmen_length_mm": ["bin_3", "bin_3", "bin_3"], + }, + dtype="string[pyarrow]", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_k_bins_discretizer_normalizes(penguins_df_default_index, new_penguins_df): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.KBinsDiscretizer, when BQML's change is in prod. + discretizer = bigframes.ml.preprocessing.KBinsDiscretizer(strategy="uniform") + discretizer.fit( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + result = discretizer.transform( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ).to_pandas() + + result = discretizer.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "kbinsdiscretizer_culmen_depth_mm": ["bin_5", "bin_4", "bin_4"], + "kbinsdiscretizer_culmen_length_mm": ["bin_3", "bin_3", "bin_3"], + "kbinsdiscretizer_flipper_length_mm": ["bin_4", "bin_2", "bin_3"], + }, + dtype="string[pyarrow]", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + +def test_k_bins_discretizer_normalizes_different_params( + penguins_df_default_index, new_penguins_df +): + # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.KBinsDiscretizer, when BQML's change is in prod. + discretizer = bigframes.ml.preprocessing.KBinsDiscretizer( + n_bins=6, strategy="uniform" + ) + discretizer.fit( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + result = discretizer.transform( + penguins_df_default_index[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ).to_pandas() + + result = discretizer.transform(new_penguins_df).to_pandas() + + # TODO: bug? feature columns seem to be in nondeterministic random order + # workaround: sort columns by name. Can't repro it in pantheon, so could + # be a bigframes issue... + result = result.reindex(sorted(result.columns), axis=1) + + expected = pd.DataFrame( + { + "kbinsdiscretizer_culmen_depth_mm": ["bin_6", "bin_4", "bin_5"], + "kbinsdiscretizer_culmen_length_mm": ["bin_3", "bin_3", "bin_3"], + "kbinsdiscretizer_flipper_length_mm": ["bin_4", "bin_2", "bin_3"], + }, + dtype="string[pyarrow]", + index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=1e-3) + + def test_one_hot_encoder_default_params(new_penguins_df): encoder = bigframes.ml.preprocessing.OneHotEncoder() encoder.fit(new_penguins_df[["species", "sex"]]) diff --git a/tests/unit/ml/test_compose.py b/tests/unit/ml/test_compose.py index 8c8fbd6ab54..60dcc75b63c 100644 --- a/tests/unit/ml/test_compose.py +++ b/tests/unit/ml/test_compose.py @@ -23,6 +23,7 @@ def test_columntransformer_init_expectedtransforms(): standard_scaler_transformer = preprocessing.StandardScaler() max_abs_scaler_transformer = preprocessing.MaxAbsScaler() min_max_scaler_transformer = preprocessing.MinMaxScaler() + k_bins_discretizer_transformer = preprocessing.KBinsDiscretizer(strategy="uniform") label_transformer = preprocessing.LabelEncoder() column_transformer = compose.ColumnTransformer( [ @@ -42,6 +43,11 @@ def test_columntransformer_init_expectedtransforms(): min_max_scaler_transformer, ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + k_bins_discretizer_transformer, + ["culmen_length_mm", "flipper_length_mm"], + ), ("label", label_transformer, "species"), ] ) @@ -54,6 +60,8 @@ def test_columntransformer_init_expectedtransforms(): ("max_abs_scale", max_abs_scaler_transformer, "flipper_length_mm"), ("min_max_scale", min_max_scaler_transformer, "culmen_length_mm"), ("min_max_scale", min_max_scaler_transformer, "flipper_length_mm"), + ("k_bins_discretizer", k_bins_discretizer_transformer, "culmen_length_mm"), + ("k_bins_discretizer", k_bins_discretizer_transformer, "flipper_length_mm"), ("label", label_transformer, "species"), ] @@ -81,6 +89,11 @@ def test_columntransformer_repr(): preprocessing.MinMaxScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) @@ -92,6 +105,9 @@ def test_columntransformer_repr(): ('max_abs_scale', MaxAbsScaler(), ['culmen_length_mm', 'flipper_length_mm']), ('min_max_scale', MinMaxScaler(), + ['culmen_length_mm', 'flipper_length_mm']), + ('k_bins_discretizer', + KBinsDiscretizer(strategy='uniform'), ['culmen_length_mm', 'flipper_length_mm'])])""" ) @@ -119,6 +135,11 @@ def test_columntransformer_repr_matches_sklearn(): preprocessing.MinMaxScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + preprocessing.KBinsDiscretizer(strategy="uniform"), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) sk_column_transformer = sklearn_compose.ColumnTransformer( @@ -143,6 +164,11 @@ def test_columntransformer_repr_matches_sklearn(): sklearn_preprocessing.MinMaxScaler(), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "k_bins_discretizer", + sklearn_preprocessing.KBinsDiscretizer(strategy="uniform"), + ["culmen_length_mm", "flipper_length_mm"], + ), ] ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index a3338e762d4..34a02edd424 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -95,6 +95,13 @@ def test_min_max_scaler_produces_correct_sql( assert sql == "ML.MIN_MAX_SCALER(col_a) OVER() AS scaled_col_a" +def test_k_bins_discretizer_produces_correct_sql( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_bucketize("col_a", [1, 2, 3, 4], "scaled_col_a") + assert sql == "ML.BUCKETIZE(col_a, [1, 2, 3, 4], FALSE) AS scaled_col_a" + + def test_one_hot_encoder_produces_correct_sql( base_sql_generator: ml_sql.BaseSqlGenerator, ): diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_discretization.py b/third_party/bigframes_vendored/sklearn/preprocessing/_discretization.py new file mode 100644 index 00000000000..0236558dd4b --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_discretization.py @@ -0,0 +1,47 @@ +# Author: Henry Lin +# Tom Dupré la Tour + +# License: BSD + +from bigframes import constants +from third_party.bigframes_vendored.sklearn.base import BaseEstimator, TransformerMixin + + +class KBinsDiscretizer(TransformerMixin, BaseEstimator): + """ + Bin continuous data into intervals. + + Args: + n_bins (int, default 5): + The number of bins to produce. Raises ValueError if ``n_bins < 2``. + strategy ({'uniform', 'quantile'}, default='quantile'): + Strategy used to define the widths of the bins. 'uniform': All bins + in each feature have identical widths. 'quantile': All bins in each + feature have the same number of points. Only `uniform` is supported now. + """ + + def fit(self, X, y=None): + """Fit the estimator. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Dataframe or Series with training data. + + y (default None): + Ignored. + + Returns: + KBinsDiscretizer: Fitted scaler. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transform(self, X): + """Discretize the data. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. + + Returns: + bigframes.dataframe.DataFrame: Transformed result.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From bd692d847567d9c31cfc7d652ed8f7d5e7aeaa64 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 4 Oct 2023 08:48:03 -0700 Subject: [PATCH 58/80] chore(main): release 0.6.0 (#72) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 19 +++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e4b2bff3c71..880f7916250 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,25 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.6.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.5.0...v0.6.0) (2023-10-04) + + +### Features + +* Add df.unstack ([#63](https://github.com/googleapis/python-bigquery-dataframes/issues/63)) ([4a84714](https://github.com/googleapis/python-bigquery-dataframes/commit/4a84714e2fb07f70c70c79f8b8da9fcb41096e33)) +* Add idxmin, idxmax to series, dataframe ([#74](https://github.com/googleapis/python-bigquery-dataframes/issues/74)) ([781307e](https://github.com/googleapis/python-bigquery-dataframes/commit/781307ec22d31a7657f8ee5c6eedc0e419450ccd)) +* Add ml.preprocessing.KBinsDiscretizer ([#81](https://github.com/googleapis/python-bigquery-dataframes/issues/81)) ([24c6256](https://github.com/googleapis/python-bigquery-dataframes/commit/24c625638984f6a84191c7a4c8ac9fb6c3cf1dca)) +* Add multi-column dataframe merge ([#73](https://github.com/googleapis/python-bigquery-dataframes/issues/73)) ([c9fa85c](https://github.com/googleapis/python-bigquery-dataframes/commit/c9fa85cc338be5e9a8dde59b255690aedbbc1127)) +* Add update and align methods to dataframe ([#57](https://github.com/googleapis/python-bigquery-dataframes/issues/57)) ([bf050cf](https://github.com/googleapis/python-bigquery-dataframes/commit/bf050cf475ad8a9e3e0ca3f896ddaf96dbe13ae3)) +* Support STRUCT data type with `Series.struct.field` to extract child fields ([#71](https://github.com/googleapis/python-bigquery-dataframes/issues/71)) ([17afac9](https://github.com/googleapis/python-bigquery-dataframes/commit/17afac9ff70a2b93ed70dc7bcce7beb9a53c2ece)) + + +### Bug Fixes + +* Avoid `403 response too large to return` error with `read_gbq` and large query results ([#77](https://github.com/googleapis/python-bigquery-dataframes/issues/77)) ([8f3b5b2](https://github.com/googleapis/python-bigquery-dataframes/commit/8f3b5b240f0f28fef92465abc53504e875d7335a)) +* Change return type of `Series.loc[scalar]` ([#40](https://github.com/googleapis/python-bigquery-dataframes/issues/40)) ([fff3d45](https://github.com/googleapis/python-bigquery-dataframes/commit/fff3d45f03ffbc7bb23143a1572e3dd157463ca9)) +* Fix df/series.iloc by list with multiindex ([#79](https://github.com/googleapis/python-bigquery-dataframes/issues/79)) ([971d091](https://github.com/googleapis/python-bigquery-dataframes/commit/971d091cac9ad662145a3d43d8f9a785eb0ccc23)) + ## [0.5.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.4.0...v0.5.0) (2023-09-28) diff --git a/bigframes/version.py b/bigframes/version.py index ad3c3082c52..238b64473a0 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.5.0" +__version__ = "0.6.0" From c0efec8956198247b27904345a795f09c80d3502 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 4 Oct 2023 10:14:08 -0700 Subject: [PATCH 59/80] feat: add aliases for several series properties (#80) * feat: add aliases for several series properties --------- Co-authored-by: Tim Swast --- bigframes/series.py | 15 +++++++++++++++ .../bigframes_vendored/pandas/core/series.py | 14 ++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/bigframes/series.py b/bigframes/series.py index 8815a6abdee..2a0b1ff1fa4 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -99,6 +99,10 @@ def shape(self) -> typing.Tuple[int]: def size(self) -> int: return self.shape[0] + @property + def ndim(self) -> int: + return 1 + @property def empty(self) -> bool: return self.shape[0] == 0 @@ -123,6 +127,13 @@ def query_job(self) -> Optional[bigquery.QueryJob]: def struct(self) -> structs.StructAccessor: return structs.StructAccessor(self._block) + @property + def T(self) -> Series: + return self.transpose() + + def transpose(self) -> Series: + return self + def _set_internal_query_job(self, query_job: bigquery.QueryJob): self._query_job = query_job @@ -362,6 +373,8 @@ def ffill(self, *, limit: typing.Optional[int] = None) -> Series: window = bigframes.core.WindowSpec(preceding=limit, following=0) return self._apply_window_op(agg_ops.LastNonNullOp(), window) + pad = ffill + def bfill(self, *, limit: typing.Optional[int] = None) -> Series: window = bigframes.core.WindowSpec(preceding=0, following=limit) return self._apply_window_op(agg_ops.FirstNonNullOp(), window) @@ -743,6 +756,8 @@ def agg(self, func: str | typing.Sequence[str]) -> scalars.Scalar | Series: agg_ops.lookup_agg_func(typing.cast(str, func)) ) + aggregate = agg + def skew(self): count = self.count() if count < 3: diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index a41a3454ca4..970ab1d8b43 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -64,6 +64,20 @@ def name(self) -> Hashable: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + @property + def T(self) -> Series: + """Return the transpose, which is by definition self.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transpose(self) -> Series: + """ + Return the transpose, which is by definition self. + + Returns: + Series + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def reset_index( self, *, From 740c45176f79d4d2f7f28cb5f6c9eeb1327c8397 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Wed, 4 Oct 2023 16:17:20 -0700 Subject: [PATCH 60/80] feat: Allow df.drop to take an index object (#68) * feat: Allow passing index objects to df.drop * remove notebook files * add first implementation for df.drop(index) * use index_columns property * don't use _expr.keys() * fix order bug and add test * fix index names * support multiindex * remove accidentally added files * add type hint * remove debug print statements * fix: fix df/series.iloc by list with multiindex * complete merge * remove unneeded isinstance * refactor _drop_by_index --------- Co-authored-by: Henry J Solberg Co-authored-by: Tim Swast --- bigframes/dataframe.py | 27 +++++++++++++- tests/system/small/test_dataframe.py | 55 ++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index eea8beb130a..8874a4edb89 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -937,7 +937,7 @@ def drop( columns = labels block = self._block - if index: + if index is not None: level_id = self._resolve_levels(level or 0)[0] if utils.is_list_like(index): @@ -947,6 +947,8 @@ def drop( block, condition_id = block.apply_unary_op( inverse_condition_id, ops.invert_op ) + elif isinstance(index, indexes.Index): + return self._drop_by_index(index) else: block, condition_id = block.apply_unary_op( level_id, ops.partial_right(ops.ne_op, index) @@ -956,10 +958,31 @@ def drop( ) if columns: block = block.drop_columns(self._sql_names(columns)) - if not index and not columns: + if index is None and not columns: raise ValueError("Must specify 'labels' or 'index'/'columns") return DataFrame(block) + def _drop_by_index(self, index: indexes.Index) -> DataFrame: + block = index._data._get_block() + block, ordering_col = block.promote_offsets() + joined_index, (get_column_left, get_column_right) = self._block.index.join( + block.index + ) + + new_ordering_col = get_column_right(ordering_col) + drop_block = joined_index._block + drop_block, drop_col = drop_block.apply_unary_op( + new_ordering_col, + ops.isnull_op, + ) + + drop_block = drop_block.filter(drop_col) + original_columns = [ + get_column_left(column) for column in self._block.value_columns + ] + drop_block = drop_block.select_columns(original_columns) + return DataFrame(drop_block) + def droplevel(self, level: LevelsType, axis: int | str = 0): axis_n = utils.get_axis_number(axis) if axis_n == 0: diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index b8616a54d65..19ea9b8ae55 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -258,6 +258,61 @@ def test_drop_index(scalars_dfs): pd.testing.assert_frame_equal(pd_result, bf_result) +def test_drop_pandas_index(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + drop_index = scalars_pandas_df.iloc[[4, 1, 2]].index + + pd_result = scalars_pandas_df.drop(index=drop_index) + bf_result = scalars_df.drop(index=drop_index).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + +def test_drop_bigframes_index(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + drop_index = scalars_df.loc[[4, 1, 2]].index + drop_pandas_index = scalars_pandas_df.loc[[4, 1, 2]].index + + pd_result = scalars_pandas_df.drop(index=drop_pandas_index) + bf_result = scalars_df.drop(index=drop_index).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + +def test_drop_bigframes_index_with_na(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + scalars_df = scalars_df.copy() + scalars_pandas_df = scalars_pandas_df.copy() + scalars_df = scalars_df.set_index("bytes_col") + scalars_pandas_df = scalars_pandas_df.set_index("bytes_col") + drop_index = scalars_df.iloc[[3, 5]].index + drop_pandas_index = scalars_pandas_df.iloc[[3, 5]].index + + pd_result = scalars_pandas_df.drop(index=drop_pandas_index) # drop_pandas_index) + bf_result = scalars_df.drop(index=drop_index).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + +def test_drop_bigframes_multiindex(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + scalars_df = scalars_df.copy() + scalars_pandas_df = scalars_pandas_df.copy() + sub_df = scalars_df.iloc[[4, 1, 2]] + sub_pandas_df = scalars_pandas_df.iloc[[4, 1, 2]] + sub_df = sub_df.set_index(["bytes_col", "numeric_col"]) + sub_pandas_df = sub_pandas_df.set_index(["bytes_col", "numeric_col"]) + drop_index = sub_df.index + drop_pandas_index = sub_pandas_df.index + + scalars_df = scalars_df.set_index(["bytes_col", "numeric_col"]) + scalars_pandas_df = scalars_pandas_df.set_index(["bytes_col", "numeric_col"]) + bf_result = scalars_df.drop(index=drop_index).to_pandas() + pd_result = scalars_pandas_df.drop(index=drop_pandas_index) + + pd.testing.assert_frame_equal(pd_result, bf_result) + + def test_drop_labels_axis_0(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs From ff3bb89cca2e29625166b4646a34ce6b889d58f9 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 5 Oct 2023 13:19:39 -0700 Subject: [PATCH 61/80] refactor: simplify ArrayValue public interface (#82) * refactor: simplify ArrayValue public interface --------- Co-authored-by: Tim Swast --- bigframes/core/__init__.py | 245 ++++++++++---------------- bigframes/core/blocks.py | 116 ++++++++---- bigframes/core/groupby/__init__.py | 4 - bigframes/core/indexes/index.py | 4 +- bigframes/core/joins/row_identity.py | 16 +- bigframes/core/joins/single_column.py | 22 +-- bigframes/dataframe.py | 6 +- bigframes/operations/base.py | 6 - bigframes/series.py | 6 +- tests/unit/test_core.py | 42 +---- 10 files changed, 212 insertions(+), 255 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 8008c1189a1..ccfd682215c 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -16,8 +16,9 @@ from dataclasses import dataclass import functools import math +import textwrap import typing -from typing import Collection, Dict, Iterable, Literal, Optional, Sequence, Tuple +from typing import Collection, Iterable, Literal, Optional, Sequence, Tuple from google.cloud import bigquery import ibis @@ -201,31 +202,27 @@ def mem_expr_from_pandas( hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), ) - @property - def table(self) -> ibis_types.Table: - return self._table - - @property - def reduced_predicate(self) -> typing.Optional[ibis_types.BooleanValue]: - """Returns the frame's predicates as an equivalent boolean value, useful where a single predicate value is preferred.""" - return ( - _reduce_predicate_list(self._predicates).name(PREDICATE_COLUMN) - if self._predicates - else None - ) - @property def columns(self) -> typing.Tuple[ibis_types.Value, ...]: return self._columns @property - def column_names(self) -> Dict[str, ibis_types.Value]: - return self._column_names + def column_ids(self) -> typing.Sequence[str]: + return tuple(self._column_names.keys()) @property def hidden_ordering_columns(self) -> typing.Tuple[ibis_types.Value, ...]: return self._hidden_ordering_columns + @property + def _reduced_predicate(self) -> typing.Optional[ibis_types.BooleanValue]: + """Returns the frame's predicates as an equivalent boolean value, useful where a single predicate value is preferred.""" + return ( + _reduce_predicate_list(self._predicates).name(PREDICATE_COLUMN) + if self._predicates + else None + ) + @property def _ibis_order(self) -> Sequence[ibis_types.Value]: """Returns a sequence of ibis values which can be directly used to order a table expression. Has direction modifiers applied.""" @@ -265,24 +262,22 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue: def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: ibis_type = typing.cast( - bigframes.dtypes.IbisDtype, self.get_any_column(key).type() + bigframes.dtypes.IbisDtype, self._get_any_column(key).type() ) return typing.cast( bigframes.dtypes.Dtype, bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), ) - def get_column(self, key: str) -> ibis_types.Value: + def _get_ibis_column(self, key: str) -> ibis_types.Value: """Gets the Ibis expression for a given column.""" - if key not in self._column_names.keys(): + if key not in self.column_ids: raise ValueError( - "Column name {} not in set of values: {}".format( - key, self._column_names.keys() - ) + "Column name {} not in set of values: {}".format(key, self.column_ids) ) return typing.cast(ibis_types.Value, self._column_names[key]) - def get_any_column(self, key: str) -> ibis_types.Value: + def _get_any_column(self, key: str) -> ibis_types.Value: """Gets the Ibis expression for a given column. Will also get hidden columns.""" all_columns = {**self._column_names, **self._hidden_ordering_column_names} if key not in all_columns.keys(): @@ -303,26 +298,11 @@ def _get_hidden_ordering_column(self, key: str) -> ibis_types.Column: ) return typing.cast(ibis_types.Column, self._hidden_ordering_column_names[key]) - def apply_limit(self, max_results: int) -> ArrayValue: - table = self._to_ibis_expr( - ordering_mode="order_by", - expose_hidden_cols=True, - ).limit(max_results) - columns = [table[column_name] for column_name in self._column_names] - hidden_ordering_columns = [ - table[column_name] for column_name in self._hidden_ordering_column_names - ] - return ArrayValue( - self._session, - table, - columns=columns, - hidden_ordering_columns=hidden_ordering_columns, - ordering=self._ordering, - ) - def filter(self, predicate_id: str, keep_null: bool = False) -> ArrayValue: """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" - condition = typing.cast(ibis_types.BooleanValue, self.get_column(predicate_id)) + condition = typing.cast( + ibis_types.BooleanValue, self._get_ibis_column(predicate_id) + ) if keep_null: condition = typing.cast( ibis_types.BooleanValue, @@ -358,7 +338,7 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue: The row numbers of result is non-deterministic, avoid to use. """ table = self._to_ibis_expr( - ordering_mode="order_by", expose_hidden_cols=True, fraction=fraction + "unordered", expose_hidden_cols=True, fraction=fraction ) columns = [table[column_name] for column_name in self._column_names] hidden_ordering_columns = [ @@ -373,7 +353,7 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue: ) @property - def offsets(self): + def _offsets(self) -> ibis_types.IntegerColumn: if not self._ordering.is_sequential: raise ValueError( "Expression does not have offsets. Generate them first using project_offsets." @@ -382,9 +362,10 @@ def offsets(self): raise ValueError( "Ordering is invalid. Marked as sequential but no total order columns." ) - return self.get_any_column(self._ordering.total_order_col.column_id) + column = self._get_any_column(self._ordering.total_order_col.column_id) + return typing.cast(ibis_types.IntegerColumn, column) - def project_offsets(self) -> ArrayValue: + def _project_offsets(self) -> ArrayValue: """Create a new expression that contains offsets. Should only be executed when offsets are needed for an operations. Has no effect on expression semantics.""" if self._ordering.is_sequential: return self @@ -414,7 +395,7 @@ def _hide_column(self, column_id) -> ArrayValue: new_name = bigframes.core.guid.generate_guid(prefix="bigframes_hidden_") expr_builder.hidden_ordering_columns = [ *self._hidden_ordering_columns, - self.get_column(column_id).name(new_name), + self._get_ibis_column(column_id).name(new_name), ] expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) return expr_builder.build() @@ -427,26 +408,28 @@ def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: ordering = self._ordering if (not ordering.is_sequential) or (not ordering.total_order_col): - return self.project_offsets().promote_offsets() + return self._project_offsets().promote_offsets() col_id = bigframes.core.guid.generate_guid() expr_builder = self.builder() expr_builder.columns = [ - self.get_any_column(ordering.total_order_col.column_id).name(col_id), + self._get_any_column(ordering.total_order_col.column_id).name(col_id), *self.columns, ] return expr_builder.build(), col_id def select_columns(self, column_ids: typing.Sequence[str]): - return self.projection([self.get_column(col_id) for col_id in column_ids]) + return self._projection( + [self._get_ibis_column(col_id) for col_id in column_ids] + ) - def projection(self, columns: Iterable[ibis_types.Value]) -> ArrayValue: + def _projection(self, columns: Iterable[ibis_types.Value]) -> ArrayValue: """Creates a new expression based on this expression with new columns.""" # TODO(swast): We might want to do validation here that columns derive # from the same table expression instead of (in addition to?) at # construction time. expr = self - for ordering_column in set(self.column_names.keys()).intersection( + for ordering_column in set(self.column_ids).intersection( [col_ref.column_id for col_ref in self._ordering.ordering_value_columns] ): # Need to hide ordering columns that are being dropped. Alternatively, could project offsets @@ -459,7 +442,7 @@ def projection(self, columns: Iterable[ibis_types.Value]) -> ArrayValue: def shape(self) -> typing.Tuple[int, int]: """Returns dimensions as (length, width) tuple.""" width = len(self.columns) - count_expr = self._to_ibis_expr(ordering_mode="unordered").count() + count_expr = self._to_ibis_expr("unordered").count() sql = self._session.ibis_client.compile(count_expr) # Support in-memory engines for hermetic unit tests. @@ -527,7 +510,7 @@ def project_unary_op( self, column_name: str, op: ops.UnaryOp, output_name=None ) -> ArrayValue: """Creates a new expression based on this expression with unary operation applied to one column.""" - value = op._as_ibis(self.get_column(column_name)).name( + value = op._as_ibis(self._get_ibis_column(column_name)).name( output_name or column_name ) return self._set_or_replace_by_id(output_name or column_name, value) @@ -541,7 +524,8 @@ def project_binary_op( ) -> ArrayValue: """Creates a new expression based on this expression with binary operation applied to two columns.""" value = op( - self.get_column(left_column_id), self.get_column(right_column_id) + self._get_ibis_column(left_column_id), + self._get_ibis_column(right_column_id), ).name(output_column_id) return self._set_or_replace_by_id(output_column_id, value) @@ -555,9 +539,9 @@ def project_ternary_op( ) -> ArrayValue: """Creates a new expression based on this expression with ternary operation applied to three columns.""" value = op( - self.get_column(col_id_1), - self.get_column(col_id_2), - self.get_column(col_id_3), + self._get_ibis_column(col_id_1), + self._get_ibis_column(col_id_2), + self._get_ibis_column(col_id_3), ).name(output_column_id) return self._set_or_replace_by_id(output_column_id, value) @@ -574,7 +558,7 @@ def aggregate( by_column_id: column id of the aggregation key, this is preserved through the transform dropna: whether null keys should be dropped """ - table = self._to_ibis_expr(ordering_mode="unordered") + table = self._to_ibis_expr("unordered") stats = { col_out: agg_op._as_ibis(table[col_in]) for col_in, agg_op, col_out in aggregations @@ -594,10 +578,10 @@ def aggregate( if dropna: for column_id in by_column_ids: expr = expr._filter( - ops.notnull_op._as_ibis(expr.get_column(column_id)) + ops.notnull_op._as_ibis(expr._get_ibis_column(column_id)) ) # Can maybe remove this as Ordering id is redundant as by_column is unique after aggregation - return expr.project_offsets() + return expr._project_offsets() else: aggregates = {**stats, ORDER_ID_COLUMN: ibis_types.literal(0)} result = table.aggregate(**aggregates) @@ -624,7 +608,7 @@ def corr_aggregate( Arguments: corr_aggregations: left_column_id, right_column_id, output_column_id tuples """ - table = self._to_ibis_expr(ordering_mode="unordered") + table = self._to_ibis_expr("unordered") stats = { col_out: table[col_left].corr(table[col_right], how="pop") for col_left, col_right, col_out in corr_aggregations @@ -664,7 +648,7 @@ def project_window_op( never_skip_nulls: will disable null skipping for operators that would otherwise do so skip_reproject_unsafe: skips the reprojection step, can be used when performing many non-dependent window operations, user responsible for not nesting window expressions, or using outputs as join, filter or aggregation keys before a reprojection """ - column = typing.cast(ibis_types.Column, self.get_column(column_name)) + column = typing.cast(ibis_types.Column, self._get_ibis_column(column_name)) window = self._ibis_window_from_spec(window_spec, allow_ties=op.handles_ties) window_op = op._as_ibis(column, window) @@ -700,26 +684,34 @@ def project_window_op( def to_sql( self, - ordering_mode: Literal[ - "order_by", "string_encoded", "offset_col", "unordered" - ] = "order_by", - order_col_name: Optional[str] = ORDER_ID_COLUMN, + offset_column: typing.Optional[str] = None, col_id_overrides: typing.Mapping[str, str] = {}, + sorted: bool = False, ) -> str: + offsets_id = offset_column or ORDER_ID_COLUMN + sql = self._session.ibis_client.compile( self._to_ibis_expr( - ordering_mode=ordering_mode, - order_col_name=order_col_name, + ordering_mode="offset_col" + if (offset_column or sorted) + else "unordered", + order_col_name=offsets_id, col_id_overrides=col_id_overrides, ) ) + if sorted: + sql = textwrap.dedent( + f""" + SELECT * EXCEPT (`{offsets_id}`) + FROM ({sql}) + ORDER BY `{offsets_id}` + """ + ) return typing.cast(str, sql) def _to_ibis_expr( self, - ordering_mode: Literal[ - "order_by", "string_encoded", "offset_col", "unordered" - ] = "order_by", + ordering_mode: Literal["string_encoded", "offset_col", "unordered"], order_col_name: Optional[str] = ORDER_ID_COLUMN, expose_hidden_cols: bool = False, fraction: Optional[float] = None, @@ -731,8 +723,6 @@ def _to_ibis_expr( ArrayValue objects are sorted, so the following options are available to reflect this in the ibis expression. - * "order_by" (Default): The output table will not have an ordering - column, however there will be an order_by clause applied to the ouput. * "offset_col": Zero-based offsets are generated as a column, this will not sort the rows however. * "string_encoded": An ordered string column is provided in output table. @@ -760,7 +750,6 @@ def _to_ibis_expr( An ibis expression representing the data help by the ArrayValue object. """ assert ordering_mode in ( - "order_by", "string_encoded", "offset_col", "unordered", @@ -775,18 +764,16 @@ def _to_ibis_expr( str ] = [] # Ordering/Filtering columns that will be dropped at end - if self.reduced_predicate is not None: - columns.append(self.reduced_predicate) + if self._reduced_predicate is not None: + columns.append(self._reduced_predicate) # Usually drop predicate as it is will be all TRUE after filtering if not expose_hidden_cols: - columns_to_drop.append(self.reduced_predicate.get_name()) + columns_to_drop.append(self._reduced_predicate.get_name()) order_columns = self._create_order_columns( ordering_mode, order_col_name, expose_hidden_cols ) columns.extend(order_columns) - if (ordering_mode == "order_by") and not expose_hidden_cols: - columns_to_drop.extend(col.get_name() for col in order_columns) # Special case for empty tables, since we can't create an empty # projection. @@ -799,15 +786,8 @@ def _to_ibis_expr( bigframes.dtypes.ibis_value_to_canonical_type(column) for column in columns ) base_table = table - if self.reduced_predicate is not None: + if self._reduced_predicate is not None: table = table.filter(base_table[PREDICATE_COLUMN]) - if ordering_mode == "order_by": - table = table.order_by( - _convert_ordering_to_table_values( - {col: base_table[col] for col in table.columns}, - self._ordering.all_ordering_columns, - ) # type: ignore - ) table = table.drop(*columns_to_drop) if col_id_overrides: table = table.relabel(col_id_overrides) @@ -826,24 +806,24 @@ def _create_order_columns( return (self._create_offset_column().name(order_col_name),) elif ordering_mode == "string_encoded": return (self._create_string_ordering_column().name(order_col_name),) - elif ordering_mode == "order_by" or expose_hidden_cols: + elif expose_hidden_cols: return self.hidden_ordering_columns return () def _create_offset_column(self) -> ibis_types.IntegerColumn: if self._ordering.total_order_col and self._ordering.is_sequential: - offsets = self.get_any_column(self._ordering.total_order_col.column_id) + offsets = self._get_any_column(self._ordering.total_order_col.column_id) return typing.cast(ibis_types.IntegerColumn, offsets) else: window = ibis.window(order_by=self._ibis_order) if self._predicates: - window = window.group_by(self.reduced_predicate) + window = window.group_by(self._reduced_predicate) offsets = ibis.row_number().over(window) return typing.cast(ibis_types.IntegerColumn, offsets) def _create_string_ordering_column(self) -> ibis_types.StringColumn: if self._ordering.total_order_col and self._ordering.is_string_encoded: - string_order_ids = self.get_any_column( + string_order_ids = self._get_any_column( self._ordering.total_order_col.column_id ) return typing.cast(ibis_types.StringColumn, string_order_ids) @@ -852,7 +832,7 @@ def _create_string_ordering_column(self) -> ibis_types.StringColumn: and self._ordering.integer_encoding.is_encoded ): # Special case: non-negative integer ordering id can be converted directly to string without regenerating row numbers - int_values = self.get_any_column(self._ordering.total_order_col.column_id) + int_values = self._get_any_column(self._ordering.total_order_col.column_id) return encode_order_string( typing.cast(ibis_types.IntegerColumn, int_values), ) @@ -860,7 +840,7 @@ def _create_string_ordering_column(self) -> ibis_types.StringColumn: # Have to build string from scratch window = ibis.window(order_by=self._ibis_order) if self._predicates: - window = window.group_by(self.reduced_predicate) + window = window.group_by(self._reduced_predicate) row_nums = typing.cast( ibis_types.IntegerColumn, ibis.row_number().over(window) ) @@ -870,7 +850,8 @@ def start_query( self, job_config: Optional[bigquery.job.QueryJobConfig] = None, max_results: Optional[int] = None, - expose_extra_columns: bool = False, + *, + sorted: bool = True, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """Execute a query and return metadata about the results.""" # TODO(swast): Cache the job ID so we can look it up again if they ask @@ -883,8 +864,7 @@ def start_query( # a LocalSession for unit testing. # TODO(swast): Add a timeout here? If the query is taking a long time, # maybe we just print the job metadata that we have so far? - table = self._to_ibis_expr(expose_hidden_cols=expose_extra_columns) - sql = self._session.ibis_client.compile(table) # type:ignore + sql = self.to_sql(sorted=True) # type:ignore return self._session._start_query( sql=sql, job_config=job_config, @@ -903,7 +883,7 @@ def _reproject_to_table(self) -> ArrayValue: recursively in projections. """ table = self._to_ibis_expr( - ordering_mode="unordered", + "unordered", expose_hidden_cols=True, ) columns = [table[column_name] for column_name in self._column_names] @@ -926,14 +906,16 @@ def _reproject_to_table(self) -> ArrayValue: def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = False): group_by: typing.List[ibis_types.Value] = ( [ - typing.cast(ibis_types.Column, _as_identity(self.get_column(column))) + typing.cast( + ibis_types.Column, _as_identity(self._get_ibis_column(column)) + ) for column in window_spec.grouping_keys ] if window_spec.grouping_keys else [] ) - if self.reduced_predicate is not None: - group_by.append(self.reduced_predicate) + if self._reduced_predicate is not None: + group_by.append(self._reduced_predicate) if window_spec.ordering: order_by = _convert_ordering_to_table_values( {**self._column_names, **self._hidden_ordering_column_names}, @@ -984,7 +966,7 @@ def unpivot( """ if how not in ("left", "right"): raise ValueError("'how' must be 'left' or 'right'") - table = self._to_ibis_expr(ordering_mode="unordered", expose_hidden_cols=True) + table = self._to_ibis_expr("unordered", expose_hidden_cols=True) row_n = len(row_labels) hidden_col_ids = self._hidden_ordering_column_names.keys() if not all( @@ -1107,7 +1089,9 @@ def unpivot( ) def assign(self, source_id: str, destination_id: str) -> ArrayValue: - return self._set_or_replace_by_id(destination_id, self.get_column(source_id)) + return self._set_or_replace_by_id( + destination_id, self._get_ibis_column(source_id) + ) def assign_constant( self, @@ -1134,74 +1118,25 @@ def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value) -> ArrayVa return self._hide_column(id)._set_or_replace_by_id(id, new_value) builder = self.builder() - if id in self.column_names: + if id in self.column_ids: builder.columns = [ val if (col_id != id) else new_value.name(id) - for col_id, val in self.column_names.items() + for col_id, val in zip(self.column_ids, self._columns) ] else: builder.columns = [*self.columns, new_value.name(id)] return builder.build() - def slice( - self, - start: typing.Optional[int] = None, - stop: typing.Optional[int] = None, - step: typing.Optional[int] = None, - ) -> ArrayValue: - if step == 0: - raise ValueError("slice step cannot be zero") - - if not step: - step = 1 - - expr_with_offsets = self.project_offsets() - - # start with True and reduce with start, stop, and step conditions - cond_list = [expr_with_offsets.offsets == expr_with_offsets.offsets] - - last_offset = expr_with_offsets.offsets.max() - - # Convert negative indexes to positive indexes - if start and start < 0: - start = last_offset + start + 1 - if stop and stop < 0: - stop = last_offset + stop + 1 - - if start is not None: - if step >= 1: - cond_list.append(expr_with_offsets.offsets >= start) - else: - cond_list.append(expr_with_offsets.offsets <= start) - if stop is not None: - if step >= 1: - cond_list.append(expr_with_offsets.offsets < stop) - else: - cond_list.append(expr_with_offsets.offsets > stop) - if step > 1: - start = start if (start is not None) else 0 - cond_list.append((expr_with_offsets.offsets - start) % step == 0) - if step < 0: - start = start if (start is not None) else last_offset - cond_list.append((start - expr_with_offsets.offsets) % (-step) == 0) - - sliced_expr = expr_with_offsets._filter( - functools.reduce(lambda x, y: x & y, cond_list) - ) - return sliced_expr if step > 0 else sliced_expr.reversed() - def cached(self, cluster_cols: typing.Sequence[str]) -> ArrayValue: """Write the ArrayValue to a session table and create a new block object that references it.""" - ibis_expr = self._to_ibis_expr( - ordering_mode="unordered", expose_hidden_cols=True - ) + ibis_expr = self._to_ibis_expr("unordered", expose_hidden_cols=True) destination = self._session._ibis_to_session_table( ibis_expr, cluster_cols=cluster_cols, api_name="cache" ) table_expression = self._session.ibis_client.table( f"{destination.project}.{destination.dataset_id}.{destination.table_id}" ) - new_columns = [table_expression[column] for column in self.column_names] + new_columns = [table_expression[column] for column in self.column_ids] new_hidden_columns = [ table_expression[column] for column in self._hidden_ordering_column_names ] diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 0161d173618..852b1fe02ab 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -152,7 +152,7 @@ def value_columns(self) -> Sequence[str]: """All value columns, mutually exclusive with index columns.""" return [ column - for column in self._expr.column_names + for column in self._expr.column_ids if column not in self.index_columns ] @@ -444,9 +444,7 @@ def _compute_and_count( # TODO(swast): Allow for dry run and timeout. expr = self._apply_value_keys_to_expr(value_keys=value_keys) - results_iterator, query_job = expr.start_query( - max_results=max_results, expose_extra_columns=True - ) + results_iterator, query_job = expr.start_query(max_results=max_results) table_size = expr._get_table_size(query_job.destination) / _BYTES_TO_MEGABYTES fraction = ( @@ -483,12 +481,6 @@ def _compute_and_count( if self.index_columns: df.set_index(list(self.index_columns), inplace=True) df.index.names = self.index.names # type: ignore - - df.drop( - [col for col in df.columns if col not in self.value_columns], - axis=1, - inplace=True, - ) elif (sampling_method == _UNIFORM) and (random_state is None): filtered_expr = self.expr._uniform_sampling(fraction) block = Block( @@ -520,12 +512,6 @@ def _compute_and_count( df.set_index(list(self.index_columns), inplace=True) df.index.names = self.index.names # type: ignore - df.drop( - [col for col in df.columns if col not in self.value_columns], - axis=1, - inplace=True, - ) - return df, total_rows, query_job def _split( @@ -1087,7 +1073,7 @@ def _normalize_expression( ): """Normalizes expression by moving index columns to left.""" value_columns = [ - col_id for col_id in expr.column_names.keys() if col_id not in index_columns + col_id for col_id in expr.column_ids if col_id not in index_columns ] if (assert_value_size is not None) and ( len(value_columns) != assert_value_size @@ -1096,20 +1082,92 @@ def _normalize_expression( return expr.select_columns([*index_columns, *value_columns]) def slice( - self: bigframes.core.blocks.Block, + self, start: typing.Optional[int] = None, stop: typing.Optional[int] = None, step: typing.Optional[int] = None, ) -> bigframes.core.blocks.Block: - sliced_expr = self.expr.slice(start=start, stop=stop, step=step) - # since this is slice, return a copy even if unchanged - block = Block( - sliced_expr, - index_columns=self.index_columns, - column_labels=self.column_labels, - index_labels=self._index_labels, + if step is None: + step = 1 + if step == 0: + raise ValueError("slice step cannot be zero") + if step < 0: + reverse_start = (-start - 1) if start else 0 + reverse_stop = (-stop - 1) if stop else None + reverse_step = -step + return self.reversed()._forward_slice( + reverse_start, reverse_stop, reverse_step + ) + return self._forward_slice(start or 0, stop, step) + + def _forward_slice(self, start: int = 0, stop=None, step: int = 1): + """Performs slice but only for positive step size.""" + if step <= 0: + raise ValueError("forward_slice only supports positive step size") + + use_postive_offsets = ( + (start > 0) + or ((stop is not None) and (stop >= 0)) + or ((step > 1) and (start >= 0)) ) - return block + use_negative_offsets = ( + (start < 0) or (stop and (stop < 0)) or ((step > 1) and (start < 0)) + ) + + block = self + + # only generate offsets that are used + positive_offsets = None + negative_offsets = None + + if use_postive_offsets: + block, positive_offsets = self.promote_offsets() + if use_negative_offsets: + block, negative_offsets = block.reversed().promote_offsets() + block = block.reversed() + + conditions = [] + if start != 0: + if start > 0: + op = ops.partial_right(ops.ge_op, start) + assert positive_offsets + block, start_cond = block.apply_unary_op(positive_offsets, op) + else: + op = ops.partial_right(ops.le_op, -start - 1) + assert negative_offsets + block, start_cond = block.apply_unary_op(negative_offsets, op) + conditions.append(start_cond) + if stop is not None: + if stop >= 0: + op = ops.partial_right(ops.lt_op, stop) + assert positive_offsets + block, stop_cond = block.apply_unary_op(positive_offsets, op) + else: + op = ops.partial_right(ops.gt_op, -stop - 1) + assert negative_offsets + block, stop_cond = block.apply_unary_op(negative_offsets, op) + conditions.append(stop_cond) + + if step > 1: + op = ops.partial_right(ops.mod_op, step) + if start >= 0: + op = ops.partial_right(ops.sub_op, start) + assert positive_offsets + block, start_diff = block.apply_unary_op(positive_offsets, op) + else: + op = ops.partial_right(ops.sub_op, -start + 1) + assert negative_offsets + block, start_diff = block.apply_unary_op(negative_offsets, op) + modulo_op = ops.partial_right(ops.mod_op, step) + block, mod = block.apply_unary_op(start_diff, modulo_op) + is_zero_op = ops.partial_right(ops.eq_op, 0) + block, step_cond = block.apply_unary_op(mod, is_zero_op) + conditions.append(step_cond) + + for cond in conditions: + block = block.filter(cond) + + return block.select_columns(self.value_columns) # Using cache to optimize for Jupyter Notebook's behavior where both '__repr__' # and '__repr_html__' are called in a single display action, reducing redundant @@ -1396,7 +1454,7 @@ def concat( ) result_block = Block( result_expr, - index_columns=list(result_expr.column_names.keys())[:index_nlevels], + index_columns=list(result_expr.column_ids)[:index_nlevels], column_labels=aligned_blocks[0].column_labels, index_labels=result_labels, ) @@ -1530,9 +1588,7 @@ def to_sql_query( # the BigQuery unicode column name feature? substitutions[old_id] = new_id - sql = array_value.to_sql( - ordering_mode="unordered", col_id_overrides=substitutions - ) + sql = array_value.to_sql(col_id_overrides=substitutions) return ( sql, new_ids[: len(idx_labels)], diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 9be7f22a714..db0843fcbc0 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -426,10 +426,6 @@ def __init__( self._value_name = value_name self._dropna = dropna # Applies to aggregations but not windowing - @property - def _value(self): - return self._block.expr.get_column(self._value_column) - def all(self) -> series.Series: return self._aggregate(agg_ops.all_op) diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index f211afe4d56..7d15e676493 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -398,9 +398,7 @@ def to_pandas(self) -> pandas.Index: """Executes deferred operations and downloads the results.""" # Project down to only the index column. So the query can be cached to visualize other data. index_columns = list(self._block.index_columns) - expr = self._expr.projection( - [self._expr.get_any_column(col) for col in index_columns] - ) + expr = self._expr.select_columns(index_columns) results, _ = expr.start_query() df = expr._session._rows_to_dataframe(results) df = df.set_index(index_columns) diff --git a/bigframes/core/joins/row_identity.py b/bigframes/core/joins/row_identity.py index 66eb2239904..156e7aef407 100644 --- a/bigframes/core/joins/row_identity.py +++ b/bigframes/core/joins/row_identity.py @@ -38,11 +38,11 @@ def join_by_row_identity( f"Only how='outer','left','inner' currently supported. {constants.FEEDBACK_LINK}" ) - if not left.table.equals(right.table): + if not left._table.equals(right._table): raise ValueError( "Cannot combine objects without an explicit join/merge key. " - f"Left based on: {left.table.compile()}, but " - f"right based on: {right.table.compile()}" + f"Left based on: {left._table.compile()}, but " + f"right based on: {right._table.compile()}" ) left_predicates = left._predicates @@ -63,11 +63,11 @@ def join_by_row_identity( left_mask = left_relative_predicates if how in ["right", "outer"] else None right_mask = right_relative_predicates if how in ["left", "outer"] else None joined_columns = [ - _mask_value(left.get_column(key), left_mask).name(map_left_id(key)) - for key in left.column_names.keys() + _mask_value(left._get_ibis_column(key), left_mask).name(map_left_id(key)) + for key in left.column_ids ] + [ - _mask_value(right.get_column(key), right_mask).name(map_right_id(key)) - for key in right.column_names.keys() + _mask_value(right._get_ibis_column(key), right_mask).name(map_right_id(key)) + for key in right.column_ids ] # If left isn't being masked, can just use left ordering @@ -108,7 +108,7 @@ def join_by_row_identity( joined_expr = core.ArrayValue( left._session, - left.table, + left._table, columns=joined_columns, hidden_ordering_columns=hidden_ordering_columns, ordering=new_ordering, diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index 2d616fc3f0c..f194b8f8c41 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -74,14 +74,14 @@ def join_by_column( if ( allow_row_identity_join and how in bigframes.core.joins.row_identity.SUPPORTED_ROW_IDENTITY_HOW - and left.table.equals(right.table) + and left._table.equals(right._table) # Make sure we're joining on exactly the same column(s), at least with # regards to value its possible that they both have the same names but # were modified in different ways. Ignore differences in the names. and all( - left.get_any_column(lcol) + left._get_any_column(lcol) .name("index") - .equals(right.get_any_column(rcol).name("index")) + .equals(right._get_any_column(rcol).name("index")) for lcol, rcol in zip(left_column_ids, right_column_ids) ) ): @@ -90,14 +90,16 @@ def join_by_column( get_column_right, ) = bigframes.core.joins.row_identity.join_by_row_identity(left, right, how=how) left_join_keys = [ - combined_expr.get_column(get_column_left(col)) for col in left_column_ids + combined_expr._get_ibis_column(get_column_left(col)) + for col in left_column_ids ] right_join_keys = [ - combined_expr.get_column(get_column_right(col)) for col in right_column_ids + combined_expr._get_ibis_column(get_column_right(col)) + for col in right_column_ids ] join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how) join_key_ids = [col.get_name() for col in join_key_cols] - combined_expr = combined_expr.projection( + combined_expr = combined_expr._projection( [*join_key_cols, *combined_expr.columns] ) if sort: @@ -119,13 +121,13 @@ def join_by_column( lmapping = { col_id: guid.generate_guid() for col_id in itertools.chain( - left.column_names, left._hidden_ordering_column_names + left.column_ids, left._hidden_ordering_column_names ) } rmapping = { col_id: guid.generate_guid() for col_id in itertools.chain( - right.column_names, right._hidden_ordering_column_names + right.column_ids, right._hidden_ordering_column_names ) } @@ -136,12 +138,12 @@ def get_column_right(col_id): return rmapping[col_id] left_table = left._to_ibis_expr( - ordering_mode="unordered", + "unordered", expose_hidden_cols=True, col_id_overrides=lmapping, ) right_table = right._to_ibis_expr( - ordering_mode="unordered", + "unordered", expose_hidden_cols=True, col_id_overrides=rmapping, ) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 8874a4edb89..0492e62c153 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -554,7 +554,7 @@ def _apply_series_binop( other._block.index, how=how ) - series_column_id = other._value.get_name() + series_column_id = other._value_column series_col = get_column_right(series_column_id) block = joined_index._block for column_id, label in zip( @@ -2382,13 +2382,11 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: if ordering_id is not None: return array_value.to_sql( - ordering_mode="offset_col", + offset_column=ordering_id, col_id_overrides=id_overrides, - order_col_name=ordering_id, ) else: return array_value.to_sql( - ordering_mode="unordered", col_id_overrides=id_overrides, ) diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index 51eaad18b9b..fc76d07edbb 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -16,7 +16,6 @@ import typing -import ibis.expr.types as ibis_types import pandas as pd import bigframes.constants as constants @@ -106,11 +105,6 @@ def __init__( if pd_series.name is None: self._block = self._block.with_column_labels([None]) - @property - def _value(self) -> ibis_types.Value: - """Private property to get Ibis expression for the value column.""" - return self._block.expr.get_column(self._value_column) - @property def _value_column(self) -> str: return self._block.value_columns[0] diff --git a/bigframes/series.py b/bigframes/series.py index 2a0b1ff1fa4..717a85a93e8 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1150,7 +1150,11 @@ def _groupby_values( def apply(self, func) -> Series: # TODO(shobs, b/274645634): Support convert_dtype, args, **kwargs # is actually a ternary op - return self._apply_unary_op(ops.RemoteFunctionOp(func)) + # Reproject as workaround to applying filter too late. This forces the filter + # to be applied before passing data to remote function, protecting from bad + # inputs causing errors. + reprojected_series = Series(self._block._force_reproject()) + return reprojected_series._apply_unary_op(ops.RemoteFunctionOp(func)) def add_prefix(self, prefix: str, axis: int | str | None = None) -> Series: return Series(self._get_block().add_prefix(prefix)) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index ee0cefb3d28..69b9e798075 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ibis import ibis.expr.types as ibis_types import pandas @@ -45,7 +44,7 @@ def test_arrayvalue_constructor_from_ibis_table_adds_all_columns(): actual = core.ArrayValue( session=session, table=ibis_table, columns=columns, ordering=ordering ) - assert actual.table is ibis_table + assert actual._table is ibis_table assert len(actual.columns) == 3 @@ -79,37 +78,12 @@ def test_arrayvalue_with_get_column(): ), total_ordering_columns=["col1"], ) - col1 = value.get_column("col1") + col1 = value._get_ibis_column("col1") assert isinstance(col1, ibis_types.Value) assert col1.get_name() == "col1" assert col1.type().is_int64() -def test_arrayvalue_to_ibis_expr_with_projection(): - value = resources.create_arrayvalue( - pandas.DataFrame( - { - "col1": [1, 2, 3], - "col2": ["a", "b", "c"], - "col3": [0.1, 0.2, 0.3], - } - ), - total_ordering_columns=["col1"], - ) - expr = value.projection( - [ - (value.table["col1"] + ibis.literal(-1)).name("int64_col"), - ibis.literal(123456789).name("literals"), - value.table["col2"].name("string_col"), - ] - ) - actual = expr._to_ibis_expr() - assert len(actual.columns) == 3 - assert actual.columns[0] == "int64_col" - assert actual.columns[1] == "literals" - assert actual.columns[2] == "string_col" - - def test_arrayvalues_to_ibis_expr_with_get_column(): value = resources.create_arrayvalue( pandas.DataFrame( @@ -121,7 +95,7 @@ def test_arrayvalues_to_ibis_expr_with_get_column(): ), total_ordering_columns=["col1"], ) - expr = value.get_column("col1") + expr = value._get_ibis_column("col1") assert expr.get_name() == "col1" assert expr.type().is_int64() @@ -138,7 +112,7 @@ def test_arrayvalues_to_ibis_expr_with_concat(): total_ordering_columns=["col1"], ) expr = value.concat([value]) - actual = expr._to_ibis_expr() + actual = expr._to_ibis_expr("unordered") assert len(actual.columns) == 3 # TODO(ashleyxu, b/299631930): test out the union expression assert actual.columns[0] == "column_0" @@ -175,7 +149,7 @@ def test_arrayvalues_to_ibis_expr_with_project_binary_op(): ) expr = value.project_binary_op("col2", "col3", ops.add_op, "col4") assert expr.columns[3].type().is_float64() - actual = expr._to_ibis_expr() + actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 4 assert actual.columns[3] == "col4" @@ -194,7 +168,7 @@ def test_arrayvalues_to_ibis_expr_with_project_ternary_op(): ) expr = value.project_ternary_op("col2", "col3", "col4", ops.where_op, "col5") assert expr.columns[4].type().is_float64() - actual = expr._to_ibis_expr() + actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 5 assert actual.columns[4] == "col5" @@ -215,7 +189,7 @@ def test_arrayvalue_to_ibis_expr_with_aggregate(): by_column_ids=["col1"], dropna=False, ) - actual = expr._to_ibis_expr() + actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 2 assert actual.columns[0] == "col1" assert actual.columns[1] == "col4" @@ -234,7 +208,7 @@ def test_arrayvalue_to_ibis_expr_with_corr_aggregate(): total_ordering_columns=["col1"], ) expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]) - actual = expr._to_ibis_expr() + actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 1 assert actual.columns[0] == "col4" assert expr.columns[0].type().is_float64() From 636a209e0853501abd50784a11a87cf7f2282ee5 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 5 Oct 2023 14:09:14 -0700 Subject: [PATCH 62/80] feat: add equals methods to series/dataframe (#76) * feat: add equals methods to series/dataframe --- bigframes/core/block_transforms.py | 33 +++++++++ bigframes/dataframe.py | 6 ++ bigframes/operations/__init__.py | 10 +++ bigframes/series.py | 8 +++ tests/system/small/test_dataframe.py | 68 +++++++++++++++++++ tests/system/small/test_series.py | 38 +++++++++++ .../bigframes_vendored/pandas/core/frame.py | 22 ++++++ 7 files changed, 185 insertions(+) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 30c79029817..4c30d7631dc 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -25,6 +25,39 @@ import bigframes.operations.aggregations as agg_ops +def equals(block1: blocks.Block, block2: blocks.Block) -> bool: + if not block1.column_labels.equals(block2.column_labels): + return False + if block1.dtypes != block2.dtypes: + return False + # TODO: More advanced expression tree traversals to short circuit actually querying data + + block1 = block1.reset_index(drop=False) + block2 = block2.reset_index(drop=False) + + joined, (lmap, rmap) = block1.index.join(block2.index, how="outer") + joined_block = joined._block + + equality_ids = [] + for lcol, rcol in zip(block1.value_columns, block2.value_columns): + lcolmapped = lmap(lcol) + rcolmapped = rmap(rcol) + joined_block, result_id = joined_block.apply_binary_op( + lcolmapped, rcolmapped, ops.eq_nulls_match_op + ) + joined_block, result_id = joined_block.apply_unary_op( + result_id, ops.partial_right(ops.fillna_op, False) + ) + equality_ids.append(result_id) + + joined_block = joined_block.select_columns(equality_ids).with_column_labels( + list(range(len(equality_ids))) + ) + stacked_block = joined_block.stack(dropna=False, sort=False) + result = stacked_block.get_stat(stacked_block.value_columns[0], agg_ops.all_op) + return typing.cast(bool, result) + + def indicate_duplicates( block: blocks.Block, columns: typing.Sequence[str], keep: str = "first" ) -> typing.Tuple[blocks.Block, str]: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 0492e62c153..87d0f21b629 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1066,6 +1066,12 @@ def rename_axis( labels = [mapper] return DataFrame(self._block.with_index_labels(labels)) + def equals(self, other: typing.Union[bigframes.series.Series, DataFrame]) -> bool: + # Must be same object type, same column dtypes, and same label values + if not isinstance(other, DataFrame): + return False + return block_ops.equals(self._block, other._block) + def assign(self, **kwargs) -> DataFrame: # TODO(garrettwu) Support list-like values. Requires ordering. # TODO(garrettwu) Support callable values. diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index bc08298eb7b..f330a703b20 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -705,6 +705,16 @@ def eq_op( return x == y +def eq_nulls_match_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + """Variant of eq_op where nulls match each other. Only use where dtypes are known to be same.""" + left = x.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$")) + right = y.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$")) + return left == right + + def ne_op( x: ibis_types.Value, y: ibis_types.Value, diff --git a/bigframes/series.py b/bigframes/series.py index 717a85a93e8..e5afe91e445 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -209,6 +209,14 @@ def rename_axis( labels = [mapper] return Series(self._block.with_index_labels(labels)) + def equals( + self, other: typing.Union[Series, bigframes.dataframe.DataFrame] + ) -> bool: + # Must be same object type, same column dtypes, and same label values + if not isinstance(other, Series): + return False + return block_ops.equals(self._block, other._block) + def reset_index( self, *, diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 19ea9b8ae55..711da10c55c 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2551,6 +2551,74 @@ def test_df_reindex_columns(scalars_df_index, scalars_pandas_df_index): ) +def test_df_equals_identical(scalars_df_index, scalars_pandas_df_index): + unsupported = [ + "geography_col", + ] + scalars_df_index = scalars_df_index.drop(columns=unsupported) + scalars_pandas_df_index = scalars_pandas_df_index.drop(columns=unsupported) + + bf_result = scalars_df_index.equals(scalars_df_index) + pd_result = scalars_pandas_df_index.equals(scalars_pandas_df_index) + + assert pd_result == bf_result + + +def test_df_equals_series(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index[["int64_col"]].equals(scalars_df_index["int64_col"]) + pd_result = scalars_pandas_df_index[["int64_col"]].equals( + scalars_pandas_df_index["int64_col"] + ) + + assert pd_result == bf_result + + +def test_df_equals_different_dtype(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_col", "int64_too"] + scalars_df_index = scalars_df_index[columns] + scalars_pandas_df_index = scalars_pandas_df_index[columns] + + bf_modified = scalars_df_index.copy() + bf_modified = bf_modified.astype("Float64") + + pd_modified = scalars_pandas_df_index.copy() + pd_modified = pd_modified.astype("Float64") + + bf_result = scalars_df_index.equals(bf_modified) + pd_result = scalars_pandas_df_index.equals(pd_modified) + + assert pd_result == bf_result + + +def test_df_equals_different_values(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_col", "int64_too"] + scalars_df_index = scalars_df_index[columns] + scalars_pandas_df_index = scalars_pandas_df_index[columns] + + bf_modified = scalars_df_index.copy() + bf_modified["int64_col"] = bf_modified.int64_col + 1 + + pd_modified = scalars_pandas_df_index.copy() + pd_modified["int64_col"] = pd_modified.int64_col + 1 + + bf_result = scalars_df_index.equals(bf_modified) + pd_result = scalars_pandas_df_index.equals(pd_modified) + + assert pd_result == bf_result + + +def test_df_equals_extra_column(scalars_df_index, scalars_pandas_df_index): + columns = ["int64_col", "int64_too"] + more_columns = ["int64_col", "int64_too", "float64_col"] + + bf_result = scalars_df_index[columns].equals(scalars_df_index[more_columns]) + pd_result = scalars_pandas_df_index[columns].equals( + scalars_pandas_df_index[more_columns] + ) + + assert pd_result == bf_result + + def test_df_reindex_like(scalars_df_index, scalars_pandas_df_index): reindex_target_bf = scalars_df_index.reindex( columns=["not_a_col", "int64_col", "int64_too"], index=[5, 1, 3, 99, 1] diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 8c1c36720bf..993df18c959 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -112,6 +112,44 @@ def test_series_get_column_default(scalars_dfs): assert result == "default_val" +def test_series_equals_identical(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index.int64_col.equals(scalars_df_index.int64_col) + pd_result = scalars_pandas_df_index.int64_col.equals( + scalars_pandas_df_index.int64_col + ) + + assert pd_result == bf_result + + +def test_series_equals_df(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_col"].equals(scalars_df_index[["int64_col"]]) + pd_result = scalars_pandas_df_index["int64_col"].equals( + scalars_pandas_df_index[["int64_col"]] + ) + + assert pd_result == bf_result + + +def test_series_equals_different_dtype(scalars_df_index, scalars_pandas_df_index): + bf_series = scalars_df_index["int64_col"] + pd_series = scalars_pandas_df_index["int64_col"] + + bf_result = bf_series.equals(bf_series.astype("Float64")) + pd_result = pd_series.equals(pd_series.astype("Float64")) + + assert pd_result == bf_result + + +def test_series_equals_different_values(scalars_df_index, scalars_pandas_df_index): + bf_series = scalars_df_index["int64_col"] + pd_series = scalars_pandas_df_index["int64_col"] + + bf_result = bf_series.equals(bf_series + 1) + pd_result = pd_series.equals(pd_series + 1) + + assert pd_result == bf_result + + def test_series_get_with_default_index(scalars_dfs): col_name = "float64_col" key = 2 diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 17d941fbdd0..644e043e83d 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -399,6 +399,28 @@ def to_orc(self, path=None, **kwargs) -> bytes | None: # ---------------------------------------------------------------------- # Unsorted + def equals(self, other) -> bool: + """ + Test whether two objects contain the same elements. + + This function allows two Series or DataFrames to be compared against + each other to see if they have the same shape and elements. NaNs in + the same location are considered equal. + + The row/column index do not need to have the same type, as long + as the values are considered equal. Corresponding columns must be of + the same dtype. + + Args: + other (Series or DataFrame): + The other Series or DataFrame to be compared with the first. + + Returns: + bool: True if all elements are the same in both objects, False + otherwise. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def assign(self, **kwargs) -> DataFrame: r""" Assign new columns to a DataFrame. From 969800d669204de4d0f2e5e61da521217e55668b Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 6 Oct 2023 19:41:23 -0700 Subject: [PATCH 63/80] fix: change the invalid url in docs (#93) * fix: change the invalid url in docs * fix: address the comment --- bigframes/ml/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bigframes/ml/base.py b/bigframes/ml/base.py index f899ac71193..f2478b1ce23 100644 --- a/bigframes/ml/base.py +++ b/bigframes/ml/base.py @@ -104,7 +104,8 @@ def predict(self, X): def register(self: _T, vertex_ai_model_id: Optional[str] = None) -> _T: """Register the model to Vertex AI. - After register, go to https://pantheon.corp.google.com/vertex-ai/models to manage the model registries. + After register, go to Google Cloud Console (https://console.cloud.google.com/vertex-ai/models) + to manage the model registries. Refer to https://cloud.google.com/vertex-ai/docs/model-registry/introduction for more options. Args: From 5b0ee39bcffb8d39e28ea10d024f7b277e62cc18 Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Mon, 9 Oct 2023 11:12:56 -0400 Subject: [PATCH 64/80] chore(deps): bump urllib3 from 1.26.12 to 1.26.17 in /.kokoro (#96) Source-Link: https://github.com/googleapis/synthtool/commit/0c7b0333f44b2b7075447f43a121a12d15a7b76a Post-Processor: gcr.io/cloud-devrel-public-resources/owlbot-python:latest@sha256:08e34975760f002746b1d8c86fdc90660be45945ee6d9db914d1508acdf9a547 Co-authored-by: Owl Bot --- .github/.OwlBot.lock.yaml | 4 ++-- .kokoro/requirements.txt | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index a9bdb1b7ac0..dd98abbdeeb 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,5 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:fac304457974bb530cc5396abd4ab25d26a469cd3bc97cbfb18c8d4324c584eb -# created: 2023-10-02T21:31:03.517640371Z + digest: sha256:08e34975760f002746b1d8c86fdc90660be45945ee6d9db914d1508acdf9a547 +# created: 2023-10-09T14:06:13.397766266Z diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 96d593c8c82..0332d3267e1 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -467,9 +467,9 @@ typing-extensions==4.4.0 \ --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \ --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e # via -r requirements.in -urllib3==1.26.12 \ - --hash=sha256:3fa96cf423e6987997fc326ae8df396db2a8b7c667747d47ddd8ecba91f4a74e \ - --hash=sha256:b930dd878d5a8afb066a637fbb35144fe7901e3b209d1cd4f524bd0e9deee997 +urllib3==1.26.17 \ + --hash=sha256:24d6a242c28d29af46c3fae832c36db3bbebcc533dd1bb549172cd739c82df21 \ + --hash=sha256:94a757d178c9be92ef5539b8840d48dc9cf1b2709c9d6b588232a055c524458b # via # requests # twine From 159231505f339173560cd802dae3fed3e63a663b Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Mon, 9 Oct 2023 17:19:26 -0700 Subject: [PATCH 65/80] docs: add more preprocessing models into the docs menu. (#97) * docs: add more preprocessing models into the docs menu. * Trigger kokoro rebuild --- docs/templates/toc.yml | 12 ++++++-- .../sklearn/preprocessing/_data.py | 2 +- .../sklearn/preprocessing/_encoder.py | 30 +++++++++---------- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 891f15a51b4..0758bb41d8a 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -121,10 +121,18 @@ - items: - name: Overview uid: bigframes.ml.preprocessing - - name: OneHotEncoder - uid: bigframes.ml.preprocessing.OneHotEncoder + - name: KBinsDiscretizer + uid: bigframes.ml.preprocessing.KBinsDiscretizer + - name: LabelEncoder + uid: bigframes.ml.preprocessing.LabelEncoder + - name: MaxAbsScaler + uid: bigframes.ml.preprocessing.MaxAbsScaler + - name: MinMaxScaler + uid: bigframes.ml.preprocessing.MinMaxScaler - name: StandardScaler uid: bigframes.ml.preprocessing.StandardScaler + - name: OneHotEncoder + uid: bigframes.ml.preprocessing.OneHotEncoder name: preprocessing name: bigframes.ml name: BigQuery DataFrames diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py index 58e16e135b7..5ce102d573f 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_data.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_data.py @@ -29,7 +29,7 @@ class StandardScaler(BaseEstimator, TransformerMixin): individual features do not more or less look like standard normally distributed data (e.g. Gaussian with 0 mean and unit variance). - Examples: + **Examples:** .. code-block:: diff --git a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py index cf660ece5d4..8da9a98c53f 100644 --- a/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py +++ b/third_party/bigframes_vendored/sklearn/preprocessing/_encoder.py @@ -17,6 +17,21 @@ class OneHotEncoder(BaseEstimator): Note that this method deviates from Scikit-Learn; instead of producing sparse binary columns, the encoding is a single column of `STRUCT`. + **Examples:** + + Given a dataset with two features, we let the encoder find the unique + values per feature and transform the data to a binary one-hot encoding. + + .. code-block:: + + from bigframes.ml.preprocessing import OneHotEncoder + import bigframes.pandas as bpd + + enc = OneHotEncoder() + X = bpd.DataFrame({"a": ["Male", "Female", "Female"], "b": ["1", "3", "2"]}) + enc.fit(X) + print(enc.transform(bpd.DataFrame({"a": ["Female", "Male"], "b": ["1", "4"]}))) + Args: drop (Optional[Literal["most_frequent"]], default None): Specifies a methodology to use to drop one of the categories per feature. @@ -37,21 +52,6 @@ class OneHotEncoder(BaseEstimator): when considering infrequent categories. If there are infrequent categories, max_categories includes the category representing the infrequent categories along with the frequent categories. Default None, set limit to 1,000,000. - - Examples: - - Given a dataset with two features, we let the encoder find the unique - values per feature and transform the data to a binary one-hot encoding. - - .. code-block:: - - from bigframes.ml.preprocessing import OneHotEncoder - import bigframes.pandas as bpd - - enc = OneHotEncoder() - X = bpd.DataFrame({"a": ["Male", "Female", "Female"], "b": ["1", "3", "2"]}) - enc.fit(X) - print(enc.transform(bpd.DataFrame({"a": ["Female", "Male"], "b": ["1", "4"]}))) """ def fit(self, X, y=None): From e1ee09b51acd75656f243a54027fd4154ef4bfd1 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 10 Oct 2023 09:56:14 -0700 Subject: [PATCH 66/80] refactor: refactor the directory of the notebooks (#89) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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://togithub.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 # 🦕 --- .../bq_dataframes_llm_code_generation.ipynb | 392 +++++++++--------- .../bq_dataframes_ml_linear_regression.ipynb | 6 +- noxfile.py | 4 +- 3 files changed, 201 insertions(+), 201 deletions(-) rename notebooks/{getting_started => generative_ai}/bq_dataframes_llm_code_generation.ipynb (95%) rename notebooks/{getting_started => regression}/bq_dataframes_ml_linear_regression.ipynb (98%) diff --git a/notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb similarity index 95% rename from notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb rename to notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index 39e2ef535cc..2e4ce3e5109 100644 --- a/notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -34,18 +34,18 @@ "\n", "\n", " \n", " \n", "
\n", - " \n", + " \n", " \"Colab Run in Colab\n", " \n", " \n", - " \n", + " \n", " \"GitHub\n", " View on GitHub\n", " \n", " \n", - " \n", + " \n", " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", @@ -162,6 +162,9 @@ }, { "cell_type": "markdown", + "metadata": { + "id": "Wbr2aVtFQBcg" + }, "source": [ "### Set up your Google Cloud project\n", "\n", @@ -183,10 +186,7 @@ " * Vertex AI API\n", "\n", "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)." - ], - "metadata": { - "id": "Wbr2aVtFQBcg" - } + ] }, { "cell_type": "markdown", @@ -350,39 +350,44 @@ }, { "cell_type": "markdown", - "source": [ - "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." - ], "metadata": { "id": "DTVtFlqeFbrU" - } + }, + "source": [ + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + ] }, { "cell_type": "markdown", + "metadata": { + "id": "6eytf4xQHzcF" + }, "source": [ "# Define the LLM model\n", "\n", "BigQuery DataFrames provides integration with [`text-bison` model of the PaLM API](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text) via Vertex AI.\n", "\n", "This section walks through a few steps required in order to use the model in your notebook." - ], - "metadata": { - "id": "6eytf4xQHzcF" - } + ] }, { "cell_type": "markdown", + "metadata": { + "id": "rS4VO1TGiO4G" + }, "source": [ "## Create a BigQuery Cloud resource connection\n", "\n", "You need to create a [Cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) to enable BigQuery DataFrames to interact with Vertex AI services." - ], - "metadata": { - "id": "rS4VO1TGiO4G" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "KFPjDM4LVh96" + }, + "outputs": [], "source": [ "CONN_NAME = \"bqdf-llm\"\n", "\n", @@ -412,15 +417,13 @@ " f\"serviceAccount:{response.cloud_resource.service_account_id}\"\n", " )\n", "print(CONN_SERVICE_ACCOUNT)" - ], - "metadata": { - "id": "KFPjDM4LVh96" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "W6l6Ol2biU9h" + }, "source": [ "## Set permissions for the service account\n", "\n", @@ -429,52 +432,52 @@ " - `roles/run.invoker`: This role is required for the connection to have read-only access to Cloud Run services that back custom/remote functions ([documentation](https://cloud.google.com/bigquery/docs/remote-functions#grant_permission_on_function)).\n", "\n", "Set these permissions by running the following `gcloud` commands:" - ], - "metadata": { - "id": "W6l6Ol2biU9h" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d8wja24SVq6s" + }, + "outputs": [], "source": [ "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/bigquery.connectionUser'\n", "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/aiplatform.user'\n", "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/run.invoker'" - ], - "metadata": { - "id": "d8wja24SVq6s" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "qUjT8nw-jIXp" + }, "source": [ "## Define the model\n", "\n", "Use `bigframes.ml.llm` to define the model:" - ], - "metadata": { - "id": "qUjT8nw-jIXp" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sdjeXFwcHfl7" + }, + "outputs": [], "source": [ "from bigframes.ml.llm import PaLM2TextGenerator\n", "\n", "session = bf.get_global_session()\n", "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", "model = PaLM2TextGenerator(session=session, connection_name=connection)" - ], - "metadata": { - "id": "sdjeXFwcHfl7" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "GbW0oCnU1s1N" + }, "source": [ "# Read data from Cloud Storage into BigQuery DataFrames\n", "\n", @@ -486,80 +489,82 @@ "* An in-memory pandas DataFrame\n", "\n", "In this tutorial, you create BigQuery DataFrames DataFrames by reading two CSV files stored in Cloud Storage, one containing a list of DataFrame API names and one containing a list of Series API names." - ], - "metadata": { - "id": "GbW0oCnU1s1N" - } + ] }, { "cell_type": "code", - "source": [ - "df_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/df.csv\")\n", - "series_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/series.csv\")" - ], + "execution_count": null, "metadata": { "id": "SchiTkQGIJog" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "df_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/df.csv\")\n", + "series_api = bf.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/series.csv\")" + ] }, { "cell_type": "markdown", - "source": [ - "Take a peek at a few rows of data for each file:" - ], "metadata": { "id": "7OBjw2nmQY3-" - } + }, + "source": [ + "Take a peek at a few rows of data for each file:" + ] }, { "cell_type": "code", - "source": [ - "df_api.head(2)" - ], + "execution_count": null, "metadata": { "id": "QCqgVCIsGGuv" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "df_api.head(2)" + ] }, { "cell_type": "code", - "source": [ - "series_api.head(2)" - ], + "execution_count": null, "metadata": { "id": "BGJnZbgEGS5-" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "series_api.head(2)" + ] }, { "cell_type": "markdown", + "metadata": { + "id": "m3ZJEsi7SUKV" + }, "source": [ "# Generate code using the LLM model\n", "\n", "Prepare the prompts and send them to the LLM model for prediction." - ], - "metadata": { - "id": "m3ZJEsi7SUKV" - } + ] }, { "cell_type": "markdown", + "metadata": { + "id": "9EMAqR37AfLS" + }, "source": [ "## Prompt design in BigQuery DataFrames\n", "\n", "Designing prompts for LLMs is a fast growing area and you can read more in [this documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/introduction-prompt-design).\n", "\n", "For this tutorial, you use a simple prompt to ask the LLM model for sample code for each of the API methods (or rows) from the last step's DataFrames. The output is the new DataFrames `df_prompt` and `series_prompt`, which contain the full prompt text." - ], - "metadata": { - "id": "9EMAqR37AfLS" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EDAaIwHpQCDZ" + }, + "outputs": [], "source": [ "df_prompt_prefix = \"Generate Pandas sample code for DataFrame.\"\n", "series_prompt_prefix = \"Generate Pandas sample code for Series.\"\n", @@ -568,83 +573,83 @@ "series_prompt = (series_prompt_prefix + series_api['API'])\n", "\n", "df_prompt.head(2)" - ], - "metadata": { - "id": "EDAaIwHpQCDZ" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "rwPLjqW2Ajzh" + }, "source": [ "## Make predictions using the LLM model\n", "\n", "Use the BigQuery DataFrames DataFrame containing the full prompt text as the input to the `predict` method. The `predict` method calls the LLM model and returns its generated text output back to two new BigQuery DataFrames DataFrames, `df_pred` and `series_pred`.\n", "\n", "Note: The predictions might take a few minutes to run." - ], - "metadata": { - "id": "rwPLjqW2Ajzh" - } + ] }, { "cell_type": "code", - "source": [ - "df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024)\n", - "series_pred = model.predict(series_prompt.to_frame(), max_output_tokens=1024)" - ], + "execution_count": null, "metadata": { "id": "6i6HkFJZa8na" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024)\n", + "series_pred = model.predict(series_prompt.to_frame(), max_output_tokens=1024)" + ] }, { "cell_type": "markdown", - "source": [ - "Once the predictions are processed, take a look at the sample output from the LLM, which provides code samples for the API names listed in the DataFrames dataset." - ], "metadata": { "id": "89cB8MW4UIdV" - } + }, + "source": [ + "Once the predictions are processed, take a look at the sample output from the LLM, which provides code samples for the API names listed in the DataFrames dataset." + ] }, { "cell_type": "code", - "source": [ - "print(df_pred['ml_generate_text_llm_result'].iloc[0])" - ], + "execution_count": null, "metadata": { "id": "9A2gw6hP_2nX" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "print(df_pred['ml_generate_text_llm_result'].iloc[0])" + ] }, { "cell_type": "markdown", + "metadata": { + "id": "Fx4lsNqMorJ-" + }, "source": [ "# Manipulate LLM output using a remote function\n", "\n", "The output that the LLM provides often contains additional text beyond the code sample itself. Using BigQuery DataFrames, you can deploy custom Python functions that process and transform this output.\n", "\n" - ], - "metadata": { - "id": "Fx4lsNqMorJ-" - } + ] }, { "cell_type": "markdown", + "metadata": { + "id": "d8L7SN03VByG" + }, "source": [ "Running the cell below creates a custom function that you can use to process the LLM output data in two ways:\n", "1. Strip the LLM text output to include only the code block.\n", "2. Substitute `import pandas as pd` with `import bigframes.pandas as bf` so that the resulting code block works with BigQuery DataFrames." - ], - "metadata": { - "id": "d8L7SN03VByG" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "GskyyUQPowBT" + }, + "outputs": [], "source": [ "@bf.remote_function([str], str, bigquery_connection=CONN_NAME)\n", "def extract_code(text: str):\n", @@ -656,166 +661,161 @@ " return res\n", " except:\n", " return \"\"" - ], - "metadata": { - "id": "GskyyUQPowBT" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "source": [ - "The custom function is deployed as a Cloud Function, and then integrated with BigQuery as a [remote function](https://cloud.google.com/bigquery/docs/remote-functions). Save both of the function names so that you can clean them up at the end of this notebook." - ], "metadata": { "id": "hVQAoqBUOJQf" - } + }, + "source": [ + "The custom function is deployed as a Cloud Function, and then integrated with BigQuery as a [remote function](https://cloud.google.com/bigquery/docs/remote-functions). Save both of the function names so that you can clean them up at the end of this notebook." + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PBlp-C-DOHRO" + }, + "outputs": [], "source": [ "CLOUD_FUNCTION_NAME = format(extract_code.bigframes_cloud_function)\n", "print(\"Cloud Function Name \" + CLOUD_FUNCTION_NAME)\n", "REMOTE_FUNCTION_NAME = format(extract_code.bigframes_remote_function)\n", "print(\"Remote Function Name \" + REMOTE_FUNCTION_NAME)" - ], - "metadata": { - "id": "PBlp-C-DOHRO" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "source": [ - "Apply the custom function to each LLM output DataFrame to get the processed results:" - ], "metadata": { "id": "4FEucaiqVs3H" - } + }, + "source": [ + "Apply the custom function to each LLM output DataFrame to get the processed results:" + ] }, { "cell_type": "code", - "source": [ - "df_code = df_pred.assign(code=df_pred['ml_generate_text_llm_result'].apply(extract_code))\n", - "series_code = series_pred.assign(code=series_pred['ml_generate_text_llm_result'].apply(extract_code))" - ], + "execution_count": null, "metadata": { "id": "bsQ9cmoWo0Ps" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "df_code = df_pred.assign(code=df_pred['ml_generate_text_llm_result'].apply(extract_code))\n", + "series_code = series_pred.assign(code=series_pred['ml_generate_text_llm_result'].apply(extract_code))" + ] }, { "cell_type": "markdown", - "source": [ - "You can see the differences by inspecting the first row of data:" - ], "metadata": { "id": "ujQVVuhfWA3y" - } + }, + "source": [ + "You can see the differences by inspecting the first row of data:" + ] }, { "cell_type": "code", - "source": [ - "print(df_code['code'].iloc[0])" - ], + "execution_count": null, "metadata": { "id": "7yWzjhGy_zcy" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "print(df_code['code'].iloc[0])" + ] }, { "cell_type": "markdown", + "metadata": { + "id": "GTRdUw-Ro5R1" + }, "source": [ "# Save the results to Cloud Storage\n", "\n", "BigQuery DataFrames lets you save a BigQuery DataFrames DataFrame as a CSV file in Cloud Storage for further use. Try that now with your processed LLM output data." - ], - "metadata": { - "id": "GTRdUw-Ro5R1" - } + ] }, { "cell_type": "markdown", - "source": [ - "Create a new Cloud Storage bucket with a unique name:" - ], "metadata": { "id": "9DQ7eiQxPTi3" - } + }, + "source": [ + "Create a new Cloud Storage bucket with a unique name:" + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "-J5LHgS6LLZ0" + }, + "outputs": [], "source": [ "import uuid\n", "BUCKET_ID = \"code-samples-\" + str(uuid.uuid1())\n", "\n", "!gsutil mb gs://{BUCKET_ID}" - ], - "metadata": { - "id": "-J5LHgS6LLZ0" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "source": [ - "Use `to_csv` to write each BigQuery DataFrames DataFrame as a CSV file in the Cloud Storage bucket:" - ], "metadata": { "id": "tyxZXj0UPYUv" - } + }, + "source": [ + "Use `to_csv` to write each BigQuery DataFrames DataFrame as a CSV file in the Cloud Storage bucket:" + ] }, { "cell_type": "code", - "source": [ - "df_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/df_code*.csv\")\n", - "series_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/series_code*.csv\")" - ], + "execution_count": null, "metadata": { "id": "Zs_b5L-4IvER" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "df_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/df_code*.csv\")\n", + "series_code[[\"code\"]].to_csv(f\"gs://{BUCKET_ID}/series_code*.csv\")" + ] }, { "cell_type": "markdown", + "metadata": { + "id": "UDBtDlrTuuh8" + }, "source": [ "You can navigate to the Cloud Storage bucket browser to download the two files and view them.\n", "\n", "Run the following cell, and then follow the link to your Cloud Storage bucket browser:" - ], - "metadata": { - "id": "UDBtDlrTuuh8" - } + ] }, { "cell_type": "code", - "source": [ - "print(f'https://console.developers.google.com/storage/browser/{BUCKET_ID}/')" - ], + "execution_count": null, "metadata": { "id": "PspCXu-qu_ND" }, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "print(f'https://console.developers.google.com/storage/browser/{BUCKET_ID}/')" + ] }, { "cell_type": "markdown", + "metadata": { + "id": "RGSvUk48RK20" + }, "source": [ "# Summary and next steps\n", "\n", "You've used BigQuery DataFrames' integration with LLM models (`bigframes.ml.llm`) to generate code samples, and have tranformed LLM output by creating and using a custom function in BigQuery DataFrames.\n", "\n", "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks)." - ], - "metadata": { - "id": "RGSvUk48RK20" - } + ] }, { "cell_type": "markdown", @@ -833,6 +833,11 @@ }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "yw7A461XLjvW" + }, + "outputs": [], "source": [ "# # Delete the BigQuery Connection\n", "# from google.cloud import bigquery_connection_v1 as bq_connection\n", @@ -840,12 +845,7 @@ "# CONNECTION_ID = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}\"\n", "# client.delete_connection(name=CONNECTION_ID)\n", "# print(f\"Deleted connection '{CONNECTION_ID}'.\")" - ], - "metadata": { - "id": "yw7A461XLjvW" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", @@ -864,22 +864,22 @@ }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "iQFo6OUBLmi3" + }, + "outputs": [], "source": [ "# # Delete the Google Cloud Storage bucket and files\n", "# ! gsutil rm -r gs://{BUCKET_ID}\n", "# print(f\"Deleted bucket '{BUCKET_ID}'.\")" - ], - "metadata": { - "id": "iQFo6OUBLmi3" - }, - "execution_count": null, - "outputs": [] + ] } ], "metadata": { "colab": { - "toc_visible": true, - "provenance": [] + "provenance": [], + "toc_visible": true }, "kernelspec": { "display_name": "Python 3", diff --git a/notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb b/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb similarity index 98% rename from notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb rename to notebooks/regression/bq_dataframes_ml_linear_regression.ipynb index d317217810f..338d6edf4fb 100644 --- a/notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb +++ b/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb @@ -35,18 +35,18 @@ "\n", "\n", " \n", " \n", "
\n", - " \n", + " \n", " \"Colab Run in Colab\n", " \n", " \n", - " \n", + " \n", " \"GitHub\n", " View on GitHub\n", " \n", " \n", - " \n", + " \n", " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", diff --git a/noxfile.py b/noxfile.py index a113e1fcde4..84e5ab11bb4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -607,8 +607,8 @@ def notebook(session): # appropriate values and omitting cleanup logic that may break # our test infrastructure. "notebooks/getting_started/getting_started_bq_dataframes.ipynb", - "notebooks/getting_started/bq_dataframes_llm_code_generation.ipynb", - "notebooks/getting_started/bq_dataframes_ml_linear_regression.ipynb", + "notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb", + "notebooks/regression/bq_dataframes_ml_linear_regression.ipynb", "notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb", "notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb", "notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb", From 228aeba09782ae2421040c7601c15d4af92790b6 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 10 Oct 2023 11:00:03 -0700 Subject: [PATCH 67/80] feat: add iat and iloc accessing by tuples of integers (#90) * feat: add iat indexers * select_columns correction * add iloc for tuples * add dataframe tests * add series tests * add third party comments --------- Co-authored-by: Henry J Solberg --- bigframes/core/indexers.py | 44 +++++++++++++++ bigframes/dataframe.py | 4 ++ bigframes/series.py | 4 ++ tests/system/small/test_dataframe.py | 55 ++++++++++++++++++- tests/system/small/test_series.py | 14 +++++ .../bigframes_vendored/pandas/core/frame.py | 10 ++++ .../bigframes_vendored/pandas/core/series.py | 10 ++++ 7 files changed, 140 insertions(+), 1 deletion(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 1a88b2abd60..a74880041ca 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -97,6 +97,16 @@ def __getitem__( return _iloc_getitem_series_or_dataframe(self._series, key) +class IatSeriesIndexer: + def __init__(self, series: bigframes.series.Series): + self._series = series + + def __getitem__(self, key: int) -> bigframes.core.scalar.Scalar: + if not isinstance(key, int): + raise ValueError("Series iAt based indexing can only have integer indexers") + return self._series.iloc[key] + + class LocDataFrameIndexer: def __init__(self, dataframe: bigframes.dataframe.DataFrame): self._dataframe = dataframe @@ -188,6 +198,28 @@ def __getitem__(self, key) -> Union[bigframes.dataframe.DataFrame, pd.Series]: return _iloc_getitem_series_or_dataframe(self._dataframe, key) +class IatDataFrameIndexer: + def __init__(self, dataframe: bigframes.dataframe.DataFrame): + self._dataframe = dataframe + + def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar: + error_message = "DataFrame.iat should be indexed by a tuple of exactly 2 ints" + # we raise TypeError or ValueError under the same conditions that pandas does + if isinstance(key, int): + raise TypeError(error_message) + if not isinstance(key, tuple): + raise ValueError(error_message) + key_values_are_ints = [isinstance(key_value, int) for key_value in key] + if not all(key_values_are_ints): + raise ValueError(error_message) + if len(key) != 2: + raise TypeError(error_message) + block = self._dataframe._block + column_block = block.select_columns([block.value_columns[key[1]]]) + column = bigframes.series.Series(column_block) + return column.iloc[key[0]] + + @typing.overload def _loc_getitem_series_or_dataframe( series_or_dataframe: bigframes.series.Series, key @@ -356,6 +388,18 @@ def _iloc_getitem_series_or_dataframe( return result_pd_df.iloc[0] elif isinstance(key, slice): return series_or_dataframe._slice(key.start, key.stop, key.step) + elif isinstance(key, tuple) and len(key) == 0: + return series_or_dataframe + elif isinstance(key, tuple) and len(key) == 1: + return _iloc_getitem_series_or_dataframe(series_or_dataframe, key[0]) + elif ( + isinstance(key, tuple) + and isinstance(series_or_dataframe, bigframes.dataframe.DataFrame) + and len(key) == 2 + ): + return series_or_dataframe.iat[key] + elif isinstance(key, tuple): + raise pd.errors.IndexingError("Too many indexers") elif pd.api.types.is_list_like(key): if len(key) == 0: return typing.cast( diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 87d0f21b629..d5b2fa86e9d 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -254,6 +254,10 @@ def loc(self) -> indexers.LocDataFrameIndexer: def iloc(self) -> indexers.ILocDataFrameIndexer: return indexers.ILocDataFrameIndexer(self) + @property + def iat(self) -> indexers.IatDataFrameIndexer: + return indexers.IatDataFrameIndexer(self) + @property def dtypes(self) -> pandas.Series: return pandas.Series(data=self._block.dtypes, index=self._block.column_labels) diff --git a/bigframes/series.py b/bigframes/series.py index e5afe91e445..56e1b43a032 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -87,6 +87,10 @@ def loc(self) -> bigframes.core.indexers.LocSeriesIndexer: def iloc(self) -> bigframes.core.indexers.IlocSeriesIndexer: return bigframes.core.indexers.IlocSeriesIndexer(self) + @property + def iat(self) -> bigframes.core.indexers.IatSeriesIndexer: + return bigframes.core.indexers.IatSeriesIndexer(self) + @property def name(self) -> blocks.Label: return self._name diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 711da10c55c..0262a31a92e 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2077,7 +2077,7 @@ def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index): @pytest.mark.parametrize( "index", - [0, 5, -2], + [0, 5, -2, (2,)], ) def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): bf_result = scalars_df_index.iloc[index] @@ -2089,6 +2089,59 @@ def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index): ) +@pytest.mark.parametrize( + "index", + [(2, 5), (5, 0), (0, 0)], +) +def test_iloc_tuple(scalars_df_index, scalars_pandas_df_index, index): + bf_result = scalars_df_index.iloc[index] + pd_result = scalars_pandas_df_index.iloc[index] + + assert bf_result == pd_result + + +@pytest.mark.parametrize( + ("index", "error"), + [ + ((1, 1, 1), pd.errors.IndexingError), + (("asd", "asd", "asd"), pd.errors.IndexingError), + (("asd"), TypeError), + ], +) +def test_iloc_tuple_errors(scalars_df_index, scalars_pandas_df_index, index, error): + with pytest.raises(error): + scalars_df_index.iloc[index] + with pytest.raises(error): + scalars_pandas_df_index.iloc[index] + + +@pytest.mark.parametrize( + "index", + [(2, 5), (5, 0), (0, 0)], +) +def test_iat(scalars_df_index, scalars_pandas_df_index, index): + bf_result = scalars_df_index.iat[index] + pd_result = scalars_pandas_df_index.iat[index] + + assert bf_result == pd_result + + +@pytest.mark.parametrize( + ("index", "error"), + [ + (0, TypeError), + ("asd", ValueError), + ((1, 2, 3), TypeError), + (("asd", "asd"), ValueError), + ], +) +def test_iat_errors(scalars_df_index, scalars_pandas_df_index, index, error): + with pytest.raises(error): + scalars_pandas_df_index.iat[index] + with pytest.raises(error): + scalars_df_index.iat[index] + + def test_iloc_single_integer_out_of_bound_error( scalars_df_index, scalars_pandas_df_index ): diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 993df18c959..802425510aa 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -1953,6 +1953,20 @@ def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, ste ) +def test_iat(scalars_df_index, scalars_pandas_df_index): + bf_result = scalars_df_index["int64_too"].iat[3] + pd_result = scalars_pandas_df_index["int64_too"].iat[3] + + assert bf_result == pd_result + + +def test_iat_error(scalars_df_index, scalars_pandas_df_index): + with pytest.raises(ValueError): + scalars_pandas_df_index["int64_too"].iat["asd"] + with pytest.raises(ValueError): + scalars_df_index["int64_too"].iat["asd"] + + def test_series_add_prefix(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index["int64_too"].add_prefix("prefix_").to_pandas() diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 644e043e83d..80a5428b36d 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2097,3 +2097,13 @@ def fillna(self, value): DataFrame: Object with missing values filled """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def iloc(self): + """Purely integer-location based indexing for selection by position.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def iat(self): + """Access a single value for a row/column pair by integer position.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 970ab1d8b43..03729922d5a 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1823,3 +1823,13 @@ def map( Series: Same index as caller. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def iloc(self): + """Purely integer-location based indexing for selection by position.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def iat(self): + """Access a single value for a row/column pair by integer position.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 97b8bec1175499c74448a4fd46b4888c4b4c35c1 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 10 Oct 2023 15:10:26 -0700 Subject: [PATCH 68/80] feat: add level param to DataFrame.stack (#88) * feat: add level param to DataFrame.stack --- bigframes/core/block_transforms.py | 2 +- bigframes/core/blocks.py | 12 +++---- bigframes/dataframe.py | 49 ++++++++++++++++++++------- tests/system/small/test_dataframe.py | 4 ++- tests/system/small/test_multiindex.py | 38 ++++++++++++++------- 5 files changed, 72 insertions(+), 33 deletions(-) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 4c30d7631dc..904da7f3122 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -53,7 +53,7 @@ def equals(block1: blocks.Block, block2: blocks.Block) -> bool: joined_block = joined_block.select_columns(equality_ids).with_column_labels( list(range(len(equality_ids))) ) - stacked_block = joined_block.stack(dropna=False, sort=False) + stacked_block = joined_block.stack() result = stacked_block.get_stat(stacked_block.value_columns[0], agg_ops.all_op) return typing.cast(bool, result) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 852b1fe02ab..9b49645c71d 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -1284,20 +1284,20 @@ def pivot( return result_block.with_column_labels(column_index) - def stack(self, how="left", dropna=True, sort=True, levels: int = 1): + def stack(self, how="left", levels: int = 1): """Unpivot last column axis level into row axis""" + if levels == 0: + return self + # These are the values that will be turned into rows col_labels, row_labels = utils.split_index(self.column_labels, levels=levels) - if dropna: - row_labels = row_labels.drop_duplicates() - if sort: - row_labels = row_labels.sort_values() + row_labels = row_labels.drop_duplicates() row_label_tuples = utils.index_as_tuples(row_labels) if col_labels is not None: - result_index = col_labels.drop_duplicates().sort_values().dropna(how="all") + result_index = col_labels.drop_duplicates().dropna(how="all") result_col_labels = utils.index_as_tuples(result_index) else: result_index = pd.Index([None]) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index d5b2fa86e9d..3ec1b4b6175 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1741,24 +1741,49 @@ def pivot( ) return DataFrame(pivot_block) - def stack(self): - # TODO: support 'level' param by simply reordering levels such that selected level is last before passing to Block.stack. - # TODO: match impl to pandas future_stack as described in pandas 2.1 release notes - stack_block = self._block.stack() - result_block = block_ops.dropna( - stack_block, stack_block.value_columns, how="all" - ) + def stack(self, level: LevelsType = -1): if not isinstance(self.columns, pandas.MultiIndex): - return bigframes.series.Series(result_block) - return DataFrame(result_block) + if level not in [0, -1, self.columns.name]: + raise IndexError(f"Invalid level {level} for single-level index") + return self._stack_mono() + return self._stack_multi(level) + + def _stack_mono(self): + result_block = self._block.stack() + return bigframes.series.Series(result_block) + + def _stack_multi(self, level: LevelsType = -1): + n_levels = self.columns.nlevels + if isinstance(level, int) or isinstance(level, str): + level = [level] + level_indices = [] + for level_ref in level: + if isinstance(level_ref, int): + if level_ref < 0: + level_indices.append(n_levels + level_ref) + else: + level_indices.append(level_ref) + else: # str + level_indices.append(self.columns.names.index(level_ref)) + + new_order = [ + *[i for i in range(n_levels) if i not in level_indices], + *level_indices, + ] + + original_columns = typing.cast(pandas.MultiIndex, self.columns) + new_columns = original_columns.reorder_levels(new_order) + + block = self._block.with_column_labels(new_columns) + + block = block.stack(levels=len(level)) + return DataFrame(block) def unstack(self): block = self._block # Special case, unstack with mono-index transpose into a series if self.index.nlevels == 1: - block = block.stack( - how="right", dropna=False, sort=False, levels=self.columns.nlevels - ) + block = block.stack(how="right", levels=self.columns.nlevels) return bigframes.series.Series(block) # Pivot by last level of index diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 0262a31a92e..9f1092d09d1 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1885,6 +1885,8 @@ def test_df_describe(scalars_dfs): def test_df_stack(scalars_dfs): + if pandas.__version__.startswith("1.") or pandas.__version__.startswith("2.0"): + pytest.skip("pandas <2.1 uses different stack implementation") scalars_df, scalars_pandas_df = scalars_dfs # To match bigquery dataframes scalars_pandas_df = scalars_pandas_df.copy() @@ -1893,7 +1895,7 @@ def test_df_stack(scalars_dfs): columns = ["int64_col", "int64_too", "rowindex_2"] bf_result = scalars_df[columns].stack().to_pandas() - pd_result = scalars_pandas_df[columns].stack() + pd_result = scalars_pandas_df[columns].stack(future_stack=True) # Pandas produces NaN, where bq dataframes produces pd.NA pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 19f1c557ef2..a1326767700 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -718,25 +718,37 @@ def test_column_multi_index_cumsum(scalars_df_index, scalars_pandas_df_index): pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) -def test_column_multi_index_stack(scalars_df_index, scalars_pandas_df_index): - columns = ["int64_too", "int64_col", "rowindex_2"] +@pytest.mark.parametrize( + ("level",), + [(["l3", "l1"],), ([-2, -1],), (["l3"],), ("l2",), (-3,)], +) +def test_column_multi_index_stack(level): + if pandas.__version__.startswith("1.") or pandas.__version__.startswith("2.0"): + pytest.skip("pandas <2.1 uses different stack implementation") + level1 = pandas.Index(["b", "a", "b"]) - # Need resulting column to be pyarrow string rather than object dtype - level2 = pandas.Index(["a", "b", "b"], dtype="string[pyarrow]") - multi_columns = pandas.MultiIndex.from_arrays([level1, level2]) - bf_df = scalars_df_index[columns].copy() - bf_df.columns = multi_columns - pd_df = scalars_pandas_df_index[columns].copy() - pd_df.columns = multi_columns + level2 = pandas.Index(["a", "b", "b"]) + level3 = pandas.Index(["b", "b", "a"]) - bf_result = bf_df.stack().to_pandas() - # Shifting sort behavior in stack - pd_result = pd_df.stack() + multi_columns = pandas.MultiIndex.from_arrays( + [level1, level2, level3], names=["l1", "l2", "l3"] + ) + pd_df = pandas.DataFrame( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + index=[5, 2, None], + columns=multi_columns, + dtype="Int64", + ) + bf_df = bpd.DataFrame(pd_df) + + bf_result = bf_df.stack(level=level).to_pandas() + # BigFrames emulates future_stack impl + pd_result = pd_df.stack(level=level, future_stack=True) # Pandas produces NaN, where bq dataframes produces pd.NA # Column ordering seems to depend on pandas version pandas.testing.assert_frame_equal( - bf_result.sort_index(axis=1), pd_result.sort_index(axis=1), check_dtype=False + bf_result, pd_result, check_dtype=False, check_index_type=False ) From 4ae4ef995348b95521c4988a8cfb3b5ac792fd69 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Wed, 11 Oct 2023 09:49:13 -0700 Subject: [PATCH 69/80] feat: use default session connection (#87) --- bigframes/_config/bigquery_options.py | 4 +- bigframes/clients.py | 24 +++++++ bigframes/ml/llm.py | 23 +++++-- bigframes/remote_function.py | 74 +++++++++------------- bigframes/session.py | 6 +- tests/system/small/ml/test_llm.py | 37 ++++++++++- tests/system/small/test_remote_function.py | 31 +++++++++ tests/unit/test_clients.py | 57 +++++++++++++++++ 8 files changed, 205 insertions(+), 51 deletions(-) create mode 100644 tests/unit/test_clients.py diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index ea1864ed5f7..eb56de826a8 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -83,12 +83,14 @@ def project(self, value: Optional[str]): @property def bq_connection(self) -> Optional[str]: - """Name of the BigQuery connection to use. + """Name of the BigQuery connection to use. Should be of the form ... You should either have the connection already created in the location you have chosen, or you should have the Project IAM Admin role to enable the service to create the connection for you if you need it. + + If this option isn't provided, or project or location aren't provided, session will use its default project/location/connection_id as default connection. """ return self._bq_connection diff --git a/bigframes/clients.py b/bigframes/clients.py index b60fcba04af..dcac611e8cf 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -29,6 +29,8 @@ ) logger = logging.getLogger(__name__) +_BIGFRAMES_DEFAULT_CONNECTION_ID = "bigframes-default-connection" + class BqConnectionManager: """Manager to handle operations with BQ connections.""" @@ -162,3 +164,25 @@ def _get_service_account_if_connection_exists( pass return service_account + + +def get_connection_name_full( + connection_name: Optional[str], default_project: str, default_location: str +) -> str: + """Retrieve the full connection name of the form ... + Use default project, location or connection_id when any of them are missing.""" + if connection_name is None: + return ( + f"{default_project}.{default_location}.{_BIGFRAMES_DEFAULT_CONNECTION_ID}" + ) + + if connection_name.count(".") == 2: + return connection_name + + if connection_name.count(".") == 1: + return f"{default_project}.{connection_name}" + + if connection_name.count(".") == 0: + return f"{default_project}.{default_location}.{connection_name}" + + raise ValueError(f"Invalid connection name format: {connection_name}.") diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index c86e5fb3b67..a61dd34e6de 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -38,8 +38,9 @@ class PaLM2TextGenerator(base.Predictor): session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): - connection to connect with remote service. str of the format ... - if None, use default connection in session context. + connection to connect with remote service. str of the format ... + if None, use default connection in session context. BigQuery DataFrame will try to create the connection and attach + permission if the connection isn't fully setup. """ def __init__( @@ -48,7 +49,14 @@ def __init__( connection_name: Optional[str] = None, ): self.session = session or bpd.get_global_session() - self.connection_name = connection_name or self.session._bq_connection + + connection_name = connection_name or self.session._bq_connection + self.connection_name = clients.get_connection_name_full( + connection_name, + default_project=self.session._project, + default_location=self.session._location, + ) + self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient ) @@ -180,7 +188,14 @@ def __init__( connection_name: Optional[str] = None, ): self.session = session or bpd.get_global_session() - self.connection_name = connection_name or self.session._bq_connection + + connection_name = connection_name or self.session._bq_connection + self.connection_name = clients.get_connection_name_full( + connection_name, + default_project=self.session._project, + default_location=self.session._location, + ) + self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient ) diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 6fc2f8e59fb..37c7a2fc64b 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -695,9 +695,12 @@ def remote_function( persistent name. """ + import bigframes.pandas as bpd + + session = session or bpd.get_global_session() # A BigQuery client is required to perform BQ operations - if not bigquery_client and session: + if not bigquery_client: bigquery_client = session.bqclient if not bigquery_client: raise ValueError( @@ -706,7 +709,7 @@ def remote_function( ) # A BigQuery connection client is required to perform BQ connection operations - if not bigquery_connection_client and session: + if not bigquery_connection_client: bigquery_connection_client = session.bqconnectionclient if not bigquery_connection_client: raise ValueError( @@ -716,8 +719,7 @@ def remote_function( # A cloud functions client is required to perform cloud functions operations if not cloud_functions_client: - if session: - cloud_functions_client = session.cloudfunctionsclient + cloud_functions_client = session.cloudfunctionsclient if not cloud_functions_client: raise ValueError( "A cloud functions client must be provided, either directly or via session. " @@ -726,8 +728,7 @@ def remote_function( # A resource manager client is required to get/set IAM operations if not resource_manager_client: - if session: - resource_manager_client = session.resourcemanagerclient + resource_manager_client = session.resourcemanagerclient if not resource_manager_client: raise ValueError( "A resource manager client must be provided, either directly or via session. " @@ -740,15 +741,10 @@ def remote_function( dataset_ref = bigquery.DatasetReference.from_string( dataset, default_project=bigquery_client.project ) - elif session: + else: dataset_ref = bigquery.DatasetReference.from_string( session._session_dataset_id, default_project=bigquery_client.project ) - else: - raise ValueError( - "Project and dataset must be provided, either directly or via session. " - f"{constants.FEEDBACK_LINK}" - ) bq_location, cloud_function_region = get_remote_function_locations( bigquery_client.location @@ -756,40 +752,30 @@ def remote_function( # A connection is required for BQ remote function # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function - if not bigquery_connection and session: - bigquery_connection = session._bq_connection # type: ignore if not bigquery_connection: + bigquery_connection = session._bq_connection # type: ignore + + bigquery_connection = clients.get_connection_name_full( + bigquery_connection, + default_project=dataset_ref.project, + default_location=bq_location, + ) + # Guaranteed to be the form of .. + ( + gcp_project_id, + bq_connection_location, + bq_connection_id, + ) = bigquery_connection.split(".") + if gcp_project_id.casefold() != dataset_ref.project.casefold(): raise ValueError( - "BigQuery connection must be provided, either directly or via session. " - f"{constants.FEEDBACK_LINK}" + "The project_id does not match BigQuery connection gcp_project_id: " + f"{dataset_ref.project}." + ) + if bq_connection_location.casefold() != bq_location.casefold(): + raise ValueError( + "The location does not match BigQuery connection location: " + f"{bq_location}." ) - - # Check connection_id with `LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID` format. - if bigquery_connection.count(".") == 1: - bq_connection_location, bq_connection_id = bigquery_connection.split(".") - if bq_connection_location.casefold() != bq_location.casefold(): - raise ValueError( - "The location does not match BigQuery connection location: " - f"{bq_location}." - ) - bigquery_connection = bq_connection_id - elif bigquery_connection.count(".") == 2: - ( - gcp_project_id, - bq_connection_location, - bq_connection_id, - ) = bigquery_connection.split(".") - if gcp_project_id.casefold() != dataset_ref.project.casefold(): - raise ValueError( - "The project_id does not match BigQuery connection gcp_project_id: " - f"{dataset_ref.project}." - ) - if bq_connection_location.casefold() != bq_location.casefold(): - raise ValueError( - "The location does not match BigQuery connection location: " - f"{bq_location}." - ) - bigquery_connection = bq_connection_id def wrapper(f): if not callable(f): @@ -808,7 +794,7 @@ def wrapper(f): dataset_ref.dataset_id, bigquery_client, bigquery_connection_client, - bigquery_connection, + bq_connection_id, resource_manager_client, ) diff --git a/bigframes/session.py b/bigframes/session.py index ac48c977cb0..a7cb78e3ffb 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -350,10 +350,14 @@ def resourcemanagerclient(self): @property def _session_dataset_id(self): """A dataset for storing temporary objects local to the session - This is a workaround for BQML models and remote functions that do not + This is a workaround for remote functions that do not yet support session-temporary instances.""" return self._session_dataset.dataset_id + @property + def _project(self): + return self.bqclient.project + def _create_and_bind_bq_session(self): """Create a BQ session and bind the session id with clients to capture BQ activities: go/bigframes-transient-data""" diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 7486277487d..e546c09f97a 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -26,7 +26,8 @@ def test_create_text_generator_model(palm2_text_generator_model): assert palm2_text_generator_model._bqml_model is not None -def test_create_text_generator_model_defaults(bq_connection): +@pytest.mark.flaky(retries=2, delay=120) +def test_create_text_generator_model_default_session(bq_connection, llm_text_pandas_df): import bigframes.pandas as bpd bpd.reset_session() @@ -36,6 +37,40 @@ def test_create_text_generator_model_defaults(bq_connection): model = llm.PaLM2TextGenerator() assert model is not None assert model._bqml_model is not None + assert model.connection_name.casefold() == "bigframes-dev.us.bigframes-rf-conn" + + llm_text_df = bpd.read_pandas(llm_text_pandas_df) + + df = model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_create_text_generator_model_default_connection(llm_text_pandas_df): + from bigframes import _config + import bigframes.pandas as bpd + + bpd.reset_session() + _config.options = _config.Options() # reset configs + + llm_text_df = bpd.read_pandas(llm_text_pandas_df) + + model = llm.PaLM2TextGenerator() + assert model is not None + assert model._bqml_model is not None + assert ( + model.connection_name.casefold() + == "bigframes-dev.us.bigframes-default-connection" + ) + + df = model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) # Marked as flaky only because BQML LLM is in preview, the service only has limited capacity, not stable enough. diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 77fb81d2c93..d024a57dedc 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -20,6 +20,7 @@ import bigframes from bigframes import remote_function as rf +import bigframes.pandas as bpd from tests.system.utils import assert_pandas_df_equal_ignore_ordering @@ -465,6 +466,36 @@ def square(x): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_default_connection(scalars_dfs, dataset_id): + @bpd.remote_function([int], int, dataset=dataset_id) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pd.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + + @pytest.mark.flaky(retries=2, delay=120) def test_dataframe_applymap(session_with_bq_connection, scalars_dfs): def add_one(x): diff --git a/tests/unit/test_clients.py b/tests/unit/test_clients.py new file mode 100644 index 00000000000..a90e5b0320b --- /dev/null +++ b/tests/unit/test_clients.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. + +import pytest + +from bigframes import clients + + +def test_get_connection_name_full_none(): + connection_name = clients.get_connection_name_full( + None, default_project="default-project", default_location="us" + ) + assert connection_name == "default-project.us.bigframes-default-connection" + + +def test_get_connection_name_full_connection_id(): + connection_name = clients.get_connection_name_full( + "connection-id", default_project="default-project", default_location="us" + ) + assert connection_name == "default-project.us.connection-id" + + +def test_get_connection_name_full_location_connection_id(): + connection_name = clients.get_connection_name_full( + "eu.connection-id", default_project="default-project", default_location="us" + ) + assert connection_name == "default-project.eu.connection-id" + + +def test_get_connection_name_full_all(): + connection_name = clients.get_connection_name_full( + "my-project.eu.connection-id", + default_project="default-project", + default_location="us", + ) + assert connection_name == "my-project.eu.connection-id" + + +def test_get_connection_name_full_raise_value_error(): + + with pytest.raises(ValueError): + clients.get_connection_name_full( + "my-project.eu.connection-id.extra_field", + default_project="default-project", + default_location="us", + ) From d39134db81249e19b266a465ac759bef1dfe229d Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 11 Oct 2023 10:58:57 -0700 Subject: [PATCH 70/80] chore(main): release 0.7.0 (#86) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 22 ++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 880f7916250..8d405d06ee7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,28 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.7.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.6.0...v0.7.0) (2023-10-11) + + +### Features + +* Add aliases for several series properties ([#80](https://github.com/googleapis/python-bigquery-dataframes/issues/80)) ([c0efec8](https://github.com/googleapis/python-bigquery-dataframes/commit/c0efec8956198247b27904345a795f09c80d3502)) +* Add equals methods to series/dataframe ([#76](https://github.com/googleapis/python-bigquery-dataframes/issues/76)) ([636a209](https://github.com/googleapis/python-bigquery-dataframes/commit/636a209e0853501abd50784a11a87cf7f2282ee5)) +* Add iat and iloc accessing by tuples of integers ([#90](https://github.com/googleapis/python-bigquery-dataframes/issues/90)) ([228aeba](https://github.com/googleapis/python-bigquery-dataframes/commit/228aeba09782ae2421040c7601c15d4af92790b6)) +* Add level param to DataFrame.stack ([#88](https://github.com/googleapis/python-bigquery-dataframes/issues/88)) ([97b8bec](https://github.com/googleapis/python-bigquery-dataframes/commit/97b8bec1175499c74448a4fd46b4888c4b4c35c1)) +* Allow df.drop to take an index object ([#68](https://github.com/googleapis/python-bigquery-dataframes/issues/68)) ([740c451](https://github.com/googleapis/python-bigquery-dataframes/commit/740c45176f79d4d2f7f28cb5f6c9eeb1327c8397)) +* Use default session connection ([#87](https://github.com/googleapis/python-bigquery-dataframes/issues/87)) ([4ae4ef9](https://github.com/googleapis/python-bigquery-dataframes/commit/4ae4ef995348b95521c4988a8cfb3b5ac792fd69)) + + +### Bug Fixes + +* Change the invalid url in docs ([#93](https://github.com/googleapis/python-bigquery-dataframes/issues/93)) ([969800d](https://github.com/googleapis/python-bigquery-dataframes/commit/969800d669204de4d0f2e5e61da521217e55668b)) + + +### Documentation + +* Add more preprocessing models into the docs menu. ([#97](https://github.com/googleapis/python-bigquery-dataframes/issues/97)) ([1592315](https://github.com/googleapis/python-bigquery-dataframes/commit/159231505f339173560cd802dae3fed3e63a663b)) + ## [0.6.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.5.0...v0.6.0) (2023-10-04) diff --git a/bigframes/version.py b/bigframes/version.py index 238b64473a0..e91e27ff76d 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.6.0" +__version__ = "0.7.0" From 1d385be1c68342a66ecb9f28c5efc83c18d0e64c Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 11 Oct 2023 20:04:14 +0000 Subject: [PATCH 71/80] fix: create session dataset for remote functions only when needed (#94) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With this change BigFrames will not create a dataset upfront at the time of session creation, but instead leave it to the components which need the dataset to create it. 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://togithub.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 # 🦕 --- bigframes/remote_function.py | 39 ++++++++++++++------ bigframes/session.py | 7 +--- tests/system/conftest.py | 23 +++++++++--- tests/system/large/test_remote_function.py | 43 ++++++++++++++++++++++ tests/system/small/test_session.py | 5 --- 5 files changed, 90 insertions(+), 27 deletions(-) diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 37c7a2fc64b..81ba26600ba 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -202,10 +202,22 @@ def create_bq_remote_function( OPTIONS ( endpoint = "{endpoint}" )""" + logger.info(f"Creating BQ remote function: {create_function_ddl}") + + # Make sure the dataset exists + dataset = bigquery.Dataset( + bigquery.DatasetReference.from_string( + self._bq_dataset, default_project=self._gcp_project_id + ) + ) + dataset.location = self._bq_location + self._bq_client.create_dataset(dataset, exists_ok=True) + # TODO: Use session._start_query() so we get progress bar query_job = self._bq_client.query(create_function_ddl) # Make an API request. query_job.result() # Wait for the job to complete. + logger.info(f"Created remote function {query_job.ddl_target_routine}") def get_cloud_function_fully_qualified_parent(self): @@ -465,17 +477,22 @@ def get_remote_function_specs(self, remote_function_name): routines = self._bq_client.list_routines( f"{self._gcp_project_id}.{self._bq_dataset}" ) - for routine in routines: - if routine.reference.routine_id == remote_function_name: - # TODO(shobs): Use first class properties when they are available - # https://github.com/googleapis/python-bigquery/issues/1552 - rf_options = routine._properties.get("remoteFunctionOptions") - if rf_options: - http_endpoint = rf_options.get("endpoint") - bq_connection = rf_options.get("connection") - if bq_connection: - bq_connection = os.path.basename(bq_connection) - break + try: + for routine in routines: + if routine.reference.routine_id == remote_function_name: + # TODO(shobs): Use first class properties when they are available + # https://github.com/googleapis/python-bigquery/issues/1552 + rf_options = routine._properties.get("remoteFunctionOptions") + if rf_options: + http_endpoint = rf_options.get("endpoint") + bq_connection = rf_options.get("connection") + if bq_connection: + bq_connection = os.path.basename(bq_connection) + break + except google.api_core.exceptions.NotFound: + # The dataset might not exist, in which case the http_endpoint doesn't, either. + # Note: list_routines doesn't make an API request until we iterate on the response object. + pass return (http_endpoint, bq_connection) diff --git a/bigframes/session.py b/bigframes/session.py index a7cb78e3ffb..6ad65000ce3 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -381,17 +381,12 @@ def _create_and_bind_bq_session(self): ] ) - # Dataset for storing BQML models and remote functions, which don't yet + # Dataset for storing remote functions, which don't yet # support proper session temporary storage yet self._session_dataset = bigquery.Dataset( f"{self.bqclient.project}.bigframes_temp_{self._location.lower().replace('-', '_')}" ) self._session_dataset.location = self._location - self._session_dataset.default_table_expiration_ms = 24 * 60 * 60 * 1000 - - # TODO: handle case when the dataset does not exist and the user does - # not have permission to create one (bigquery.datasets.create IAM) - self.bqclient.create_dataset(self._session_dataset, exists_ok=True) def close(self): """Terminated the BQ session, otherwises the session will be terminated automatically after diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 3153bd1559e..ed22a3e8dac 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -134,15 +134,28 @@ def cleanup_datasets(bigquery_client: bigquery.Client) -> None: ) +def get_dataset_id(project_id: str): + "Get a fully qualified dataset id belonging to the given project." + dataset_id = f"{project_id}.{prefixer.create_prefix()}_dataset_id" + return dataset_id + + @pytest.fixture(scope="session") def dataset_id(bigquery_client: bigquery.Client): """Create (and cleanup) a temporary dataset.""" - project_id = bigquery_client.project - dataset_id = f"{project_id}.{prefixer.create_prefix()}_dataset_id" - dataset = bigquery.Dataset(dataset_id) - bigquery_client.create_dataset(dataset) + dataset_id = get_dataset_id(bigquery_client.project) + bigquery_client.create_dataset(dataset_id) + yield dataset_id + bigquery_client.delete_dataset(dataset_id, delete_contents=True) + + +@pytest.fixture +def dataset_id_not_created(bigquery_client: bigquery.Client): + """Return a temporary dataset object without creating it, and clean it up + after it has been used.""" + dataset_id = get_dataset_id(bigquery_client.project) yield dataset_id - bigquery_client.delete_dataset(dataset, delete_contents=True) + bigquery_client.delete_dataset(dataset_id, delete_contents=True) @pytest.fixture(scope="session") diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index 2f231f40c95..f2700991828 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -408,6 +408,49 @@ def add_one(x): ) +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_explicit_dataset_not_created( + session, scalars_dfs, dataset_id_not_created, bq_cf_connection, functions_client +): + try: + + @session.remote_function( + [int], + int, + dataset_id_not_created, + bq_cf_connection, + reuse=False, + ) + def square(x): + return x * x + + scalars_df, scalars_pandas_df = scalars_dfs + + bf_int64_col = scalars_df["int64_col"] + bf_int64_col_filter = bf_int64_col.notnull() + bf_int64_col_filtered = bf_int64_col[bf_int64_col_filter] + bf_result_col = bf_int64_col_filtered.apply(square) + bf_result = ( + bf_int64_col_filtered.to_frame().assign(result=bf_result_col).to_pandas() + ) + + pd_int64_col = scalars_pandas_df["int64_col"] + pd_int64_col_filter = pd_int64_col.notnull() + pd_int64_col_filtered = pd_int64_col[pd_int64_col_filter] + pd_result_col = pd_int64_col_filtered.apply(lambda x: x * x) + # TODO(shobs): Figure why pandas .apply() changes the dtype, i.e. + # pd_int64_col_filtered.dtype is Int64Dtype() + # pd_int64_col_filtered.apply(lambda x: x * x).dtype is int64. + # For this test let's force the pandas dtype to be same as bigframes' dtype. + pd_result_col = pd_result_col.astype(pandas.Int64Dtype()) + pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col) + + assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets(session.bqclient, functions_client, square) + + @pytest.mark.flaky(retries=2, delay=120) def test_remote_udf_referring_outside_var( session, scalars_dfs, dataset_id, bq_cf_connection, functions_client diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 53ddfa3c495..7655325bfcc 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -894,11 +894,6 @@ def test_session_id(session): # TODO(chelsealin): Verify the session id can be binded with a load job. -def test_session_dataset_exists_and_configured(session: bigframes.Session): - dataset = session.bqclient.get_dataset(session._session_dataset_id) - assert dataset.default_table_expiration_ms == 24 * 60 * 60 * 1000 - - @pytest.mark.flaky(retries=2) def test_to_close_session(): session = bigframes.Session() From a8c286f0995cc8cf2a4c44fb51855773ecf71f72 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 12 Oct 2023 07:13:04 +0000 Subject: [PATCH 72/80] feat: Support compression in `to_parquet` (#91) * feat: Support compression in to_parquet This changes the default behavior from no compression to snappy compression. * feat: Support compression in to_parquet BREAKING CHANGE: The default behavior of to_parquet is changing from no compression to 'snappy' compression. * fix exception message, add tests for not supported compression techniques --- bigframes/dataframe.py | 18 +++- tests/system/small/test_session.py | 85 ++++++++++++++++++- .../bigframes_vendored/pandas/core/frame.py | 5 ++ 3 files changed, 105 insertions(+), 3 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 3ec1b4b6175..015a7642f8b 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -21,6 +21,7 @@ import typing from typing import ( Callable, + Dict, Iterable, List, Literal, @@ -2270,7 +2271,13 @@ def to_numpy( __array__ = to_numpy - def to_parquet(self, path: str, *, index: bool = True) -> None: + def to_parquet( + self, + path: str, + *, + compression: Optional[Literal["snappy", "gzip"]] = "snappy", + index: bool = True, + ) -> None: # TODO(swast): Can we support partition columns argument? # TODO(chelsealin): Support local file paths. # TODO(swast): Some warning that wildcard is recommended for large @@ -2282,6 +2289,13 @@ def to_parquet(self, path: str, *, index: bool = True) -> None: if "*" not in path: raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) + if compression not in {None, "snappy", "gzip"}: + raise ValueError("'{0}' is not valid for compression".format(compression)) + + export_options: Dict[str, Union[bool, str]] = {} + if compression: + export_options["compression"] = compression.upper() + result_table = self._run_io_query( index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID ) @@ -2289,7 +2303,7 @@ def to_parquet(self, path: str, *, index: bool = True) -> None: f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path, format="PARQUET", - export_options={}, + export_options=export_options, ) _, query_job = self._block.expr._session._start_query(export_data_statement) self._set_internal_query_job(query_job) diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 7655325bfcc..bfe9bc8d0f9 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -793,7 +793,7 @@ def test_read_parquet_gcs(session: bigframes.Session, scalars_dfs, gcs_folder): scalars_df, _ = scalars_dfs # Include wildcard so that multiple files can be written/read if > 1 GB. # https://cloud.google.com/bigquery/docs/exporting-data#exporting_data_into_one_or_more_files - path = gcs_folder + "test_read_parquet_gcs*.parquet" + path = gcs_folder + test_read_parquet_gcs.__name__ + "*.parquet" df_in: bigframes.dataframe.DataFrame = scalars_df.copy() # GEOGRAPHY not supported in parquet export. df_in = df_in.drop(columns="geography_col") @@ -823,6 +823,89 @@ def test_read_parquet_gcs(session: bigframes.Session, scalars_dfs, gcs_folder): pd.testing.assert_frame_equal(pd_df_in, pd_df_out) +@pytest.mark.parametrize( + "compression", + [ + None, + "gzip", + "snappy", + ], +) +def test_read_parquet_gcs_compressed( + session: bigframes.Session, scalars_dfs, gcs_folder, compression +): + scalars_df, _ = scalars_dfs + # Include wildcard so that multiple files can be written/read if > 1 GB. + # https://cloud.google.com/bigquery/docs/exporting-data#exporting_data_into_one_or_more_files + path = ( + gcs_folder + + test_read_parquet_gcs_compressed.__name__ + + (f"_{compression}" if compression else "") + + "*.parquet" + ) + df_in: bigframes.dataframe.DataFrame = scalars_df.copy() + # GEOGRAPHY not supported in parquet export. + df_in = df_in.drop(columns="geography_col") + # Make sure we can also serialize the order. + df_write = df_in.reset_index(drop=False) + df_write.index.name = f"ordering_id_{random.randrange(1_000_000)}" + df_write.to_parquet(path, compression=compression, index=True) + + df_out = ( + session.read_parquet(path) + # Restore order. + .set_index(df_write.index.name).sort_index() + # Restore index. + .set_index(typing.cast(str, df_in.index.name)) + ) + + # DATETIME gets loaded as TIMESTAMP in parquet. See: + # https://cloud.google.com/bigquery/docs/exporting-data#parquet_export_details + df_out = df_out.assign( + datetime_col=df_out["datetime_col"].astype("timestamp[us][pyarrow]") + ) + + # Make sure we actually have at least some values before comparing. + assert df_out.size != 0 + pd_df_in = df_in.to_pandas() + pd_df_out = df_out.to_pandas() + pd.testing.assert_frame_equal(pd_df_in, pd_df_out) + + +@pytest.mark.parametrize( + "compression", + [ + "brotli", + "lz4", + "zstd", + "unknown", + ], +) +def test_read_parquet_gcs_compression_not_supported( + session: bigframes.Session, scalars_dfs, gcs_folder, compression +): + scalars_df, _ = scalars_dfs + # Include wildcard so that multiple files can be written/read if > 1 GB. + # https://cloud.google.com/bigquery/docs/exporting-data#exporting_data_into_one_or_more_files + path = ( + gcs_folder + + test_read_parquet_gcs_compression_not_supported.__name__ + + (f"_{compression}" if compression else "") + + "*.parquet" + ) + df_in: bigframes.dataframe.DataFrame = scalars_df.copy() + # GEOGRAPHY not supported in parquet export. + df_in = df_in.drop(columns="geography_col") + # Make sure we can also serialize the order. + df_write = df_in.reset_index(drop=False) + df_write.index.name = f"ordering_id_{random.randrange(1_000_000)}" + + with pytest.raises( + ValueError, match=f"'{compression}' is not valid for compression" + ): + df_write.to_parquet(path, compression=compression, index=True) + + def test_read_json_gcs_bq_engine(session, scalars_dfs, gcs_folder): scalars_df, _ = scalars_dfs path = gcs_folder + "test_read_json_gcs_bq_engine_w_index*.json" diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 80a5428b36d..e54f984d59b 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -129,6 +129,7 @@ def to_parquet( self, path: str, *, + compression: Optional[Literal["snappy", "gzip"]] = "snappy", index: bool = True, ) -> None: """Write a DataFrame to the binary Parquet format. @@ -143,6 +144,10 @@ def to_parquet( If the data size is more than 1GB, you must use a wildcard to export the data into multiple files and the size of the files varies. + compression (str, default 'snappy'): + Name of the compression to use. Use ``None`` for no compression. + Supported options: ``'gzip'``, ``'snappy'``. + index (bool, default True): If ``True``, include the dataframe's index(es) in the file output. If ``False``, they will not be written to the file. From 1ff2755bfda121977ebde153450dc0c5a0042281 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 17:40:13 +0000 Subject: [PATCH 73/80] chore(main): release 0.8.0 (#99) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit :robot: I have created a release *beep* *boop* --- ## [0.8.0](https://togithub.com/googleapis/python-bigquery-dataframes/compare/v0.7.0...v0.8.0) (2023-10-12) ### ⚠ BREAKING CHANGES * The default behavior of `to_parquet` is changing from no compression to `'snappy'` compression. ### Features * Support compression in `to_parquet` ([a8c286f](https://togithub.com/googleapis/python-bigquery-dataframes/commit/a8c286f0995cc8cf2a4c44fb51855773ecf71f72)) ### Bug Fixes * Create session dataset for remote functions only when needed ([#94](https://togithub.com/googleapis/python-bigquery-dataframes/issues/94)) ([1d385be](https://togithub.com/googleapis/python-bigquery-dataframes/commit/1d385be1c68342a66ecb9f28c5efc83c18d0e64c)) --- This PR was generated with [Release Please](https://togithub.com/googleapis/release-please). See [documentation](https://togithub.com/googleapis/release-please#release-please). --- CHANGELOG.md | 16 ++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d405d06ee7..e34b073999a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,22 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.8.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.7.0...v0.8.0) (2023-10-12) + + +### ⚠ BREAKING CHANGES + +* The default behavior of `to_parquet` is changing from no compression to `'snappy'` compression. + +### Features + +* Support compression in `to_parquet` ([a8c286f](https://github.com/googleapis/python-bigquery-dataframes/commit/a8c286f0995cc8cf2a4c44fb51855773ecf71f72)) + + +### Bug Fixes + +* Create session dataset for remote functions only when needed ([#94](https://github.com/googleapis/python-bigquery-dataframes/issues/94)) ([1d385be](https://github.com/googleapis/python-bigquery-dataframes/commit/1d385be1c68342a66ecb9f28c5efc83c18d0e64c)) + ## [0.7.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.6.0...v0.7.0) (2023-10-11) diff --git a/bigframes/version.py b/bigframes/version.py index e91e27ff76d..974fbf1ac9b 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.7.0" +__version__ = "0.8.0" From 3b51a363615e292cb5442eee18ea76a0a3999d2e Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Thu, 12 Oct 2023 11:56:15 -0700 Subject: [PATCH 74/80] chore: address last comments of PR#87 (#102) PR#87 https://togithub.com/googleapis/python-bigquery-dataframes/pull/87 --- bigframes/clients.py | 41 +++++++++++++++--------------------- bigframes/ml/llm.py | 16 +++++++------- bigframes/remote_function.py | 2 +- bigframes/session.py | 4 +++- tests/unit/test_clients.py | 16 ++++---------- 5 files changed, 33 insertions(+), 46 deletions(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index dcac611e8cf..4ba9d93d694 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -29,8 +29,6 @@ ) logger = logging.getLogger(__name__) -_BIGFRAMES_DEFAULT_CONNECTION_ID = "bigframes-default-connection" - class BqConnectionManager: """Manager to handle operations with BQ connections.""" @@ -46,6 +44,23 @@ def __init__( self._bq_connection_client = bq_connection_client self._cloud_resource_manager_client = cloud_resource_manager_client + @classmethod + def resolve_full_connection_name( + cls, connection_name: str, default_project: str, default_location: str + ) -> str: + """Retrieve the full connection name of the form ... + Use default project, location or connection_id when any of them are missing.""" + if connection_name.count(".") == 2: + return connection_name + + if connection_name.count(".") == 1: + return f"{default_project}.{connection_name}" + + if connection_name.count(".") == 0: + return f"{default_project}.{default_location}.{connection_name}" + + raise ValueError(f"Invalid connection name format: {connection_name}.") + def create_bq_connection( self, project_id: str, location: str, connection_id: str, iam_role: str ): @@ -164,25 +179,3 @@ def _get_service_account_if_connection_exists( pass return service_account - - -def get_connection_name_full( - connection_name: Optional[str], default_project: str, default_location: str -) -> str: - """Retrieve the full connection name of the form ... - Use default project, location or connection_id when any of them are missing.""" - if connection_name is None: - return ( - f"{default_project}.{default_location}.{_BIGFRAMES_DEFAULT_CONNECTION_ID}" - ) - - if connection_name.count(".") == 2: - return connection_name - - if connection_name.count(".") == 1: - return f"{default_project}.{connection_name}" - - if connection_name.count(".") == 0: - return f"{default_project}.{default_location}.{connection_name}" - - raise ValueError(f"Invalid connection name format: {connection_name}.") diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index a61dd34e6de..d78f467537c 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -49,17 +49,17 @@ def __init__( connection_name: Optional[str] = None, ): self.session = session or bpd.get_global_session() + self._bq_connection_manager = clients.BqConnectionManager( + self.session.bqconnectionclient, self.session.resourcemanagerclient + ) connection_name = connection_name or self.session._bq_connection - self.connection_name = clients.get_connection_name_full( + self.connection_name = self._bq_connection_manager.resolve_full_connection_name( connection_name, default_project=self.session._project, default_location=self.session._location, ) - self._bq_connection_manager = clients.BqConnectionManager( - self.session.bqconnectionclient, self.session.resourcemanagerclient - ) self._bqml_model_factory = globals.bqml_model_factory() self._bqml_model: core.BqmlModel = self._create_bqml_model() @@ -188,17 +188,17 @@ def __init__( connection_name: Optional[str] = None, ): self.session = session or bpd.get_global_session() + self._bq_connection_manager = clients.BqConnectionManager( + self.session.bqconnectionclient, self.session.resourcemanagerclient + ) connection_name = connection_name or self.session._bq_connection - self.connection_name = clients.get_connection_name_full( + self.connection_name = self._bq_connection_manager.resolve_full_connection_name( connection_name, default_project=self.session._project, default_location=self.session._location, ) - self._bq_connection_manager = clients.BqConnectionManager( - self.session.bqconnectionclient, self.session.resourcemanagerclient - ) self._bqml_model_factory = globals.bqml_model_factory() self._bqml_model: core.BqmlModel = self._create_bqml_model() diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index 81ba26600ba..fd9aec825fa 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -772,7 +772,7 @@ def remote_function( if not bigquery_connection: bigquery_connection = session._bq_connection # type: ignore - bigquery_connection = clients.get_connection_name_full( + bigquery_connection = clients.BqConnectionManager.resolve_full_connection_name( bigquery_connection, default_project=dataset_ref.project, default_location=bq_location, diff --git a/bigframes/session.py b/bigframes/session.py index 6ad65000ce3..4f509f0704a 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -97,6 +97,8 @@ _BIGQUERYCONNECTION_REGIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com" _BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com" +_BIGFRAMES_DEFAULT_CONNECTION_ID = "bigframes-default-connection" + _MAX_CLUSTER_COLUMNS = 4 # TODO(swast): Need to connect to regional endpoints when performing remote @@ -321,7 +323,7 @@ def __init__( ), ) - self._bq_connection = context.bq_connection + self._bq_connection = context.bq_connection or _BIGFRAMES_DEFAULT_CONNECTION_ID # Now that we're starting the session, don't allow the options to be # changed. diff --git a/tests/unit/test_clients.py b/tests/unit/test_clients.py index a90e5b0320b..f89cc21397c 100644 --- a/tests/unit/test_clients.py +++ b/tests/unit/test_clients.py @@ -17,29 +17,22 @@ from bigframes import clients -def test_get_connection_name_full_none(): - connection_name = clients.get_connection_name_full( - None, default_project="default-project", default_location="us" - ) - assert connection_name == "default-project.us.bigframes-default-connection" - - def test_get_connection_name_full_connection_id(): - connection_name = clients.get_connection_name_full( + connection_name = clients.BqConnectionManager.resolve_full_connection_name( "connection-id", default_project="default-project", default_location="us" ) assert connection_name == "default-project.us.connection-id" def test_get_connection_name_full_location_connection_id(): - connection_name = clients.get_connection_name_full( + connection_name = clients.BqConnectionManager.resolve_full_connection_name( "eu.connection-id", default_project="default-project", default_location="us" ) assert connection_name == "default-project.eu.connection-id" def test_get_connection_name_full_all(): - connection_name = clients.get_connection_name_full( + connection_name = clients.BqConnectionManager.resolve_full_connection_name( "my-project.eu.connection-id", default_project="default-project", default_location="us", @@ -48,9 +41,8 @@ def test_get_connection_name_full_all(): def test_get_connection_name_full_raise_value_error(): - with pytest.raises(ValueError): - clients.get_connection_name_full( + clients.BqConnectionManager.resolve_full_connection_name( "my-project.eu.connection-id.extra_field", default_project="default-project", default_location="us", From e325fbb1c91e040d87df10f7d4d5ce53f7c052cb Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Thu, 12 Oct 2023 14:15:29 -0700 Subject: [PATCH 75/80] feat: send BigQuery cancel request when canceling bigframes process (#103) Co-authored-by: Henry J Solberg --- bigframes/formatting_helpers.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py index 82e2510e2af..f3e3a76ce42 100644 --- a/bigframes/formatting_helpers.py +++ b/bigframes/formatting_helpers.py @@ -153,6 +153,11 @@ def wait_for_query_job( except api_core_exceptions.GoogleAPICallError as exc: add_feedback_link(exc) raise + except KeyboardInterrupt: + query_job.cancel() + print(f"Requested cancelation for {query_job.job_type} {query_job.job_id}...") + # begin the cancel request before immediately rethrowing + raise def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): @@ -190,6 +195,11 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): except api_core_exceptions.GoogleAPICallError as exc: add_feedback_link(exc) raise + except KeyboardInterrupt: + job.cancel() + print(f"Requested cancelation for {job.job_type} {job.job_id}...") + # begin the cancel request before immediately rethrowing + raise def get_job_url(query_job: GenericJob): From 36693bff398c23e179d9bde95d52cbaddaf85c45 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Thu, 12 Oct 2023 15:02:15 -0700 Subject: [PATCH 76/80] feat!: rename `bigframes.pandas.reset_session` to `close_session` (#101) --- README.rst | 4 ++-- bigframes/__init__.py | 4 ++-- bigframes/_config/bigquery_options.py | 2 +- bigframes/core/global_session.py | 2 +- bigframes/pandas/__init__.py | 4 ++-- .../bq_dataframes_llm_code_generation.ipynb | 2 +- .../getting_started_bq_dataframes.ipynb | 2 +- .../bq_dataframes_ml_linear_regression.ipynb | 2 +- samples/snippets/quickstart_test.py | 2 +- samples/snippets/remote_function_test.py | 2 +- tests/system/small/ml/test_llm.py | 6 +++--- tests/system/small/test_pandas_options.py | 12 ++++++------ tests/unit/test_pandas.py | 6 +++--- 13 files changed, 25 insertions(+), 25 deletions(-) diff --git a/README.rst b/README.rst index 77c42e43251..b2ad1ee0179 100644 --- a/README.rst +++ b/README.rst @@ -63,7 +63,7 @@ auto-populates ``bf.options.bigquery.location`` if the user starts with directly or in a SQL statement. If you want to reset the location of the created DataFrame or Series objects, -you can reset the session by executing ``bigframes.pandas.reset_session()``. +you can close the session by executing ``bigframes.pandas.close_session()``. After that, you can reuse ``bigframes.pandas.options.bigquery.location`` to specify another location. @@ -335,7 +335,7 @@ sessions ; when this happens, you can’t use previously created DataFrame or Series objects and must re-create them using a new BigQuery DataFrames session. You can do this by running -``bigframes.pandas.reset_session()`` and then re-running the BigQuery +``bigframes.pandas.close_session()`` and then re-running the BigQuery DataFrames expressions. diff --git a/bigframes/__init__.py b/bigframes/__init__.py index 3e54a6d0903..8f41790072b 100644 --- a/bigframes/__init__.py +++ b/bigframes/__init__.py @@ -16,7 +16,7 @@ from bigframes._config import options from bigframes._config.bigquery_options import BigQueryOptions -from bigframes.core.global_session import get_global_session, reset_session +from bigframes.core.global_session import close_session, get_global_session from bigframes.session import connect, Session from bigframes.version import __version__ @@ -24,7 +24,7 @@ "options", "BigQueryOptions", "get_global_session", - "reset_session", + "close_session", "connect", "Session", "__version__", diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index eb56de826a8..453e7f5dff4 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -23,7 +23,7 @@ SESSION_STARTED_MESSAGE = ( "Cannot change '{attribute}' once a session has started. " - "Call bigframes.pandas.reset_session() first, if you are using the bigframes.pandas API." + "Call bigframes.pandas.close_session() first, if you are using the bigframes.pandas API." ) diff --git a/bigframes/core/global_session.py b/bigframes/core/global_session.py index 68529981cda..1f960839a0a 100644 --- a/bigframes/core/global_session.py +++ b/bigframes/core/global_session.py @@ -24,7 +24,7 @@ _global_session_lock = threading.Lock() -def reset_session() -> None: +def close_session() -> None: """Start a fresh session the next time a function requires a session. Closes the current session if it was already started. diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index cb278345907..ded4760fa3b 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -447,7 +447,7 @@ def read_gbq_function(function_name: str): # Session management APIs get_global_session = global_session.get_global_session -reset_session = global_session.reset_session +close_session = global_session.close_session # Use __all__ to let type checkers know what is part of the public API. @@ -478,5 +478,5 @@ def read_gbq_function(function_name: str): "options", # Session management APIs "get_global_session", - "reset_session", + "close_session", ] diff --git a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index 2e4ce3e5109..0f113b84c6d 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -354,7 +354,7 @@ "id": "DTVtFlqeFbrU" }, "source": [ - "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.close_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." ] }, { diff --git a/notebooks/getting_started/getting_started_bq_dataframes.ipynb b/notebooks/getting_started/getting_started_bq_dataframes.ipynb index 6936e1cf59b..6cc6acc9935 100644 --- a/notebooks/getting_started/getting_started_bq_dataframes.ipynb +++ b/notebooks/getting_started/getting_started_bq_dataframes.ipynb @@ -383,7 +383,7 @@ "id": "pDfrKwMKE_dK" }, "source": [ - "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.close_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." ] }, { diff --git a/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb b/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb index 338d6edf4fb..675416f6ea8 100644 --- a/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb +++ b/notebooks/regression/bq_dataframes_ml_linear_regression.ipynb @@ -370,7 +370,7 @@ "id": "D21CoOlfFTYI" }, "source": [ - "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.reset_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." + "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bf.close_session()`. After that, you can reuse `bf.options.bigquery.location` to specify another location." ] }, { diff --git a/samples/snippets/quickstart_test.py b/samples/snippets/quickstart_test.py index 6b0c69de999..bbe4a8b3c49 100644 --- a/samples/snippets/quickstart_test.py +++ b/samples/snippets/quickstart_test.py @@ -23,7 +23,7 @@ def test_quickstart( capsys: pytest.CaptureFixture[str], ) -> None: # We need a fresh session since we're modifying connection options. - bigframes.pandas.reset_session() + bigframes.pandas.close_session() # TODO(swast): Get project from environment so contributors can run tests. quickstart.run_quickstart("bigframes-dev") diff --git a/samples/snippets/remote_function_test.py b/samples/snippets/remote_function_test.py index 8b51e46b45a..e1317c6ac08 100644 --- a/samples/snippets/remote_function_test.py +++ b/samples/snippets/remote_function_test.py @@ -23,7 +23,7 @@ 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.reset_session() + bigframes.pandas.close_session() # TODO(swast): Get project from environment so contributors can run tests. remote_function.run_remote_function_and_read_gbq_function("bigframes-dev") diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index e546c09f97a..b7257dde1be 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -30,7 +30,7 @@ def test_create_text_generator_model(palm2_text_generator_model): def test_create_text_generator_model_default_session(bq_connection, llm_text_pandas_df): import bigframes.pandas as bpd - bpd.reset_session() + bpd.close_session() bpd.options.bigquery.bq_connection = bq_connection bpd.options.bigquery.location = "us" @@ -53,7 +53,7 @@ def test_create_text_generator_model_default_connection(llm_text_pandas_df): from bigframes import _config import bigframes.pandas as bpd - bpd.reset_session() + bpd.close_session() _config.options = _config.Options() # reset configs llm_text_df = bpd.read_pandas(llm_text_pandas_df) @@ -130,7 +130,7 @@ def test_create_embedding_generator_model(palm2_embedding_generator_model): def test_create_text_embedding_generator_model_defaults(bq_connection): import bigframes.pandas as bpd - bpd.reset_session() + bpd.close_session() bpd.options.bigquery.bq_connection = bq_connection bpd.options.bigquery.location = "us" diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index 956b29ae12f..ca67710d4ea 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -26,7 +26,7 @@ @pytest.fixture(autouse=True) def reset_default_session_and_location(): - bpd.reset_session() + bpd.close_session() bpd.options.bigquery.location = None @@ -79,8 +79,8 @@ def test_read_gbq_start_sets_session_location( ): read_method(query) - # Reset global session to start over - bpd.reset_session() + # Close global session to start over + bpd.close_session() # There should still be the previous location set in the bigquery options assert bpd.options.bigquery.location == tokyo_location @@ -254,7 +254,7 @@ def test_read_gbq_must_comply_with_set_location_non_US( assert df is not None -def test_reset_session_after_credentials_need_reauthentication(monkeypatch): +def test_close_session_after_credentials_need_reauthentication(monkeypatch): # Use a simple test query to verify that default session works to interact # with BQ test_query = "SELECT 1" @@ -288,8 +288,8 @@ def test_reset_session_after_credentials_need_reauthentication(monkeypatch): with pytest.raises(google.auth.exceptions.RefreshError): bpd.read_gbq(test_query) - # Now verify that resetting the session works - bpd.reset_session() + # Now verify that closing the session works + bpd.close_session() assert bigframes.core.global_session._global_session is None # Now verify that use is able to start over diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index 2325fc96a0a..5d4f69c7c01 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -116,7 +116,7 @@ def test_pandas_attribute(): assert bpd.ArrowDtype is pd.ArrowDtype -def test_reset_session_after_bq_session_ended(monkeypatch): +def test_close_session_after_bq_session_ended(monkeypatch): bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) bqclient.project = "test-project" session = resources.create_bigquery_session( @@ -144,7 +144,7 @@ def test_reset_session_after_bq_session_ended(monkeypatch): bpd.read_gbq("SELECT 1") # Even though the query to stop the session raises an exception, we should - # still be able to reset it without raising an error to the user. - bpd.reset_session() + # still be able to close it without raising an error to the user. + bpd.close_session() assert "CALL BQ.ABORT_SESSION('JUST_A_TEST')" in bqclient.query.call_args.args[0] assert bigframes.core.global_session._global_session is None From 855616a1472f1f3daa69b2dc3f84d4398537ca88 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 12 Oct 2023 16:04:40 -0700 Subject: [PATCH 77/80] refactor: all ArrayValue ops return only ArrayValue (#92) * refactor: all ArrayValue ops return only ArrayValue * copyright notice --------- Co-authored-by: Tim Swast --- bigframes/core/__init__.py | 15 +- bigframes/core/block_transforms.py | 8 +- bigframes/core/blocks.py | 125 +++++++++------- bigframes/core/indexes/index.py | 101 +++++++++---- bigframes/core/joins/name_resolution.py | 46 ++++++ bigframes/core/joins/row_identity.py | 44 +++--- bigframes/core/joins/single_column.py | 182 +++++------------------- bigframes/dataframe.py | 31 ++-- bigframes/operations/__init__.py | 10 ++ bigframes/operations/base.py | 4 +- bigframes/series.py | 6 +- 11 files changed, 290 insertions(+), 282 deletions(-) create mode 100644 bigframes/core/joins/name_resolution.py diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index ccfd682215c..6c78a07f3ba 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -211,8 +211,8 @@ def column_ids(self) -> typing.Sequence[str]: return tuple(self._column_names.keys()) @property - def hidden_ordering_columns(self) -> typing.Tuple[ibis_types.Value, ...]: - return self._hidden_ordering_columns + def _hidden_column_ids(self) -> typing.Sequence[str]: + return tuple(self._hidden_ordering_column_names.keys()) @property def _reduced_predicate(self) -> typing.Optional[ibis_types.BooleanValue]: @@ -400,7 +400,7 @@ def _hide_column(self, column_id) -> ArrayValue: expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) return expr_builder.build() - def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: + def promote_offsets(self, col_id: str) -> ArrayValue: """ Convenience function to promote copy of column offsets to a value column. Can be used to reset index. """ @@ -408,16 +408,15 @@ def promote_offsets(self) -> typing.Tuple[ArrayValue, str]: ordering = self._ordering if (not ordering.is_sequential) or (not ordering.total_order_col): - return self._project_offsets().promote_offsets() - col_id = bigframes.core.guid.generate_guid() + return self._project_offsets().promote_offsets(col_id) expr_builder = self.builder() expr_builder.columns = [ self._get_any_column(ordering.total_order_col.column_id).name(col_id), *self.columns, ] - return expr_builder.build(), col_id + return expr_builder.build() - def select_columns(self, column_ids: typing.Sequence[str]): + def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue: return self._projection( [self._get_ibis_column(col_id) for col_id in column_ids] ) @@ -807,7 +806,7 @@ def _create_order_columns( elif ordering_mode == "string_encoded": return (self._create_string_ordering_column().name(order_col_name),) elif expose_hidden_cols: - return self.hidden_ordering_columns + return self._hidden_ordering_columns return () def _create_offset_column(self) -> ibis_types.IntegerColumn: diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 904da7f3122..b0f05f47984 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -40,8 +40,8 @@ def equals(block1: blocks.Block, block2: blocks.Block) -> bool: equality_ids = [] for lcol, rcol in zip(block1.value_columns, block2.value_columns): - lcolmapped = lmap(lcol) - rcolmapped = rmap(rcol) + lcolmapped = lmap[lcol] + rcolmapped = rmap[rcol] joined_block, result_id = joined_block.apply_binary_op( lcolmapped, rcolmapped, ops.eq_nulls_match_op ) @@ -563,8 +563,8 @@ def align_rows( joined_index, (get_column_left, get_column_right) = left_block.index.join( right_block.index, how=join ) - left_columns = [get_column_left(col) for col in left_block.value_columns] - right_columns = [get_column_right(col) for col in right_block.value_columns] + left_columns = [get_column_left[col] for col in left_block.value_columns] + right_columns = [get_column_right[col] for col in right_block.value_columns] left_block = joined_index._block.select_columns(left_columns) right_block = joined_index._block.select_columns(right_columns) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 9b49645c71d..4548fca5934 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -39,6 +39,7 @@ import bigframes.core.guid as guid import bigframes.core.indexes as indexes import bigframes.core.joins as joins +import bigframes.core.joins.name_resolution as join_names import bigframes.core.ordering as ordering import bigframes.core.utils import bigframes.core.utils as utils @@ -97,7 +98,8 @@ def __init__( "'index_columns' and 'index_labels' must have equal length" ) if len(index_columns) == 0: - expr, new_index_col_id = expr.promote_offsets() + new_index_col_id = guid.generate_guid() + expr = expr.promote_offsets(new_index_col_id) index_columns = [new_index_col_id] self._index_columns = tuple(index_columns) # Index labels don't need complicated hierarchical access so can store as tuple @@ -260,7 +262,8 @@ def reset_index(self, drop: bool = True) -> Block: from Index classes that point to this block. """ block = self - expr, new_index_col_id = self._expr.promote_offsets() + new_index_col_id = guid.generate_guid() + expr = self._expr.promote_offsets(new_index_col_id) if drop: # Even though the index might be part of the ordering, keep that # ordering expression as reset_index shouldn't change the row @@ -833,7 +836,8 @@ def aggregate_all_and_stack( else: # axis_n == 1 # using offsets as identity to group on. # TODO: Allow to promote identity/total_order columns instead for better perf - expr_with_offsets, offset_col = self.expr.promote_offsets() + offset_col = guid.generate_guid() + expr_with_offsets = self.expr.promote_offsets(offset_col) stacked_expr = expr_with_offsets.unpivot( row_labels=self.column_labels.to_list(), index_col_ids=[guid.generate_guid()], @@ -952,9 +956,10 @@ def aggregate( ] by_column_labels = self._get_labels_for_columns(by_value_columns) labels = (*by_column_labels, *aggregate_labels) - result_expr_pruned, offsets_id = result_expr.select_columns( + offsets_id = guid.generate_guid() + result_expr_pruned = result_expr.select_columns( [*by_value_columns, *output_col_ids] - ).promote_offsets() + ).promote_offsets(offsets_id) return ( Block( @@ -975,7 +980,8 @@ def get_stat(self, column_id: str, stat: agg_ops.AggregateOp): aggregations = [(column_id, stat, stat.name) for stat in stats_to_fetch] expr = self.expr.aggregate(aggregations) - expr, offset_index_id = expr.promote_offsets() + offset_index_id = guid.generate_guid() + expr = expr.promote_offsets(offset_index_id) block = Block( expr, index_columns=[offset_index_id], @@ -999,7 +1005,8 @@ def get_corr_stat(self, column_id_left: str, column_id_right: str): ) ] expr = self.expr.corr_aggregate(corr_aggregations) - expr, offset_index_id = expr.promote_offsets() + offset_index_id = guid.generate_guid() + expr = expr.promote_offsets(offset_index_id) block = Block( expr, index_columns=[offset_index_id], @@ -1197,7 +1204,8 @@ def retrieve_repr_request_results( return formatted_df, count, query_job def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]: - expr, result_id = self._expr.promote_offsets() + result_id = guid.generate_guid() + expr = self._expr.promote_offsets(result_id) return ( Block( expr, @@ -1471,67 +1479,76 @@ def merge( "outer", "right", ], - left_col_ids: typing.Sequence[str], - right_col_ids: typing.Sequence[str], + left_join_ids: typing.Sequence[str], + right_join_ids: typing.Sequence[str], sort: bool, suffixes: tuple[str, str] = ("_x", "_y"), ) -> Block: - ( - joined_expr, - coalesced_join_cols, - (get_column_left, get_column_right), - ) = joins.join_by_column( + joined_expr = joins.join_by_column( self.expr, - left_col_ids, + left_join_ids, other.expr, - right_col_ids, + right_join_ids, how=how, - sort=sort, ) + get_column_left, get_column_right = join_names.JOIN_NAME_REMAPPER( + self.expr.column_ids, other.expr.column_ids + ) + result_columns = [] + matching_join_labels = [] + + coalesced_ids = [] + for left_id, right_id in zip(left_join_ids, right_join_ids): + coalesced_id = guid.generate_guid() + joined_expr = joined_expr.project_binary_op( + get_column_left[left_id], + get_column_right[right_id], + ops.coalesce_op, + coalesced_id, + ) + coalesced_ids.append(coalesced_id) + + for col_id in self.value_columns: + if col_id in left_join_ids: + key_part = left_join_ids.index(col_id) + matching_right_id = right_join_ids[key_part] + if ( + self.col_id_to_label[col_id] + == other.col_id_to_label[matching_right_id] + ): + matching_join_labels.append(self.col_id_to_label[col_id]) + result_columns.append(coalesced_ids[key_part]) + else: + result_columns.append(get_column_left[col_id]) + else: + result_columns.append(get_column_left[col_id]) + for col_id in other.value_columns: + if col_id in right_join_ids: + key_part = right_join_ids.index(col_id) + if other.col_id_to_label[matching_right_id] in matching_join_labels: + pass + else: + result_columns.append(get_column_right[col_id]) + else: + result_columns.append(get_column_right[col_id]) - # which join key parts should be coalesced - merge_join_key_mask = [ - str(self.col_id_to_label[left_id]) == str(other.col_id_to_label[right_id]) - for left_id, right_id in zip(left_col_ids, right_col_ids) - ] - labels_to_coalesce = [ - self.col_id_to_label[col_id] - for i, col_id in enumerate(left_col_ids) - if merge_join_key_mask[i] - ] - - def left_col_mapping(col_id: str) -> str: - if col_id in left_col_ids: - join_key_part = left_col_ids.index(col_id) - if merge_join_key_mask[join_key_part]: - return coalesced_join_cols[join_key_part] - return get_column_left(col_id) - - def right_col_mapping(col_id: str) -> typing.Optional[str]: - if col_id in right_col_ids: - join_key_part = right_col_ids.index(col_id) - if merge_join_key_mask[join_key_part]: - return None - return get_column_right(col_id) - - left_columns = [left_col_mapping(col_id) for col_id in self.value_columns] - - right_columns = [ - typing.cast(str, right_col_mapping(col_id)) - for col_id in other.value_columns - if right_col_mapping(col_id) - ] + if sort: + # sort uses coalesced join keys always + joined_expr = joined_expr.order_by( + [ordering.OrderingColumnReference(col_id) for col_id in coalesced_ids], + stable=True, + ) - expr = joined_expr.select_columns([*left_columns, *right_columns]) + joined_expr = joined_expr.select_columns(result_columns) labels = utils.merge_column_labels( self.column_labels, other.column_labels, - coalesce_labels=labels_to_coalesce, + coalesce_labels=matching_join_labels, suffixes=suffixes, ) - # Constructs default index - expr, offset_index_id = expr.promote_offsets() + offset_index_id = guid.generate_guid() + expr = joined_expr.promote_offsets(offset_index_id) return Block(expr, index_columns=[offset_index_id], column_labels=labels) def _force_reproject(self) -> Block: diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 7d15e676493..677bb8529c2 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -17,7 +17,7 @@ from __future__ import annotations import typing -from typing import Callable, Sequence, Tuple, Union +from typing import Mapping, Sequence, Tuple, Union import numpy as np import pandas @@ -27,6 +27,7 @@ import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks import bigframes.core.joins as joins +import bigframes.core.joins.name_resolution as join_names import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.dtypes @@ -413,7 +414,7 @@ def join( how="left", sort=False, block_identity_join: bool = False, - ) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: + ) -> Tuple[IndexValue, Tuple[Mapping[str, str], Mapping[str, str]],]: if not isinstance(other, IndexValue): # TODO(swast): We need to improve this error message to be more # actionable for the user. For example, it's possible they @@ -456,27 +457,34 @@ def join_mono_indexed( how="left", sort=False, block_identity_join: bool = False, -) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: - ( - combined_expr, - joined_index_col_names, - (get_column_left, get_column_right), - ) = joins.join_by_column( +) -> Tuple[IndexValue, Tuple[Mapping[str, str], Mapping[str, str]],]: + left_expr = left._block.expr + right_expr = right._block.expr + get_column_left, get_column_right = join_names.JOIN_NAME_REMAPPER( + left_expr.column_ids, right_expr.column_ids + ) + combined_expr = joins.join_by_column( left._block.expr, left._block.index_columns, right._block.expr, right._block.index_columns, how=how, - sort=sort, allow_row_identity_join=(not block_identity_join), ) # Drop original indices from each side. and used the coalesced combination generated by the join. - left_indices = [get_column_left(col_id) for col_id in left._block.index_columns] - right_indices = [get_column_right(col_id) for col_id in right._block.index_columns] - combined_expr = combined_expr.drop_columns(left_indices).drop_columns(right_indices) + left_index = get_column_left[left._block.index_columns[0]] + right_index = get_column_right[right._block.index_columns[0]] + # Drop original indices from each side. and used the coalesced combination generated by the join. + combined_expr, coalesced_join_cols = coalesce_columns( + combined_expr, [left_index], [right_index], how=how + ) + if sort: + combined_expr = combined_expr.order_by( + [order.OrderingColumnReference(col_id) for col_id in coalesced_join_cols] + ) block = blocks.Block( combined_expr, - index_columns=[*joined_index_col_names], + index_columns=coalesced_join_cols, column_labels=[*left._block.column_labels, *right._block.column_labels], index_labels=[left.name] if left.name == right.name else [None], ) @@ -493,7 +501,7 @@ def join_multi_indexed( how="left", sort=False, block_identity_join: bool = False, -) -> Tuple[IndexValue, Tuple[Callable[[str], str], Callable[[str], str]],]: +) -> Tuple[IndexValue, Tuple[Mapping[str, str], Mapping[str, str]],]: if not (left.is_uniquely_named() and right.is_uniquely_named()): raise ValueError("Joins not supported on indices with non-unique level names") @@ -508,25 +516,33 @@ def join_multi_indexed( right_join_ids = [right.resolve_level_name(name) for name in common_names] names_fully_match = len(left_only_names) == 0 and len(right_only_names) == 0 - ( - combined_expr, - joined_index_col_names, - (get_column_left, get_column_right), - ) = joins.join_by_column( - left._block.expr, + + left_expr = left._block.expr + right_expr = right._block.expr + get_column_left, get_column_right = join_names.JOIN_NAME_REMAPPER( + left_expr.column_ids, right_expr.column_ids + ) + + combined_expr = joins.join_by_column( + left_expr, left_join_ids, - right._block.expr, + right_expr, right_join_ids, how=how, - sort=sort, # If we're only joining on a subset of the index columns, we need to # perform a true join. - allow_row_identity_join=names_fully_match and not block_identity_join, + allow_row_identity_join=(names_fully_match and not block_identity_join), ) + left_ids_post_join = [get_column_left[id] for id in left_join_ids] + right_ids_post_join = [get_column_right[id] for id in right_join_ids] # Drop original indices from each side. and used the coalesced combination generated by the join. - combined_expr = combined_expr.drop_columns( - [get_column_left(col) for col in left_join_ids] - ).drop_columns([get_column_right(col) for col in right_join_ids]) + combined_expr, coalesced_join_cols = coalesce_columns( + combined_expr, left_ids_post_join, right_ids_post_join, how=how + ) + if sort: + combined_expr = combined_expr.order_by( + [order.OrderingColumnReference(col_id) for col_id in coalesced_join_cols] + ) if left.nlevels == 1: index_labels = right.names @@ -536,12 +552,13 @@ def join_multi_indexed( index_labels = [*common_names, *left_only_names, *right_only_names] def resolve_label_id(label: blocks.Label) -> str: + # if name is shared between both blocks, coalesce the values if label in common_names: - return joined_index_col_names[common_names.index(label)] + return coalesced_join_cols[common_names.index(label)] if label in left_only_names: - return get_column_left(left.resolve_level_name(label)) + return get_column_left[left.resolve_level_name(label)] if label in right_only_names: - return get_column_right(right.resolve_level_name(label)) + return get_column_right[right.resolve_level_name(label)] raise ValueError(f"Unexpected label: {label}") index_columns = [resolve_label_id(label) for label in index_labels] @@ -556,3 +573,29 @@ def resolve_label_id(label: blocks.Label) -> str: typing.cast(IndexValue, block.index), (get_column_left, get_column_right), ) + + +def coalesce_columns( + expr: core.ArrayValue, + left_ids: typing.Sequence[str], + right_ids: typing.Sequence[str], + how: str, +) -> Tuple[core.ArrayValue, Sequence[str]]: + result_ids = [] + for left_id, right_id in zip(left_ids, right_ids): + if how == "left" or how == "inner": + result_ids.append(left_id) + expr = expr.drop_columns([right_id]) + elif how == "right": + result_ids.append(right_id) + expr = expr.drop_columns([left_id]) + elif how == "outer": + coalesced_id = bigframes.core.guid.generate_guid() + expr = expr.project_binary_op( + left_id, right_id, ops.coalesce_op, coalesced_id + ) + expr = expr.drop_columns([left_id, right_id]) + result_ids.append(coalesced_id) + else: + raise ValueError(f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}") + return expr, result_ids diff --git a/bigframes/core/joins/name_resolution.py b/bigframes/core/joins/name_resolution.py new file mode 100644 index 00000000000..df946b3a590 --- /dev/null +++ b/bigframes/core/joins/name_resolution.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. +from __future__ import annotations + +from typing import Mapping, Sequence, Tuple + + +class JoinNameRemapper: + def __init__(self, namespace: str) -> None: + self._namespace = namespace + + def __call__( + self, left_column_ids: Sequence[str], right_column_ids: Sequence[str] + ) -> Tuple[Mapping[str, str], Mapping[str, str]]: + """ + When joining column ids from different namespaces, this function defines how names are remapped. + + Take care to map value column ids and hidden column ids in separate namespaces. This is important because value + column ids must be deterministic as they are referenced by dependent operators. The generation of hidden ids is + dependent on compilation context, and should be completely separated from value column id mappings. + """ + # This naming strategy depends on the number of value columns in source tables. + # This means column id mappings must be adjusted if pushing operations above or below join in transformation + new_left_ids = { + col: f"{self._namespace}_l_{i}" for i, col in enumerate(left_column_ids) + } + new_right_ids = { + col: f"{self._namespace}_r_{i}" for i, col in enumerate(right_column_ids) + } + return new_left_ids, new_right_ids + + +# Defines how column ids are remapped, regardless of join strategy or ordering mode +# Use this remapper for all value column remappings. +JOIN_NAME_REMAPPER = JoinNameRemapper("bfjoin") diff --git a/bigframes/core/joins/row_identity.py b/bigframes/core/joins/row_identity.py index 156e7aef407..76e456ec942 100644 --- a/bigframes/core/joins/row_identity.py +++ b/bigframes/core/joins/row_identity.py @@ -18,20 +18,20 @@ import functools import typing -from typing import Callable, Tuple import ibis import ibis.expr.types as ibis_types import bigframes.constants as constants import bigframes.core as core +import bigframes.core.joins.name_resolution as naming SUPPORTED_ROW_IDENTITY_HOW = {"outer", "left", "inner"} def join_by_row_identity( left: core.ArrayValue, right: core.ArrayValue, *, how: str -) -> Tuple[core.ArrayValue, Tuple[Callable[[str], str], Callable[[str], str]],]: +) -> core.ArrayValue: """Compute join when we are joining by row identity not a specific column.""" if how not in SUPPORTED_ROW_IDENTITY_HOW: raise NotImplementedError( @@ -62,31 +62,42 @@ def join_by_row_identity( left_mask = left_relative_predicates if how in ["right", "outer"] else None right_mask = right_relative_predicates if how in ["left", "outer"] else None + + # Public mapping must use JOIN_NAME_REMAPPER to stay in sync with consumers of join result + lpublicmapping, rpublicmapping = naming.JOIN_NAME_REMAPPER( + left.column_ids, right.column_ids + ) + lhiddenmapping, rhiddenmapping = naming.JoinNameRemapper(namespace="hidden")( + left._hidden_column_ids, right._hidden_column_ids + ) + map_left_id = {**lpublicmapping, **lhiddenmapping} + map_right_id = {**rpublicmapping, **rhiddenmapping} + joined_columns = [ - _mask_value(left._get_ibis_column(key), left_mask).name(map_left_id(key)) + _mask_value(left._get_ibis_column(key), left_mask).name(map_left_id[key]) for key in left.column_ids ] + [ - _mask_value(right._get_ibis_column(key), right_mask).name(map_right_id(key)) + _mask_value(right._get_ibis_column(key), right_mask).name(map_right_id[key]) for key in right.column_ids ] # If left isn't being masked, can just use left ordering if not left_mask: col_mapping = { - order_ref.column_id: map_left_id(order_ref.column_id) + order_ref.column_id: map_left_id[order_ref.column_id] for order_ref in left._ordering.ordering_value_columns } new_ordering = left._ordering.with_column_remap(col_mapping) else: ordering_columns = [ - col_ref.with_name(map_left_id(col_ref.column_id)) + col_ref.with_name(map_left_id[col_ref.column_id]) for col_ref in left._ordering.ordering_value_columns ] + [ - col_ref.with_name(map_right_id(col_ref.column_id)) + col_ref.with_name(map_right_id[col_ref.column_id]) for col_ref in right._ordering.ordering_value_columns ] left_total_order_cols = frozenset( - map_left_id(col) for col in left._ordering.total_ordering_columns + map_left_id[col] for col in left._ordering.total_ordering_columns ) # Assume that left ordering is sufficient since 1:1 join over same base table join_total_order_cols = left_total_order_cols @@ -95,12 +106,12 @@ def join_by_row_identity( ) hidden_ordering_columns = [ - left._get_hidden_ordering_column(key.column_id).name(map_left_id(key.column_id)) + left._get_hidden_ordering_column(key.column_id).name(map_left_id[key.column_id]) for key in left._ordering.ordering_value_columns if key.column_id in left._hidden_ordering_column_names.keys() ] + [ right._get_hidden_ordering_column(key.column_id).name( - map_right_id(key.column_id) + map_right_id[key.column_id] ) for key in right._ordering.ordering_value_columns if key.column_id in right._hidden_ordering_column_names.keys() @@ -114,18 +125,7 @@ def join_by_row_identity( ordering=new_ordering, predicates=combined_predicates, ) - return joined_expr, ( - lambda key: map_left_id(key), - lambda key: map_right_id(key), - ) - - -def map_left_id(left_side_id): - return f"{left_side_id}_x" - - -def map_right_id(right_side_id): - return f"{right_side_id}_y" + return joined_expr def _mask_value( diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/joins/single_column.py index f194b8f8c41..0c0e2008b58 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/joins/single_column.py @@ -16,17 +16,15 @@ from __future__ import annotations -import itertools import typing -from typing import Callable, Literal, Tuple +from typing import Literal, Mapping import ibis import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types -import bigframes.constants as constants import bigframes.core as core -import bigframes.core.guid as guid +import bigframes.core.joins.name_resolution as naming import bigframes.core.joins.row_identity import bigframes.core.ordering @@ -43,13 +41,8 @@ def join_by_column( "outer", "right", ], - sort: bool = False, allow_row_identity_join: bool = True, -) -> Tuple[ - core.ArrayValue, - typing.Sequence[str], - Tuple[Callable[[str], str], Callable[[str], str]], -]: +) -> core.ArrayValue: """Join two expressions by column equality. Arguments: @@ -62,14 +55,9 @@ def join_by_column( If True, allow matching by row identity. Set to False to always perform a true JOIN in generated SQL. Returns: - The joined expression and the objects needed to interpret it. - - * ArrayValue: Joined table with all columns from left and right. - * Sequence[str]: Column IDs of the coalesced join columns. Sometimes either the - left/right table will have missing rows. This column pulls the - non-NULL value from either left/right. - * Tuple[Callable, Callable]: For a given column ID from left or right, - respectively, return the new column id from the combined expression. + The joined expression. The resulting columns will be, in order, + first the coalesced join keys, then, all the left columns, and + finally, all the right columns. """ if ( allow_row_identity_join @@ -85,71 +73,33 @@ def join_by_column( for lcol, rcol in zip(left_column_ids, right_column_ids) ) ): - combined_expr, ( - get_column_left, - get_column_right, - ) = bigframes.core.joins.row_identity.join_by_row_identity(left, right, how=how) - left_join_keys = [ - combined_expr._get_ibis_column(get_column_left(col)) - for col in left_column_ids - ] - right_join_keys = [ - combined_expr._get_ibis_column(get_column_right(col)) - for col in right_column_ids - ] - join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how) - join_key_ids = [col.get_name() for col in join_key_cols] - combined_expr = combined_expr._projection( - [*join_key_cols, *combined_expr.columns] - ) - if sort: - combined_expr = combined_expr.order_by( - [ - core.OrderingColumnReference(join_col_id) - for join_col_id in join_key_ids - ] - ) - return ( - combined_expr, - join_key_ids, - ( - get_column_left, - get_column_right, - ), + return bigframes.core.joins.row_identity.join_by_row_identity( + left, right, how=how ) else: - lmapping = { - col_id: guid.generate_guid() - for col_id in itertools.chain( - left.column_ids, left._hidden_ordering_column_names - ) - } - rmapping = { - col_id: guid.generate_guid() - for col_id in itertools.chain( - right.column_ids, right._hidden_ordering_column_names - ) - } - - def get_column_left(col_id): - return lmapping[col_id] - - def get_column_right(col_id): - return rmapping[col_id] + # Value column mapping must use JOIN_NAME_REMAPPER to stay in sync with consumers of join result + l_public_mapping, r_public_mapping = naming.JOIN_NAME_REMAPPER( + left.column_ids, right.column_ids + ) + l_hidden_mapping, r_hidden_mapping = naming.JoinNameRemapper( + namespace="hidden" + )(left._hidden_column_ids, right._hidden_column_ids) + l_mapping = {**l_public_mapping, **l_hidden_mapping} + r_mapping = {**r_public_mapping, **r_hidden_mapping} left_table = left._to_ibis_expr( "unordered", expose_hidden_cols=True, - col_id_overrides=lmapping, + col_id_overrides=l_mapping, ) right_table = right._to_ibis_expr( "unordered", expose_hidden_cols=True, - col_id_overrides=rmapping, + col_id_overrides=r_mapping, ) join_conditions = [ - value_to_join_key(left_table[lmapping[left_index]]) - == value_to_join_key(right_table[rmapping[right_index]]) + value_to_join_key(left_table[l_mapping[left_index]]) + == value_to_join_key(right_table[r_mapping[right_index]]) for left_index, right_index in zip(left_column_ids, right_column_ids) ] @@ -158,97 +108,39 @@ def get_column_right(col_id): right_table, predicates=join_conditions, how=how, - lname="{name}_x", - rname="{name}_y", ) # Preserve ordering accross joins. ordering = join_orderings( left._ordering, right._ordering, - get_column_left, - get_column_right, + l_mapping, + r_mapping, left_order_dominates=(how != "right"), ) - left_join_keys = [ - combined_table[get_column_left(col)] for col in left_column_ids - ] - right_join_keys = [ - combined_table[get_column_right(col)] for col in right_column_ids - ] - join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how) # We could filter out the original join columns, but predicates/ordering # might still reference them in implicit joins. - columns = ( - join_key_cols - + [combined_table[get_column_left(col.get_name())] for col in left.columns] - + [ - combined_table[get_column_right(col.get_name())] - for col in right.columns - ] - ) + columns = [ + combined_table[l_mapping[col.get_name()]] for col in left.columns + ] + [combined_table[r_mapping[col.get_name()]] for col in right.columns] hidden_ordering_columns = [ *[ - combined_table[get_column_left(col.get_name())] - for col in left.hidden_ordering_columns + combined_table[l_hidden_mapping[col.get_name()]] + for col in left._hidden_ordering_columns ], *[ - combined_table[get_column_right(col.get_name())] - for col in right.hidden_ordering_columns + combined_table[r_hidden_mapping[col.get_name()]] + for col in right._hidden_ordering_columns ], ] - combined_expr = core.ArrayValue( + return core.ArrayValue( left._session, combined_table, columns=columns, hidden_ordering_columns=hidden_ordering_columns, ordering=ordering, ) - if sort: - combined_expr = combined_expr.order_by( - [ - core.OrderingColumnReference(join_key_col.get_name()) - for join_key_col in join_key_cols - ] - ) - return ( - combined_expr, - [key.get_name() for key in join_key_cols], - (get_column_left, get_column_right), - ) - - -def get_coalesced_join_cols( - left_join_cols: typing.Iterable[ibis_types.Value], - right_join_cols: typing.Iterable[ibis_types.Value], - how: str, -) -> typing.List[ibis_types.Value]: - join_key_cols: list[ibis_types.Value] = [] - for left_col, right_col in zip(left_join_cols, right_join_cols): - if how == "left" or how == "inner": - join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) - elif how == "right": - join_key_cols.append(right_col.name(guid.generate_guid(prefix="index_"))) - elif how == "outer": - # The left index and the right index might contain null values, for - # example due to an outer join with different numbers of rows. Coalesce - # these to take the index value from either column. - # Use a random name in case the left index and the right index have the - # same name. In such a case, _x and _y suffixes will already be used. - # Don't need to coalesce if they are exactly the same column. - if left_col.name("index").equals(right_col.name("index")): - join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_"))) - else: - join_key_cols.append( - ibis.coalesce( - left_col, - right_col, - ).name(guid.generate_guid(prefix="index_")) - ) - else: - raise ValueError(f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}") - return join_key_cols def value_to_join_key(value: ibis_types.Value): @@ -261,16 +153,16 @@ def value_to_join_key(value: ibis_types.Value): def join_orderings( left: core.ExpressionOrdering, right: core.ExpressionOrdering, - left_id_mapping: Callable[[str], str], - right_id_mapping: Callable[[str], str], + left_id_mapping: Mapping[str, str], + right_id_mapping: Mapping[str, str], left_order_dominates: bool = True, ) -> core.ExpressionOrdering: left_ordering_refs = [ - ref.with_name(left_id_mapping(ref.column_id)) + ref.with_name(left_id_mapping[ref.column_id]) for ref in left.all_ordering_columns ] right_ordering_refs = [ - ref.with_name(right_id_mapping(ref.column_id)) + ref.with_name(right_id_mapping[ref.column_id]) for ref in right.all_ordering_columns ] if left_order_dominates: @@ -279,10 +171,10 @@ def join_orderings( joined_refs = [*right_ordering_refs, *left_ordering_refs] left_total_order_cols = frozenset( - [left_id_mapping(id) for id in left.total_ordering_columns] + [left_id_mapping[id] for id in left.total_ordering_columns] ) right_total_order_cols = frozenset( - [right_id_mapping(id) for id in right.total_ordering_columns] + [right_id_mapping[id] for id in right.total_ordering_columns] ) return core.ExpressionOrdering( ordering_value_columns=joined_refs, diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 015a7642f8b..c91ddffada8 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -419,7 +419,7 @@ def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame: get_column_right, ) = self._block.index.join(key._block.index, how="left") block = combined_index._block - filter_col_id = get_column_right(key._value_column) + filter_col_id = get_column_right[key._value_column] block = block.filter(filter_col_id) block = block.drop_columns([filter_col_id]) return DataFrame(block) @@ -560,18 +560,18 @@ def _apply_series_binop( ) series_column_id = other._value_column - series_col = get_column_right(series_column_id) + series_col = get_column_right[series_column_id] block = joined_index._block for column_id, label in zip( self._block.value_columns, self._block.column_labels ): block, _ = block.apply_binary_op( - get_column_left(column_id), + get_column_left[column_id], series_col, op, result_label=label, ) - block = block.drop_columns([get_column_left(column_id)]) + block = block.drop_columns([get_column_left[column_id]]) block = block.drop_columns([series_col]) block = block.with_index_labels(self.index.names) @@ -603,22 +603,22 @@ def _apply_dataframe_binop( left_col_id = self._block.value_columns[left_index] right_col_id = other._block.value_columns[right_index] block, result_col_id = block.apply_binary_op( - get_column_left(left_col_id), - get_column_right(right_col_id), + get_column_left[left_col_id], + get_column_right[right_col_id], op, ) binop_result_ids.append(result_col_id) elif left_index >= 0: left_col_id = self._block.value_columns[left_index] block, result_col_id = block.apply_unary_op( - get_column_left(left_col_id), + get_column_left[left_col_id], ops.partial_right(op, None), ) binop_result_ids.append(result_col_id) elif right_index >= 0: right_col_id = other._block.value_columns[right_index] block, result_col_id = block.apply_unary_op( - get_column_right(right_col_id), + get_column_right[right_col_id], ops.partial_left(op, None), ) binop_result_ids.append(result_col_id) @@ -974,7 +974,7 @@ def _drop_by_index(self, index: indexes.Index) -> DataFrame: block.index ) - new_ordering_col = get_column_right(ordering_col) + new_ordering_col = get_column_right[ordering_col] drop_block = joined_index._block drop_block, drop_col = drop_block.apply_unary_op( new_ordering_col, @@ -983,7 +983,7 @@ def _drop_by_index(self, index: indexes.Index) -> DataFrame: drop_block = drop_block.filter(drop_col) original_columns = [ - get_column_left(column) for column in self._block.value_columns + get_column_left[column] for column in self._block.value_columns ] drop_block = drop_block.select_columns(original_columns) return DataFrame(drop_block) @@ -1119,7 +1119,8 @@ def _assign_single_item( # local_df is likely (but not guarunteed) to be cached locally # since the original list came from memory and so is probably < MAX_INLINE_DF_SIZE - this_expr, this_offsets_col_id = self._get_block()._expr.promote_offsets() + this_offsets_col_id = bigframes.core.guid.generate_guid() + this_expr = self._get_block()._expr.promote_offsets(this_offsets_col_id) block = blocks.Block( expr=this_expr, index_labels=self.index.names, @@ -1156,10 +1157,10 @@ def _assign_series_join_on_index( ) column_ids = [ - get_column_left(col_id) for col_id in self._block.cols_matching_label(label) + get_column_left[col_id] for col_id in self._block.cols_matching_label(label) ] block = joined_index._block - source_column = get_column_right(series._value_column) + source_column = get_column_right[series._value_column] # Replace each column matching the label for column_id in column_ids: @@ -2032,8 +2033,8 @@ def _groupby_series( key._block.index, how="inner" if dropna else "left" ) col_ids = [ - *[get_column_left(value) for value in col_ids], - get_column_right(key._value_column), + *[get_column_left[value] for value in col_ids], + get_column_right[key._value_column], ] block = combined_index._block else: diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index f330a703b20..a29dd36c72a 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -913,6 +913,16 @@ def ge_op( return x >= y +def coalesce_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + if x.name("name").equals(y.name("name")): + return x + else: + return ibis.coalesce(x, y) + + @short_circuit_nulls(ibis_dtypes.int) def floordiv_op( x: ibis_types.Value, diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index fc76d07edbb..b9abb2cc038 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -196,8 +196,8 @@ def _align_n( get_column_right, ) = block.index.join(other._block.index, how=how) value_ids = [ - *[get_column_left(value) for value in value_ids], - get_column_right(other._value_column), + *[get_column_left[value] for value in value_ids], + get_column_right[other._value_column], ] block = combined_index._block else: diff --git a/bigframes/series.py b/bigframes/series.py index 56e1b43a032..4f2f73bd605 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1136,10 +1136,10 @@ def _groupby_values( key._block.index, how="inner" if dropna else "left" ) - value_col = get_column_left(self._value_column) + value_col = get_column_left[self._value_column] grouping_cols = [ - *[get_column_left(value) for value in grouping_cols], - get_column_right(key._value_column), + *[get_column_left[value] for value in grouping_cols], + get_column_right[key._value_column], ] block = combined_index._block else: From db51fe340f644a0d7c911c11d92c8299a4be3446 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 13 Oct 2023 12:20:38 -0700 Subject: [PATCH 78/80] docs: add open-source link in API doc (#106) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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://togithub.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 # 🦕 --- README.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/README.rst b/README.rst index b2ad1ee0179..7b4f1986bcd 100644 --- a/README.rst +++ b/README.rst @@ -13,6 +13,7 @@ BigQuery DataFrames is an open-source package. You can run Documentation ------------- +* `BigQuery DataFrames source code (GitHub) `_ * `BigQuery DataFrames sample notebooks `_ * `BigQuery DataFrames API reference `_ * `BigQuery documentation `_ From 1b3f3a5374915b2833c6c1ac05670e9708f07bff Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 13 Oct 2023 14:12:20 -0700 Subject: [PATCH 79/80] docs: update ML overview API doc (#105) --- README.rst | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/README.rst b/README.rst index 7b4f1986bcd..5ddb4a7639a 100644 --- a/README.rst +++ b/README.rst @@ -95,10 +95,18 @@ using the and the `bigframes.ml.compose module `_. BigQuery DataFrames offers the following transformations: -* Use the `OneHotEncoder class `_ - in the ``bigframes.ml.preprocessing`` module to transform categorical values into numeric format. +* Use the `KBinsDiscretizer class `_ + in the ``bigframes.ml.preprocessing`` module to bin continuous data into intervals. +* Use the `LabelEncoder class `_ + in the ``bigframes.ml.preprocessing`` module to normalize the target labels as integer values. +* Use the `MaxAbsScaler class `_ + in the ``bigframes.ml.preprocessing`` module to scale each feature to the range ``[-1, 1]`` by its maximum absolute value. +* Use the `MinMaxScaler class `_ + in the ``bigframes.ml.preprocessing`` module to standardize features by scaling each feature to the range ``[0, 1]``. * Use the `StandardScaler class `_ in the ``bigframes.ml.preprocessing`` module to standardize features by removing the mean and scaling to unit variance. +* Use the `OneHotEncoder class `_ + in the ``bigframes.ml.preprocessing`` module to transform categorical values into numeric format. * Use the `ColumnTransformer class `_ in the ``bigframes.ml.compose`` module to apply transformers to DataFrames columns. From 752a1d606f3e79554bffa7c4f456ee938ec446bc Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Fri, 13 Oct 2023 15:44:13 -0700 Subject: [PATCH 80/80] style: improve cancellation string (#111) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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://togithub.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 # 🦕 --- bigframes/formatting_helpers.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py index f3e3a76ce42..6851bdd2bd4 100644 --- a/bigframes/formatting_helpers.py +++ b/bigframes/formatting_helpers.py @@ -155,7 +155,10 @@ def wait_for_query_job( raise except KeyboardInterrupt: query_job.cancel() - print(f"Requested cancelation for {query_job.job_type} {query_job.job_id}...") + print( + f"Requested cancellation for {query_job.job_type.capitalize()}" + f" job {query_job.job_id} in location {query_job.location}..." + ) # begin the cancel request before immediately rethrowing raise @@ -197,7 +200,10 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): raise except KeyboardInterrupt: job.cancel() - print(f"Requested cancelation for {job.job_type} {job.job_id}...") + print( + f"Requested cancellation for {job.job_type.capitalize()}" + f" job {job.job_id} in location {job.location}..." + ) # begin the cancel request before immediately rethrowing raise