From 765a7af6db9ede660dbae65a665ff8896da4ba49 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 14 Dec 2022 20:23:12 +0000 Subject: [PATCH 001/197] 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 002/197] 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 003/197] 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 004/197] 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 005/197] 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 006/197] 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 007/197] 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 008/197] 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 009/197] 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 010/197] 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 011/197] 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 012/197] 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 013/197] 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 014/197] 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 015/197] 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 016/197] 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 017/197] 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 018/197] 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 019/197] 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 020/197] 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 021/197] 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 022/197] 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 023/197] 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 024/197] 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 025/197] 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 026/197] 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 027/197] 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 028/197] 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 029/197] 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 030/197] 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 031/197] 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 032/197] 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 033/197] 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 034/197] 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 035/197] 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 036/197] 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 037/197] 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 038/197] 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 039/197] 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 040/197] 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 041/197] 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 042/197] 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 043/197] 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 044/197] 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 045/197] 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 046/197] 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 047/197] 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 048/197] 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 049/197] 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 050/197] 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 051/197] 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 052/197] 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 053/197] 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 054/197] 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 055/197] 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 056/197] 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 057/197] 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 058/197] 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 059/197] 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 060/197] 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 061/197] 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 062/197] 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 063/197] 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 064/197] 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 065/197] 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 066/197] 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", - " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", + " \n", + " \n", " \n", " \n", - " \n", - " \n", + " \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", + " \n", " \n", " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", " \n", - " \n", + " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", " \n", - " \n", + " \n", " \n", " \n", - " \n", + " \n", " \n", - " \n", - " \n", - " \n", - " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \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 067/197] 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 068/197] 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 069/197] 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 070/197] 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 071/197] 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 072/197] 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 073/197] 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 074/197] 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 075/197] 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 076/197] 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 077/197] 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 078/197] 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 079/197] 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 080/197] 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 From e6c0cd1777736e0fa7285da59625fbac487573bd Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 16 Oct 2023 16:30:16 -0500 Subject: [PATCH 081/197] perf: if primary keys are defined, `read_gbq` avoids copying table data (#112) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit We make the same uniqueness assumption as the query engine and use these columns as the total ordering. 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 internal issue b/305260214 🦕 --- bigframes/core/indexers.py | 3 +- bigframes/session.py | 110 +++++++++++++++++++++++++---- noxfile.py | 1 - tests/system/conftest.py | 31 ++++++++ tests/system/small/test_session.py | 25 +++++++ 5 files changed, 153 insertions(+), 17 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index a74880041ca..09f0d5956c1 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -21,6 +21,7 @@ import pandas as pd import bigframes.constants as constants +import bigframes.core.blocks import bigframes.core.guid as guid import bigframes.core.indexes as indexes import bigframes.core.scalar @@ -214,7 +215,7 @@ def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar: raise ValueError(error_message) if len(key) != 2: raise TypeError(error_message) - block = self._dataframe._block + block: bigframes.core.blocks.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]] diff --git a/bigframes/session.py b/bigframes/session.py index 4f509f0704a..6c1160c88e6 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -498,6 +498,8 @@ def read_gbq_query( See also: :meth:`Session.read_gbq`. """ + # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so + # these docstrings are inline. return self._read_gbq_query( query=query, index_col=index_col, @@ -515,8 +517,6 @@ def _read_gbq_query( 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: @@ -561,6 +561,8 @@ def read_gbq_table( See also: :meth:`Session.read_gbq`. """ + # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so + # these docstrings are inline. return self._read_gbq_table( query=query, index_col=index_col, @@ -569,6 +571,62 @@ def read_gbq_table( api_name="read_gbq_table", ) + def _read_gbq_table_to_ibis_with_total_ordering( + self, + table_ref: bigquery.table.TableReference, + *, + api_name: str, + ) -> Tuple[ibis_types.Table, Optional[Sequence[str]]]: + """Create a read-only Ibis table expression representing a table. + + If we can get a total ordering from the table, such as via primary key + column(s), then return those too so that ordering generation can be + avoided. + """ + if table_ref.dataset_id.upper() == "_SESSION": + # _SESSION tables aren't supported by the tables.get REST API. + return ( + self.ibis_client.sql( + f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" + ), + None, + ) + + table_expression = self.ibis_client.table( + table_ref.table_id, + database=f"{table_ref.project}.{table_ref.dataset_id}", + ) + + # If there are primary keys defined, the query engine assumes these + # columns are unique, even if the constraint is not enforced. We make + # the same assumption and use these columns as the total ordering keys. + table = self.bqclient.get_table(table_ref) + + # TODO(b/305264153): Use public properties to fetch primary keys once + # added to google-cloud-bigquery. + primary_keys = ( + table._properties.get("tableConstraints", {}) + .get("primaryKey", {}) + .get("columns") + ) + + if not primary_keys: + return table_expression, None + else: + # Read from a snapshot since we won't have to copy the table data to create a total ordering. + job_config = bigquery.QueryJobConfig() + job_config.labels["bigframes-api"] = api_name + current_timestamp = list( + self.bqclient.query( + "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", + job_config=job_config, + ).result() + )[0][0] + table_expression = self.ibis_client.sql( + bigframes_io.create_snapshot_sql(table_ref, current_timestamp) + ) + return table_expression, primary_keys + def _read_gbq_table( self, query: str, @@ -581,24 +639,19 @@ def _read_gbq_table( 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 # 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: - table_expression = self.ibis_client.table( - table_ref.table_id, - database=f"{table_ref.project}.{table_ref.dataset_id}", - ) + ( + table_expression, + total_ordering_cols, + ) = self._read_gbq_table_to_ibis_with_total_ordering( + table_ref, + api_name=api_name, + ) for key in col_order: if key not in table_expression.columns: @@ -624,7 +677,34 @@ def _read_gbq_table( ordering = None is_total_ordering = False - if len(index_cols) != 0: + if total_ordering_cols is not None: + # Note: currently, this a table has a total ordering only when the + # primary key(s) are set on a table. The query engine assumes such + # columns are unique, even if not enforced. + is_total_ordering = True + ordering = core.ExpressionOrdering( + ordering_value_columns=[ + core.OrderingColumnReference(column_id) + for column_id in total_ordering_cols + ], + total_ordering_columns=frozenset(total_ordering_cols), + ) + + if len(index_cols) != 0: + index_labels = typing.cast(List[Optional[str]], index_cols) + else: + # Use the total_ordering_cols to project offsets to use as the default index. + table_expression = table_expression.order_by(index_cols) + default_index_id = guid.generate_guid("bigframes_index_") + default_index_col = ( + ibis.row_number().cast(ibis_dtypes.int64).name(default_index_id) + ) + table_expression = table_expression.mutate( + **{default_index_id: default_index_col} + ) + index_cols = [default_index_id] + index_labels = [None] + elif 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 ( diff --git a/noxfile.py b/noxfile.py index 84e5ab11bb4..54ccdb9a877 100644 --- a/noxfile.py +++ b/noxfile.py @@ -89,7 +89,6 @@ "system", "doctest", "cover", - "release_dry_run", ] # Error if a python version is missing diff --git a/tests/system/conftest.py b/tests/system/conftest.py index ed22a3e8dac..f36a29b0ab1 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -17,6 +17,7 @@ import logging import math import pathlib +import textwrap import typing from typing import Dict, Optional @@ -795,6 +796,36 @@ def penguins_randomforest_classifier_model_name( return model_name +@pytest.fixture(scope="session") +def usa_names_grouped_table( + session: bigframes.Session, dataset_id_permanent +) -> bigquery.Table: + """Provides a table with primary key(s) set.""" + table_id = f"{dataset_id_permanent}.usa_names_grouped" + try: + return session.bqclient.get_table(table_id) + except google.cloud.exceptions.NotFound: + query = textwrap.dedent( + f""" + CREATE TABLE `{dataset_id_permanent}.usa_names_grouped` + ( + total_people INT64, + name STRING, + gender STRING, + year INT64, + PRIMARY KEY(name, gender, year) NOT ENFORCED + ) + AS + SELECT SUM(`number`) AS total_people, name, gender, year + FROM `bigquery-public-data.usa_names.usa_1910_2013` + GROUP BY name, gender, year + """ + ) + job = session.bqclient.query(query) + job.result() + return session.bqclient.get_table(table_id) + + @pytest.fixture() def deferred_repr(): bigframes.options.display.repr_mode = "deferred" diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index bfe9bc8d0f9..127a88a760b 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -20,6 +20,7 @@ from typing import List import google.api_core.exceptions +import google.cloud.bigquery as bigquery import numpy as np import pandas as pd import pytest @@ -231,6 +232,30 @@ def test_read_gbq_w_anonymous_query_results_table(session: bigframes.Session): pd.testing.assert_frame_equal(result, expected, check_dtype=False) +def test_read_gbq_w_primary_keys_table( + session: bigframes.Session, usa_names_grouped_table: bigquery.Table +): + table = usa_names_grouped_table + # TODO(b/305264153): Use public properties to fetch primary keys once + # added to google-cloud-bigquery. + primary_keys = ( + table._properties.get("tableConstraints", {}) + .get("primaryKey", {}) + .get("columns") + ) + assert len(primary_keys) != 0 + + df = session.read_gbq(f"{table.project}.{table.dataset_id}.{table.table_id}") + result = df.head(100).to_pandas() + + # Verify that the DataFrame is already sorted by primary keys. + sorted_result = result.sort_values(primary_keys) + pd.testing.assert_frame_equal(result, sorted_result) + + # Verify that we're working from a snapshot rather than a copy of the table. + assert "FOR SYSTEM_TIME AS OF TIMESTAMP" in df.sql + + @pytest.mark.parametrize( ("query_or_table", "max_results"), [ From 752b01ff9df114c54ed58eb96956e9ce34a8ed47 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Mon, 16 Oct 2023 17:10:00 -0700 Subject: [PATCH 082/197] feat: add AtIndexer getitems (#107) * feat: add AtIndexer getitems * fix third party docstrings * use loc from at --------- Co-authored-by: Henry J Solberg --- bigframes/core/indexers.py | 24 +++++++++++++++++++ bigframes/dataframe.py | 4 ++++ bigframes/series.py | 4 ++++ tests/system/small/test_dataframe.py | 23 ++++++++++++++++++ tests/system/small/test_series.py | 10 ++++++++ .../bigframes_vendored/pandas/core/frame.py | 5 ++++ .../bigframes_vendored/pandas/core/series.py | 5 ++++ 7 files changed, 75 insertions(+) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 09f0d5956c1..01c1162ae6d 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -108,6 +108,16 @@ def __getitem__(self, key: int) -> bigframes.core.scalar.Scalar: return self._series.iloc[key] +class AtSeriesIndexer: + def __init__(self, series: bigframes.series.Series): + self._series = series + + def __getitem__( + self, key: LocSingleKey + ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: + return self._series.loc[key] + + class LocDataFrameIndexer: def __init__(self, dataframe: bigframes.dataframe.DataFrame): self._dataframe = dataframe @@ -221,6 +231,20 @@ def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar: return column.iloc[key[0]] +class AtDataFrameIndexer: + def __init__(self, dataframe: bigframes.dataframe.DataFrame): + self._dataframe = dataframe + + def __getitem__( + self, key: tuple + ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: + if not isinstance(key, tuple): + raise TypeError( + "DataFrame.at should be indexed by a (row label, column name) tuple." + ) + return self._dataframe.loc[key] + + @typing.overload def _loc_getitem_series_or_dataframe( series_or_dataframe: bigframes.series.Series, key diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c91ddffada8..5740d2c4dc8 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -259,6 +259,10 @@ def iloc(self) -> indexers.ILocDataFrameIndexer: def iat(self) -> indexers.IatDataFrameIndexer: return indexers.IatDataFrameIndexer(self) + @property + def at(self) -> indexers.AtDataFrameIndexer: + return indexers.AtDataFrameIndexer(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 4f2f73bd605..84ca2a578f8 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -91,6 +91,10 @@ def iloc(self) -> bigframes.core.indexers.IlocSeriesIndexer: def iat(self) -> bigframes.core.indexers.IatSeriesIndexer: return bigframes.core.indexers.IatSeriesIndexer(self) + @property + def at(self) -> bigframes.core.indexers.AtSeriesIndexer: + return bigframes.core.indexers.AtSeriesIndexer(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 9f1092d09d1..ecafb7c1bf5 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2196,6 +2196,29 @@ def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index ) +def test_at_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.at[index, "int64_too"] + pd_result = scalars_pandas_df_index.at[index, "int64_too"] + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_at_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.at[index, "string_col"] + pd_result = scalars_pandas_df_index.at[index, "string_col"] + assert bf_result == pd_result + + def test_loc_setitem_bool_series_scalar_new_col(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs bf_df = scalars_df.copy() diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 802425510aa..f2ced841daf 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -1953,6 +1953,16 @@ def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, ste ) +def test_at(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["string_col"].at[index] + pd_result = scalars_pandas_df_index["string_col"].at[index] + + assert bf_result == pd_result + + 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] diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index e54f984d59b..621d052cb8e 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2112,3 +2112,8 @@ def iloc(self): def iat(self): """Access a single value for a row/column pair by integer position.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def at(self): + """Access a single value for a row/column label pair.""" + 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 03729922d5a..ec3c07bfa85 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1833,3 +1833,8 @@ def iloc(self): def iat(self): """Access a single value for a row/column pair by integer position.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def at(self): + """Access a single value for a row/column label pair.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From ec10c4a5a7833c42e28fe9e7b734bc0c4fb84b6e Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 17 Oct 2023 02:55:40 +0000 Subject: [PATCH 083/197] feat: Support external packages in `remote_function` (#98) * feat: Support external packages in `remote_function` * Update code sample demonstrating external packages for `remote_function` * GCF customization for hackathon --- bigframes/pandas/__init__.py | 2 + bigframes/remote_function.py | 49 ++++++++++++++++------ bigframes/session.py | 8 +++- samples/snippets/remote_function.py | 18 +++++--- tests/system/large/test_remote_function.py | 45 ++++++++++++++++++++ 5 files changed, 103 insertions(+), 19 deletions(-) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index ded4760fa3b..833f0d1d1dd 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -400,6 +400,7 @@ def remote_function( bigquery_connection: Optional[str] = None, reuse: bool = True, name: Optional[str] = None, + packages: Optional[Sequence[str]] = None, ): return global_session.with_default_session( bigframes.session.Session.remote_function, @@ -409,6 +410,7 @@ def remote_function( bigquery_connection=bigquery_connection, reuse=reuse, name=name, + packages=packages, ) diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index fd9aec825fa..c82ba840564 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -100,9 +100,12 @@ def get_remote_function_locations(bq_location): return bq_location, cloud_function_region -def _get_hash(def_): +def _get_hash(def_, package_requirements=None): "Get hash (32 digits alphanumeric) of a function." def_repr = cloudpickle.dumps(def_, protocol=_pickle_protocol_version) + if package_requirements: + for p in sorted(package_requirements): + def_repr += p.encode() return hashlib.md5(def_repr).hexdigest() @@ -129,18 +132,18 @@ class IbisSignature(NamedTuple): output_type: IbisDataType -def get_cloud_function_name(def_, uniq_suffix=None): +def get_cloud_function_name(def_, uniq_suffix=None, package_requirements=None): "Get a name for the cloud function for the given user defined function." - cf_name = _get_hash(def_) + cf_name = _get_hash(def_, package_requirements) 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): +def get_remote_function_name(def_, uniq_suffix=None, package_requirements=None): "Get a name for the BQ remote function for the given user defined function." - bq_rf_name = _get_hash(def_) + bq_rf_name = _get_hash(def_, package_requirements) bq_rf_name = f"bigframes_{bq_rf_name}" # for identification if uniq_suffix: bq_rf_name = f"{bq_rf_name}_{uniq_suffix}" @@ -200,7 +203,8 @@ def create_bq_remote_function( RETURNS {bq_function_return_type} REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}` OPTIONS ( - endpoint = "{endpoint}" + endpoint = "{endpoint}", + max_batching_rows = 1000 )""" logger.info(f"Creating BQ remote function: {create_function_ddl}") @@ -320,11 +324,14 @@ def {handler_func_name}(request): return handler_func_name - def generate_cloud_function_code(self, def_, dir): + def generate_cloud_function_code(self, def_, dir, package_requirements=None): """Generate the cloud function code for a given user defined function.""" # requirements.txt requirements = ["cloudpickle >= 2.1.0"] + if package_requirements: + requirements.extend(package_requirements) + requirements = sorted(requirements) requirements_txt = os.path.join(dir, "requirements.txt") with open(requirements_txt, "w") as f: f.write("\n".join(requirements)) @@ -333,12 +340,14 @@ def generate_cloud_function_code(self, def_, dir): entry_point = self.generate_cloud_function_main_code(def_, dir) return entry_point - def create_cloud_function(self, def_, cf_name): + def create_cloud_function(self, def_, cf_name, package_requirements=None): """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) + entry_point = self.generate_cloud_function_code( + def_, dir, package_requirements + ) archive_path = shutil.make_archive(dir, "zip", dir) # We are creating cloud function source code from the currently running @@ -392,6 +401,9 @@ def create_cloud_function(self, def_, cf_name): function.build_config.source.storage_source.object_ = ( upload_url_response.storage_source.object_ ) + function.service_config = functions_v2.ServiceConfig() + function.service_config.available_memory = "1024M" + function.service_config.timeout_seconds = 600 create_function_request.function = function # Create the cloud function and wait for it to be ready to use @@ -422,6 +434,7 @@ def provision_bq_remote_function( output_type, reuse, name, + package_requirements, ): """Provision a BigQuery remote function.""" # If reuse of any existing function with the same name (indicated by the @@ -435,19 +448,25 @@ def provision_bq_remote_function( # Derive the name of the cloud function underlying the intended BQ # remote function - cloud_function_name = get_cloud_function_name(def_, uniq_suffix) + cloud_function_name = get_cloud_function_name( + def_, uniq_suffix, package_requirements + ) cf_endpoint = self.get_cloud_function_endpoint(cloud_function_name) # Create the cloud function if it does not exist if not cf_endpoint: - cf_endpoint = self.create_cloud_function(def_, cloud_function_name) + cf_endpoint = self.create_cloud_function( + def_, cloud_function_name, package_requirements + ) else: logger.info(f"Cloud function {cloud_function_name} already exists.") # 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) + remote_function_name = get_remote_function_name( + def_, uniq_suffix, package_requirements + ) rf_endpoint, rf_conn = self.get_remote_function_specs(remote_function_name) # Create the BQ remote function in following circumstances: @@ -619,6 +638,7 @@ def remote_function( bigquery_connection: Optional[str] = None, reuse: bool = True, name: Optional[str] = None, + packages: Optional[Sequence[str]] = None, ): """Decorator to turn a user defined function into a BigQuery remote function. @@ -710,6 +730,10 @@ def remote_function( 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. + packages (str[], Optional): + Explicit name of the external package dependencies. Each dependency + is added to the `requirements.txt` as is, and can be of the form + supported in https://pip.pypa.io/en/stable/reference/requirements-file-format/. """ import bigframes.pandas as bpd @@ -821,6 +845,7 @@ def wrapper(f): ibis_signature.output_type, reuse, name, + packages, ) node = remote_function_node(dataset_ref.routine(rf_name), ibis_signature) diff --git a/bigframes/session.py b/bigframes/session.py index 6c1160c88e6..fa5b4153507 100644 --- a/bigframes/session.py +++ b/bigframes/session.py @@ -1413,6 +1413,7 @@ def remote_function( bigquery_connection: Optional[str] = None, reuse: bool = True, name: Optional[str] = None, + packages: Optional[Sequence[str]] = None, ): """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. @@ -1467,7 +1468,7 @@ def remote_function( 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 + 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. @@ -1482,6 +1483,10 @@ def remote_function( 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. + packages (str[], Optional): + Explicit name of the external package dependencies. Each dependency + is added to the `requirements.txt` as is, and can be of the form + supported in https://pip.pypa.io/en/stable/reference/requirements-file-format/. 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 @@ -1499,6 +1504,7 @@ def remote_function( bigquery_connection=bigquery_connection, reuse=reuse, name=name, + packages=packages, ) def read_gbq_function( diff --git a/samples/snippets/remote_function.py b/samples/snippets/remote_function.py index 9998a23eb2c..646d7b0c307 100644 --- a/samples/snippets/remote_function.py +++ b/samples/snippets/remote_function.py @@ -89,19 +89,25 @@ def get_bucket(num): # say we consider the `species`, `island` and `sex` of the penguins # sensitive information and want to redact that by replacing with their hash # code instead. Let's define another scalar custom function and decorate it - # as a remote function + # as a remote function. The custom function in this example has external + # package dependency, which can be specified via `packages` parameter. @bpd.remote_function( - [str], str, bigquery_connection="bigframes-rf-conn", reuse=False + [str], + str, + bigquery_connection="bigframes-rf-conn", + reuse=False, + packages=["cryptography"], ) def get_hash(input): - import hashlib + from cryptography.fernet import Fernet # handle missing value if input is None: input = "" - encoded_input = input.encode() - hash = hashlib.md5(encoded_input) - return hash.hexdigest() + + key = Fernet.generate_key() + f = Fernet(key) + return f.encrypt(input.encode()).decode() # We can use this remote function in another `pandas`-like API `map` that # can be applied on a DataFrame diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index f2700991828..730a1dbde4e 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -916,6 +916,51 @@ def square(x): ) +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_with_external_package_dependencies( + session, scalars_dfs, dataset_id, bq_cf_connection, functions_client +): + try: + + def pd_np_foo(x): + import numpy as mynp + import pandas as mypd + + return mypd.Series([x, mynp.sqrt(mynp.abs(x))]).sum() + + # Create the remote function with the name provided explicitly + pd_np_foo_remote = session.remote_function( + [int], + float, + dataset_id, + bq_cf_connection, + reuse=False, + packages=["numpy", "pandas >= 2.0.0"], + )(pd_np_foo) + + # 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(pd_np_foo_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(pd_np_foo) + pd_result = pd_int64_col.to_frame().assign(result=pd_result_col) + + # pandas result is non-nullable type float64, make it Float64 before + # comparing for the purpose of this test + pd_result.result = pd_result.result.astype(pandas.Float64Dtype()) + + 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, pd_np_foo_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 From a6dab9cdb7dd0e56c93ca96b665ab1be1baac5e5 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 17 Oct 2023 17:28:19 -0500 Subject: [PATCH 084/197] docs: add documentation for `Series.struct.field` and `Series.struct.explode` (#114) --- bigframes/core/indexers.py | 2 +- bigframes/operations/structs.py | 4 ++-- docs/reference/bigframes.pandas/series.rst | 8 +++++++ docs/templates/toc.yml | 2 ++ .../bigframes_vendored/pandas/core/series.py | 21 +++++++++++++++++++ 5 files changed, 34 insertions(+), 3 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 01c1162ae6d..f4c4f9011e0 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -225,7 +225,7 @@ def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar: raise ValueError(error_message) if len(key) != 2: raise TypeError(error_message) - block: bigframes.core.blocks.Block = self._dataframe._block + block: bigframes.core.blocks.Block = self._dataframe._block # type: ignore column_block = block.select_columns([block.value_columns[key[1]]]) column = bigframes.series.Series(column_block) return column.iloc[key[0]] diff --git a/bigframes/operations/structs.py b/bigframes/operations/structs.py index 80d51115d01..506a5577094 100644 --- a/bigframes/operations/structs.py +++ b/bigframes/operations/structs.py @@ -25,7 +25,7 @@ import third_party.bigframes_vendored.pandas.core.arrays.arrow.accessors as vendoracessors -class StructField(bigframes.operations.UnaryOp): +class _StructField(bigframes.operations.UnaryOp): def __init__(self, name_or_index: str | int): self._name_or_index = name_or_index @@ -44,7 +44,7 @@ class 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)) + series = self._apply_unary_op(_StructField(name_or_index)) if isinstance(name_or_index, str): name = name_or_index else: diff --git a/docs/reference/bigframes.pandas/series.rst b/docs/reference/bigframes.pandas/series.rst index b179da9ca13..e212904f3ff 100644 --- a/docs/reference/bigframes.pandas/series.rst +++ b/docs/reference/bigframes.pandas/series.rst @@ -34,3 +34,11 @@ String handling :members: :inherited-members: :undoc-members: + +Struct handling +^^^^^^^^^^^^^^^ + +.. automodule:: bigframes.operations.structs + :members: + :inherited-members: + :undoc-members: diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 0758bb41d8a..4fe2ec1a6ad 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -39,6 +39,8 @@ uid: bigframes.operations.datetimes.DatetimeMethods - name: StringMethods uid: bigframes.operations.strings.StringMethods + - name: StructAccessor + uid: bigframes.operations.structs.StructAccessor name: Series - name: Window uid: bigframes.core.window.Window diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index ec3c07bfa85..bd1f9a9a185 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -22,6 +22,23 @@ class Series(NDFrame): # type: ignore[misc] def dt(self): """ Accessor object for datetime-like properties of the Series values. + + Returns: + bigframes.operations.datetimes.DatetimeMethods: + An accessor containing datetime methods. + + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @property + def struct(self): + """ + Accessor object for struct properties of the Series values. + + Returns: + bigframes.operations.structs.StructAccessor: + An accessor containing struct methods. + """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1733,6 +1750,10 @@ 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. + + Returns: + bigframes.operations.strings.StringMethods: + An accessor containing string methods. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 52d64ffdbbab16b1d94974b543ce9080be1ec0d1 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 17 Oct 2023 18:20:17 -0500 Subject: [PATCH 085/197] feat: add `bigframes.options.bigquery.application_name` for partner attribution (#117) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Because `session.py` was getting long, this also refactors `session.py` to separate client construction in a separate module. Fixes internal issue 305950924 🦕 --- bigframes/_config/bigquery_options.py | 19 ++ bigframes/pandas/__init__.py | 4 +- bigframes/{session.py => session/__init__.py} | 170 +-------------- bigframes/session/clients.py | 196 ++++++++++++++++++ tests/unit/_config/test_bigquery_options.py | 4 + tests/unit/resources.py | 3 +- tests/unit/session/__init__.py | 13 ++ tests/unit/session/test_clients.py | 114 ++++++++++ tests/unit/{ => session}/test_session.py | 2 +- 9 files changed, 360 insertions(+), 165 deletions(-) rename bigframes/{session.py => session/__init__.py} (90%) create mode 100644 bigframes/session/clients.py create mode 100644 tests/unit/session/__init__.py create mode 100644 tests/unit/session/test_clients.py rename tests/unit/{ => session}/test_session.py (98%) diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index 453e7f5dff4..d0cce9492bb 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -37,14 +37,33 @@ def __init__( location: Optional[str] = None, bq_connection: Optional[str] = None, use_regional_endpoints: bool = False, + application_name: Optional[str] = None, ): self._credentials = credentials self._project = project self._location = location self._bq_connection = bq_connection self._use_regional_endpoints = use_regional_endpoints + self._application_name = application_name self._session_started = False + @property + def application_name(self) -> Optional[str]: + """The application name to amend to the user-agent sent to Google APIs. + + Recommended format is ``"appplication-name/major.minor.patch_version"`` + or ``"(gpn:PartnerName;)"`` for official Google partners. + """ + return self._application_name + + @application_name.setter + def application_name(self, value: Optional[str]): + if self._session_started and self._application_name != value: + raise ValueError( + SESSION_STARTED_MESSAGE.format(attribute="application_name") + ) + self._application_name = value + @property def credentials(self) -> Optional[google.auth.credentials.Credentials]: """The OAuth2 Credentials to use for this client.""" diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 833f0d1d1dd..1b9144fb625 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -51,6 +51,7 @@ import bigframes.dataframe import bigframes.series import bigframes.session +import bigframes.session.clients 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 @@ -180,11 +181,12 @@ def _set_default_session_location_if_possible(query): ): return - clients_provider = bigframes.session.ClientsProvider( + clients_provider = bigframes.session.clients.ClientsProvider( project=options.bigquery.project, location=options.bigquery.location, use_regional_endpoints=options.bigquery.use_regional_endpoints, credentials=options.bigquery.credentials, + application_name=options.bigquery.application_name, ) bqclient = clients_provider.bqclient diff --git a/bigframes/session.py b/bigframes/session/__init__.py similarity index 90% rename from bigframes/session.py rename to bigframes/session/__init__.py index fa5b4153507..1031fde9b5b 100644 --- a/bigframes/session.py +++ b/bigframes/session/__init__.py @@ -61,7 +61,6 @@ ReadPickleBuffer, StorageOptions, ) -import pydata_google_auth import bigframes._config.bigquery_options as bigquery_options import bigframes.constants as constants @@ -75,6 +74,7 @@ import bigframes.formatting_helpers as formatting_helpers from bigframes.remote_function import read_gbq_function as bigframes_rgf from bigframes.remote_function import remote_function as bigframes_rf +import bigframes.session.clients import bigframes.version # Even though the ibis.backends.bigquery.registry import is unused, it's needed @@ -85,18 +85,6 @@ 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__}" -_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" - _BIGFRAMES_DEFAULT_CONNECTION_ID = "bigframes-default-connection" _MAX_CLUSTER_COLUMNS = 4 @@ -122,149 +110,6 @@ def _is_query(query_or_table: str) -> bool: return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None -def _get_default_credentials_with_project(): - return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False) - - -class ClientsProvider: - """Provides client instances necessary to perform cloud operations.""" - - 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 not project: - raise ValueError( - "Project must be set to initialize BigQuery client. " - "Try setting `bigframes.options.bigquery.project` first." - ) - - self._project = project - self._location = location - self._use_regional_endpoints = use_regional_endpoints - self._credentials = credentials - - # cloud clients initialized for lazy load - self._bqclient = None - self._bqconnectionclient = None - self._bqstorageclient = None - self._cloudfunctionsclient = None - self._resourcemanagerclient = None - - @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 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( third_party_pandas_gbq.GBQIOMixin, third_party_pandas_parquet.ParquetIOMixin, @@ -279,14 +124,14 @@ class Session( 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): + clients_provider (bigframes.session.bigframes.session.clients.ClientsProvider): An object providing client library objects. """ def __init__( self, context: Optional[bigquery_options.BigQueryOptions] = None, - clients_provider: Optional[ClientsProvider] = None, + clients_provider: Optional[bigframes.session.clients.ClientsProvider] = None, ): if context is None: context = bigquery_options.BigQueryOptions() @@ -306,11 +151,12 @@ def __init__( if clients_provider: self._clients_provider = clients_provider else: - self._clients_provider = ClientsProvider( + self._clients_provider = bigframes.session.clients.ClientsProvider( project=context.project, location=self._location, use_regional_endpoints=context.use_regional_endpoints, credentials=context.credentials, + application_name=context.application_name, ) self._create_and_bind_bq_session() @@ -319,7 +165,7 @@ def __init__( ibis.bigquery.connect( project_id=context.project, client=self.bqclient, - storage_client=self.bqstorageclient, + storage_client=self.bqstoragereadclient, ), ) @@ -338,8 +184,8 @@ def bqconnectionclient(self): return self._clients_provider.bqconnectionclient @property - def bqstorageclient(self): - return self._clients_provider.bqstorageclient + def bqstoragereadclient(self): + return self._clients_provider.bqstoragereadclient @property def cloudfunctionsclient(self): diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py new file mode 100644 index 00000000000..544f74265fe --- /dev/null +++ b/bigframes/session/clients.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. + +"""Clients manages the connection to Google APIs.""" + +import os +import typing +from typing import Optional + +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.functions_v2 +import google.cloud.resourcemanager_v3 +import pydata_google_auth + +import bigframes.version + +_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" + + +def _get_default_credentials_with_project(): + return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False) + + +class ClientsProvider: + """Provides client instances necessary to perform cloud operations.""" + + def __init__( + self, + project: Optional[str], + location: Optional[str], + use_regional_endpoints: Optional[bool], + credentials: Optional[google.auth.credentials.Credentials], + application_name: Optional[str], + ): + 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. " + "Try setting `bigframes.options.bigquery.project` first." + ) + + self._application_name = ( + f"{_APPLICATION_NAME} {application_name}" + if application_name + else _APPLICATION_NAME + ) + self._project = project + self._location = location + self._use_regional_endpoints = use_regional_endpoints + self._credentials = credentials + + # cloud clients initialized for lazy load + self._bqclient = None + self._bqconnectionclient = None + self._bqstoragereadclient = None + self._cloudfunctionsclient = None + self._resourcemanagerclient = None + + @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=self._application_name + ) + self._bqclient = bigquery.Client( + client_info=bq_info, + client_options=bq_options, + credentials=self._credentials, + project=self._project, + location=self._location, + ) + + 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=self._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 bqstoragereadclient(self): + if not self._bqstoragereadclient: + 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=self._application_name + ) + self._bqstoragereadclient = ( + google.cloud.bigquery_storage_v1.BigQueryReadClient( + client_info=bqstorage_info, + client_options=bqstorage_options, + credentials=self._credentials, + ) + ) + + return self._bqstoragereadclient + + @property + def cloudfunctionsclient(self): + if not self._cloudfunctionsclient: + functions_info = google.api_core.gapic_v1.client_info.ClientInfo( + user_agent=self._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=self._application_name + ) + self._resourcemanagerclient = ( + google.cloud.resourcemanager_v3.ProjectsClient( + credentials=self._credentials, client_info=resourcemanager_info + ) + ) + + return self._resourcemanagerclient diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index aeee0583195..e5b6cfe2f1b 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -22,11 +22,13 @@ @pytest.mark.parametrize( ["attribute", "original_value", "new_value"], [ + ("application_name", None, "test-partner"), # For credentials, the match is by reference. ("credentials", object(), object()), ("location", "us-east1", "us-central1"), ("project", "my-project", "my-other-project"), ("bq_connection", "path/to/connection/1", "path/to/connection/2"), + ("use_regional_endpoints", False, True), ], ) def test_setter_raises_if_session_started(attribute, original_value, new_value): @@ -53,10 +55,12 @@ def test_setter_raises_if_session_started(attribute, original_value, new_value): [ (attribute,) for attribute in [ + "application_name", "credentials", "location", "project", "bq_connection", + "use_regional_endpoints", ] ], ) diff --git a/tests/unit/resources.py b/tests/unit/resources.py index c8ed6e86ed6..0a68600a359 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -22,6 +22,7 @@ import bigframes import bigframes.core as core +import bigframes.session.clients """Utilities for creating test resources.""" @@ -37,7 +38,7 @@ def create_bigquery_session( bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) bqclient.project = "test-project" - clients_provider = mock.create_autospec(bigframes.session.ClientsProvider) + clients_provider = mock.create_autospec(bigframes.session.clients.ClientsProvider) type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient) clients_provider._credentials = credentials diff --git a/tests/unit/session/__init__.py b/tests/unit/session/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/tests/unit/session/__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/session/test_clients.py b/tests/unit/session/test_clients.py new file mode 100644 index 00000000000..f1b2a5045a3 --- /dev/null +++ b/tests/unit/session/test_clients.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. + +from typing import Optional +import unittest.mock as mock + +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 +import google.cloud.bigquery_connection_v1 +import google.cloud.bigquery_storage_v1 +import google.cloud.functions_v2 +import google.cloud.resourcemanager_v3 + +import bigframes.session.clients as clients +import bigframes.version + + +def create_clients_provider(application_name: Optional[str] = None): + credentials = mock.create_autospec(google.auth.credentials.Credentials) + return clients.ClientsProvider( + project="test-project", + location="test-region", + use_regional_endpoints=False, + credentials=credentials, + application_name=application_name, + ) + + +def monkeypatch_client_constructors(monkeypatch): + bqclient = mock.create_autospec(google.cloud.bigquery.Client) + bqclient.return_value = bqclient + monkeypatch.setattr(google.cloud.bigquery, "Client", bqclient) + + bqconnectionclient = mock.create_autospec( + google.cloud.bigquery_connection_v1.ConnectionServiceClient + ) + bqconnectionclient.return_value = bqconnectionclient + monkeypatch.setattr( + google.cloud.bigquery_connection_v1, + "ConnectionServiceClient", + bqconnectionclient, + ) + + bqstoragereadclient = mock.create_autospec( + google.cloud.bigquery_storage_v1.BigQueryReadClient + ) + bqstoragereadclient.return_value = bqstoragereadclient + monkeypatch.setattr( + google.cloud.bigquery_storage_v1, "BigQueryReadClient", bqstoragereadclient + ) + + cloudfunctionsclient = mock.create_autospec( + google.cloud.functions_v2.FunctionServiceClient + ) + cloudfunctionsclient.return_value = cloudfunctionsclient + monkeypatch.setattr( + google.cloud.functions_v2, "FunctionServiceClient", cloudfunctionsclient + ) + + resourcemanagerclient = mock.create_autospec( + google.cloud.resourcemanager_v3.ProjectsClient + ) + resourcemanagerclient.return_value = resourcemanagerclient + monkeypatch.setattr( + google.cloud.resourcemanager_v3, "ProjectsClient", resourcemanagerclient + ) + + +def assert_constructed_w_user_agent(mock_client: mock.Mock, expected_user_agent: str): + assert ( + expected_user_agent + in mock_client.call_args.kwargs["client_info"].to_user_agent() + ) + + +def assert_clients_w_user_agent( + provider: clients.ClientsProvider, expected_user_agent: str +): + assert_constructed_w_user_agent(provider.bqclient, expected_user_agent) + assert_constructed_w_user_agent(provider.bqconnectionclient, expected_user_agent) + assert_constructed_w_user_agent(provider.bqstoragereadclient, expected_user_agent) + assert_constructed_w_user_agent(provider.cloudfunctionsclient, expected_user_agent) + assert_constructed_w_user_agent(provider.resourcemanagerclient, expected_user_agent) + + +def test_user_agent_default(monkeypatch): + monkeypatch_client_constructors(monkeypatch) + provider = create_clients_provider(application_name=None) + assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") + + +def test_user_agent_custom(monkeypatch): + monkeypatch_client_constructors(monkeypatch) + provider = create_clients_provider(application_name="(gpn:testpartner;)") + assert_clients_w_user_agent(provider, "(gpn:testpartner;)") + + # We still need to include attribution to bigframes, even if there's also a + # partner using the package. + assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") diff --git a/tests/unit/test_session.py b/tests/unit/session/test_session.py similarity index 98% rename from tests/unit/test_session.py rename to tests/unit/session/test_session.py index e39a316e5b6..18fd42e0f31 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/session/test_session.py @@ -20,7 +20,7 @@ import bigframes -from . import resources +from .. import resources @pytest.mark.parametrize("missing_parts_table_id", [(""), ("table")]) From 68e3cd37258084d045ea1075e5e61df12c28faac Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 17 Oct 2023 18:36:28 -0700 Subject: [PATCH 086/197] fix: fix TODOs for loc multiindex (#113) Co-authored-by: Henry J Solberg --- bigframes/core/indexers.py | 59 ++++++++++++++++++---------- tests/system/small/test_dataframe.py | 34 ++++++++++++++++ tests/system/small/test_series.py | 34 ++++++++++++++++ 3 files changed, 107 insertions(+), 20 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index f4c4f9011e0..4f5a9471b99 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, Union +from typing import List, Tuple, Union import ibis import pandas as pd @@ -271,40 +271,59 @@ def _loc_getitem_series_or_dataframe( 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_") + if len(series_or_dataframe.index.names) > 1: + temp_name = series_or_dataframe.index.names[0] 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 + key = typing.cast(List, key) + if len(key) == 0: return typing.cast( 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 + if pd.api.types.is_list_like(key[0]): + original_index_names = series_or_dataframe.index.names + num_index_cols = len(original_index_names) + + entry_col_count_correct = [len(entry) == num_index_cols for entry in key] + if not all(entry_col_count_correct): + # pandas usually throws TypeError in these cases- tuple causes IndexError, but that + # seems like unintended behavior + raise TypeError( + "All entries must be of equal length when indexing by list of listlikes" + ) + temporary_index_names = [ + guid.generate_guid(prefix="temp_loc_index_") + for _ in range(len(original_index_names)) + ] + index_cols_dict = {} + for i in range(num_index_cols): + index_name = temporary_index_names[i] + values = [entry[i] for entry in key] + index_cols_dict[index_name] = values + keys_df = bigframes.dataframe.DataFrame(index_cols_dict) + keys_df = keys_df.set_index(temporary_index_names, drop=True) + keys_df = keys_df.rename_axis(original_index_names) + else: + # 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}) + 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): if (key.start is None) and (key.stop is None) and (key.step is None): diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index ecafb7c1bf5..309e8df4f0d 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2787,6 +2787,22 @@ def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_list_multiindex(scalars_df_index, scalars_pandas_df_index): + scalars_df_multiindex = scalars_df_index.set_index(["string_col", "int64_col"]) + scalars_pandas_df_multiindex = scalars_pandas_df_index.set_index( + ["string_col", "int64_col"] + ) + index_list = [("Hello, World!", -234892), ("Hello, World!", 123456789)] + + bf_result = scalars_df_multiindex.loc[index_list] + pd_result = scalars_pandas_df_multiindex.loc[index_list] + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + def test_iloc_list(scalars_df_index, scalars_pandas_df_index): index_list = [0, 0, 0, 5, 4, 7] @@ -2863,6 +2879,24 @@ def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_bf_series_multiindex(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_multiindex = scalars_df_index.set_index(["string_col", "int64_col"]) + scalars_pandas_df_multiindex = scalars_pandas_df_index.set_index( + ["string_col", "int64_col"] + ) + + bf_result = scalars_df_multiindex.loc[bf_string_series] + pd_result = scalars_pandas_df_multiindex.loc[pd_string_series] + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + 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 diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index f2ced841daf..bd9edbb1ca0 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2624,6 +2624,22 @@ def test_loc_list_integer_index(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_list_multiindex(scalars_df_index, scalars_pandas_df_index): + scalars_df_multiindex = scalars_df_index.set_index(["string_col", "int64_col"]) + scalars_pandas_df_multiindex = scalars_pandas_df_index.set_index( + ["string_col", "int64_col"] + ) + index_list = [("Hello, World!", -234892), ("Hello, World!", 123456789)] + + bf_result = scalars_df_multiindex.int64_too.loc[index_list] + pd_result = scalars_pandas_df_multiindex.int64_too.loc[index_list] + + pd.testing.assert_series_equal( + bf_result.to_pandas(), + pd_result, + ) + + def test_iloc_list(scalars_df_index, scalars_pandas_df_index): index_list = [0, 0, 0, 5, 4, 7] @@ -2681,6 +2697,24 @@ def test_loc_bf_series_string_index(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_bf_series_multiindex(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_multiindex = scalars_df_index.set_index(["string_col", "int64_col"]) + scalars_pandas_df_multiindex = scalars_pandas_df_index.set_index( + ["string_col", "int64_col"] + ) + + bf_result = scalars_df_multiindex.int64_too.loc[bf_string_series] + pd_result = scalars_pandas_df_multiindex.int64_too.loc[pd_string_series] + + pd.testing.assert_series_equal( + bf_result.to_pandas(), + 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 From 9238fadcfa7e843be6564813ff3131893b79f8b0 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 18 Oct 2023 11:48:18 -0500 Subject: [PATCH 087/197] feat: use ArrowDtype for STRUCT columns in `to_pandas` (#85) 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/core/blocks.py | 2 + bigframes/dtypes.py | 17 + noxfile.py | 2 + tests/system/small/test_dataframe_io.py | 23 +- .../google_cloud_bigquery/LICENSE | 202 +++++++++ .../google_cloud_bigquery/__init__.py | 13 + .../google_cloud_bigquery/_pandas_helpers.py | 158 +++++++ .../google_cloud_bigquery/tests/__init__.py | 13 + .../tests/unit/__init__.py | 13 + .../tests/unit/test_pandas_helpers.py | 413 ++++++++++++++++++ 10 files changed, 853 insertions(+), 3 deletions(-) create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/LICENSE create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/__init__.py create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/_pandas_helpers.py create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/tests/__init__.py create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/__init__.py create mode 100644 third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/test_pandas_helpers.py diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 4548fca5934..8966b6189bd 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -377,7 +377,9 @@ def _to_dataframe( cls, result, schema: typing.Mapping[str, bigframes.dtypes.Dtype] ) -> pd.DataFrame: """Convert BigQuery data to pandas DataFrame with specific dtypes.""" + dtypes = bigframes.dtypes.to_pandas_dtypes_overrides(result.schema) df = result.to_dataframe( + dtypes=dtypes, bool_dtype=pd.BooleanDtype(), int_dtype=pd.Int64Dtype(), float_dtype=pd.Float64Dtype(), diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 46a7a1cb503..da221a95aca 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -19,6 +19,7 @@ from typing import Any, Dict, Iterable, Literal, Tuple, Union import geopandas as gpd # type: ignore +import google.cloud.bigquery as bigquery import ibis import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types @@ -27,6 +28,7 @@ import pyarrow as pa import bigframes.constants as constants +import third_party.bigframes_vendored.google_cloud_bigquery._pandas_helpers as gcb3p_pandas_helpers # Type hints for Pandas dtypes supported by BigQuery DataFrame Dtype = Union[ @@ -401,3 +403,18 @@ def cast_ibis_value( raise TypeError( f"Unsupported cast {value.type()} to {to_type}. {constants.FEEDBACK_LINK}" ) + + +def to_pandas_dtypes_overrides(schema: Iterable[bigquery.SchemaField]) -> Dict: + """For each STRUCT field, make sure we specify the full type to use.""" + # TODO(swast): Also override ARRAY fields. + dtypes = {} + for field in schema: + if field.field_type == "RECORD" and field.mode != "REPEATED": + # TODO(swast): We're using a private API here. Would likely be + # better if we called `to_arrow()` and converted to a pandas + # DataFrame ourselves from that. + dtypes[field.name] = pd.ArrowDtype( + gcb3p_pandas_helpers.bq_to_arrow_data_type(field) + ) + return dtypes diff --git a/noxfile.py b/noxfile.py index 54ccdb9a877..1864da9fe7d 100644 --- a/noxfile.py +++ b/noxfile.py @@ -185,6 +185,7 @@ def run_unit(session, install_test_extra): # Run py.test against the unit tests. tests_path = os.path.join("tests", "unit") + third_party_tests_path = os.path.join("third_party", "bigframes_vendored") session.run( "py.test", "--quiet", @@ -196,6 +197,7 @@ def run_unit(session, install_test_extra): "--cov-report=term-missing", "--cov-fail-under=0", tests_path, + third_party_tests_path, *session.posargs, ) diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index 3886b85f40c..d60083a837e 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -16,6 +16,7 @@ import google.api_core.exceptions import pandas as pd +import pyarrow as pa import pytest from tests.system.utils import ( @@ -44,7 +45,7 @@ def test_to_pandas_w_correct_dtypes(scalars_df_default_index): def test_to_pandas_array_struct_correct_result(session): - """In future, we should support arrays and structs with arrow types. + """In future, we should support arrays with arrow types. For now we fall back to the current connector behavior of converting to Python objects""" df = session.read_gbq( @@ -59,11 +60,27 @@ def test_to_pandas_array_struct_correct_result(session): expected = pd.DataFrame( { "array_column": [[1, 3, 2]], - "struct_column": [{"string_field": "a", "float_field": 1.2}], + "struct_column": pd.Series( + [{"string_field": "a", "float_field": 1.2}], + dtype=pd.ArrowDtype( + pa.struct( + [ + ("string_field", pa.string()), + ("float_field", pa.float64()), + ] + ) + ), + ), } ) expected.index = expected.index.astype("Int64") - pd.testing.assert_frame_equal(result, expected) + pd.testing.assert_series_equal(result.dtypes, expected.dtypes) + pd.testing.assert_series_equal(result["array_column"], expected["array_column"]) + # assert_series_equal not implemented for struct columns yet. Compare + # values as Python objects, instead. + pd.testing.assert_series_equal( + result["struct_column"].astype("O"), expected["struct_column"].astype("O") + ) @pytest.mark.parametrize( diff --git a/third_party/bigframes_vendored/google_cloud_bigquery/LICENSE b/third_party/bigframes_vendored/google_cloud_bigquery/LICENSE new file mode 100644 index 00000000000..d6456956733 --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/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. diff --git a/third_party/bigframes_vendored/google_cloud_bigquery/__init__.py b/third_party/bigframes_vendored/google_cloud_bigquery/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/__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/third_party/bigframes_vendored/google_cloud_bigquery/_pandas_helpers.py b/third_party/bigframes_vendored/google_cloud_bigquery/_pandas_helpers.py new file mode 100644 index 00000000000..5e2a7a7ef0f --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/_pandas_helpers.py @@ -0,0 +1,158 @@ +# Original: https://github.com/googleapis/python-bigquery/blob/main/google/cloud/bigquery/_pandas_helpers.py +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helper functions for connecting BigQuery and pandas.""" + +import warnings + +import google.cloud.bigquery.schema as schema +import pyarrow + + +def pyarrow_datetime(): + return pyarrow.timestamp("us", tz=None) + + +def pyarrow_numeric(): + return pyarrow.decimal128(38, 9) + + +def pyarrow_bignumeric(): + # 77th digit is partial. + # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#decimal_types + return pyarrow.decimal256(76, 38) + + +def pyarrow_time(): + return pyarrow.time64("us") + + +def pyarrow_timestamp(): + return pyarrow.timestamp("us", tz="UTC") + + +# This dictionary is duplicated in bigquery_storage/test/unite/test_reader.py +# When modifying it be sure to update it there as well. +BQ_TO_ARROW_SCALARS = { + "BOOL": pyarrow.bool_, + "BOOLEAN": pyarrow.bool_, + "BYTES": pyarrow.binary, + "DATE": pyarrow.date32, + "DATETIME": pyarrow_datetime, + "FLOAT": pyarrow.float64, + "FLOAT64": pyarrow.float64, + "GEOGRAPHY": pyarrow.string, + "INT64": pyarrow.int64, + "INTEGER": pyarrow.int64, + "NUMERIC": pyarrow_numeric, + "STRING": pyarrow.string, + "TIME": pyarrow_time, + "TIMESTAMP": pyarrow_timestamp, + "BIGNUMERIC": pyarrow_bignumeric, +} +ARROW_SCALAR_IDS_TO_BQ = { + # https://arrow.apache.org/docs/python/api/datatypes.html#type-classes + pyarrow.bool_().id: "BOOL", + pyarrow.int8().id: "INT64", + pyarrow.int16().id: "INT64", + pyarrow.int32().id: "INT64", + pyarrow.int64().id: "INT64", + pyarrow.uint8().id: "INT64", + pyarrow.uint16().id: "INT64", + pyarrow.uint32().id: "INT64", + pyarrow.uint64().id: "INT64", + pyarrow.float16().id: "FLOAT64", + pyarrow.float32().id: "FLOAT64", + pyarrow.float64().id: "FLOAT64", + pyarrow.time32("ms").id: "TIME", + pyarrow.time64("ns").id: "TIME", + pyarrow.timestamp("ns").id: "TIMESTAMP", + pyarrow.date32().id: "DATE", + pyarrow.date64().id: "DATETIME", # because millisecond resolution + pyarrow.binary().id: "BYTES", + pyarrow.string().id: "STRING", # also alias for pyarrow.utf8() + # The exact scale and precision don't matter. Only the type ID matters, + # and it's the same for all decimal128/decimal256 instances. + pyarrow.decimal128(38, scale=9).id: "NUMERIC", + pyarrow.decimal256(76, scale=38).id: "BIGNUMERIC", +} + + +BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA = { + "GEOGRAPHY": { + b"ARROW:extension:name": b"google:sqlType:geography", + b"ARROW:extension:metadata": b'{"encoding": "WKT"}', + }, + "DATETIME": {b"ARROW:extension:name": b"google:sqlType:datetime"}, +} + + +def bq_to_arrow_struct_data_type(field): + arrow_fields = [] + for subfield in field.fields: + arrow_subfield = bq_to_arrow_field(subfield) + if arrow_subfield: + arrow_fields.append(arrow_subfield) + else: + # Could not determine a subfield type. Fallback to type + # inference. + return None + return pyarrow.struct(arrow_fields) + + +def bq_to_arrow_data_type(field): + """Return the Arrow data type, corresponding to a given BigQuery column. + + Returns: + None: if default Arrow type inspection should be used. + """ + if field.mode is not None and field.mode.upper() == "REPEATED": + inner_type = bq_to_arrow_data_type( + schema.SchemaField(field.name, field.field_type, fields=field.fields) + ) + if inner_type: + return pyarrow.list_(inner_type) + return None + + field_type_upper = field.field_type.upper() if field.field_type else "" + if field_type_upper in schema._STRUCT_TYPES: + return bq_to_arrow_struct_data_type(field) + + data_type_constructor = BQ_TO_ARROW_SCALARS.get(field_type_upper) + if data_type_constructor is None: + return None + return data_type_constructor() + + +def bq_to_arrow_field(bq_field, array_type=None): + """Return the Arrow field, corresponding to a given BigQuery column. + + Returns: + None: if the Arrow type cannot be determined. + """ + arrow_type = bq_to_arrow_data_type(bq_field) + if arrow_type is not None: + if array_type is not None: + arrow_type = array_type # For GEOGRAPHY, at least initially + is_nullable = bq_field.mode.upper() == "NULLABLE" + metadata = BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA.get( + bq_field.field_type.upper() if bq_field.field_type else "" + ) + return pyarrow.field( + bq_field.name, arrow_type, nullable=is_nullable, metadata=metadata + ) + + warnings.warn("Unable to determine type for field '{}'.".format(bq_field.name)) + return None diff --git a/third_party/bigframes_vendored/google_cloud_bigquery/tests/__init__.py b/third_party/bigframes_vendored/google_cloud_bigquery/tests/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/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/third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/__init__.py b/third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/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/third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/test_pandas_helpers.py b/third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/test_pandas_helpers.py new file mode 100644 index 00000000000..dc4a09cc541 --- /dev/null +++ b/third_party/bigframes_vendored/google_cloud_bigquery/tests/unit/test_pandas_helpers.py @@ -0,0 +1,413 @@ +# Original: https://github.com/googleapis/python-bigquery/blob/main/tests/unit/test__pandas_helpers.py +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +import warnings + +from google.cloud.bigquery import schema +import pyarrow +import pyarrow.parquet +import pyarrow.types +import pytest + + +@pytest.fixture +def module_under_test(): + from third_party.bigframes_vendored.google_cloud_bigquery import _pandas_helpers + + return _pandas_helpers + + +def is_none(value): + return value is None + + +def is_datetime(type_): + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#datetime-type + return all_( + pyarrow.types.is_timestamp, + lambda type_: type_.unit == "us", + lambda type_: type_.tz is None, + )(type_) + + +def is_numeric(type_): + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#numeric-type + return all_( + pyarrow.types.is_decimal, + lambda type_: type_.precision == 38, + lambda type_: type_.scale == 9, + )(type_) + + +def is_bignumeric(type_): + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#numeric-type + return all_( + pyarrow.types.is_decimal, + lambda type_: type_.precision == 76, + lambda type_: type_.scale == 38, + )(type_) + + +def is_timestamp(type_): + # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#timestamp-type + return all_( + pyarrow.types.is_timestamp, + lambda type_: type_.unit == "us", + lambda type_: type_.tz == "UTC", + )(type_) + + +def do_all(functions, value): + return all((func(value) for func in functions)) + + +def all_(*functions): + return functools.partial(do_all, functions) + + +def test_is_datetime(): + assert is_datetime(pyarrow.timestamp("us", tz=None)) + assert not is_datetime(pyarrow.timestamp("ms", tz=None)) + assert not is_datetime(pyarrow.timestamp("us", tz="UTC")) + assert not is_datetime(pyarrow.timestamp("ns", tz="UTC")) + assert not is_datetime(pyarrow.string()) + + +def test_do_all(): + assert do_all((lambda _: True, lambda _: True), None) + assert not do_all((lambda _: True, lambda _: False), None) + assert not do_all((lambda _: False,), None) + + +def test_all_(): + assert all_(lambda _: True, lambda _: True)(None) + assert not all_(lambda _: True, lambda _: False)(None) + + +@pytest.mark.parametrize( + "bq_type,bq_mode,is_correct_type", + [ + ("STRING", "NULLABLE", pyarrow.types.is_string), + ("STRING", None, pyarrow.types.is_string), + ("string", "NULLABLE", pyarrow.types.is_string), + ("StRiNg", "NULLABLE", pyarrow.types.is_string), + ("BYTES", "NULLABLE", pyarrow.types.is_binary), + ("INTEGER", "NULLABLE", pyarrow.types.is_int64), + ("INT64", "NULLABLE", pyarrow.types.is_int64), + ("FLOAT", "NULLABLE", pyarrow.types.is_float64), + ("FLOAT64", "NULLABLE", pyarrow.types.is_float64), + ("NUMERIC", "NULLABLE", is_numeric), + pytest.param( + "BIGNUMERIC", + "NULLABLE", + is_bignumeric, + ), + ("BOOLEAN", "NULLABLE", pyarrow.types.is_boolean), + ("BOOL", "NULLABLE", pyarrow.types.is_boolean), + ("TIMESTAMP", "NULLABLE", is_timestamp), + ("DATE", "NULLABLE", pyarrow.types.is_date32), + ("TIME", "NULLABLE", pyarrow.types.is_time64), + ("DATETIME", "NULLABLE", is_datetime), + ("GEOGRAPHY", "NULLABLE", pyarrow.types.is_string), + ("UNKNOWN_TYPE", "NULLABLE", is_none), + # Use pyarrow.list_(item_type) for repeated (array) fields. + ( + "STRING", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_string(type_.value_type), + ), + ), + ( + "STRING", + "repeated", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_string(type_.value_type), + ), + ), + ( + "STRING", + "RePeAtEd", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_string(type_.value_type), + ), + ), + ( + "BYTES", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_binary(type_.value_type), + ), + ), + ( + "INTEGER", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_int64(type_.value_type), + ), + ), + ( + "INT64", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_int64(type_.value_type), + ), + ), + ( + "FLOAT", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_float64(type_.value_type), + ), + ), + ( + "FLOAT64", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_float64(type_.value_type), + ), + ), + ( + "NUMERIC", + "REPEATED", + all_(pyarrow.types.is_list, lambda type_: is_numeric(type_.value_type)), + ), + pytest.param( + "BIGNUMERIC", + "REPEATED", + all_(pyarrow.types.is_list, lambda type_: is_bignumeric(type_.value_type)), + ), + ( + "BOOLEAN", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_boolean(type_.value_type), + ), + ), + ( + "BOOL", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_boolean(type_.value_type), + ), + ), + ( + "TIMESTAMP", + "REPEATED", + all_(pyarrow.types.is_list, lambda type_: is_timestamp(type_.value_type)), + ), + ( + "DATE", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_date32(type_.value_type), + ), + ), + ( + "TIME", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_time64(type_.value_type), + ), + ), + ( + "DATETIME", + "REPEATED", + all_(pyarrow.types.is_list, lambda type_: is_datetime(type_.value_type)), + ), + ( + "GEOGRAPHY", + "REPEATED", + all_( + pyarrow.types.is_list, + lambda type_: pyarrow.types.is_string(type_.value_type), + ), + ), + ("RECORD", "REPEATED", is_none), + ("UNKNOWN_TYPE", "REPEATED", is_none), + ], +) +def test_bq_to_arrow_data_type(module_under_test, bq_type, bq_mode, is_correct_type): + field = schema.SchemaField("ignored_name", bq_type, mode=bq_mode) + actual = module_under_test.bq_to_arrow_data_type(field) + assert is_correct_type(actual) + + +@pytest.mark.parametrize("bq_type", ["RECORD", "record", "STRUCT", "struct"]) +def test_bq_to_arrow_data_type_w_struct(module_under_test, bq_type): + fields = ( + schema.SchemaField("field01", "STRING"), + schema.SchemaField("field02", "BYTES"), + schema.SchemaField("field03", "INTEGER"), + schema.SchemaField("field04", "INT64"), + schema.SchemaField("field05", "FLOAT"), + schema.SchemaField("field06", "FLOAT64"), + schema.SchemaField("field07", "NUMERIC"), + schema.SchemaField("field08", "BIGNUMERIC"), + schema.SchemaField("field09", "BOOLEAN"), + schema.SchemaField("field10", "BOOL"), + schema.SchemaField("field11", "TIMESTAMP"), + schema.SchemaField("field12", "DATE"), + schema.SchemaField("field13", "TIME"), + schema.SchemaField("field14", "DATETIME"), + schema.SchemaField("field15", "GEOGRAPHY"), + ) + + field = schema.SchemaField("ignored_name", bq_type, mode="NULLABLE", fields=fields) + actual = module_under_test.bq_to_arrow_data_type(field) + + expected = ( + pyarrow.field("field01", pyarrow.string()), + pyarrow.field("field02", pyarrow.binary()), + pyarrow.field("field03", pyarrow.int64()), + pyarrow.field("field04", pyarrow.int64()), + pyarrow.field("field05", pyarrow.float64()), + pyarrow.field("field06", pyarrow.float64()), + pyarrow.field("field07", module_under_test.pyarrow_numeric()), + pyarrow.field("field08", module_under_test.pyarrow_bignumeric()), + pyarrow.field("field09", pyarrow.bool_()), + pyarrow.field("field10", pyarrow.bool_()), + pyarrow.field("field11", module_under_test.pyarrow_timestamp()), + pyarrow.field("field12", pyarrow.date32()), + pyarrow.field("field13", module_under_test.pyarrow_time()), + pyarrow.field("field14", module_under_test.pyarrow_datetime()), + pyarrow.field("field15", pyarrow.string()), + ) + expected = pyarrow.struct(expected) + + assert pyarrow.types.is_struct(actual) + assert actual.num_fields == len(fields) + assert actual.equals(expected) + + +@pytest.mark.parametrize("bq_type", ["RECORD", "record", "STRUCT", "struct"]) +def test_bq_to_arrow_data_type_w_array_struct(module_under_test, bq_type): + fields = ( + schema.SchemaField("field01", "STRING"), + schema.SchemaField("field02", "BYTES"), + schema.SchemaField("field03", "INTEGER"), + schema.SchemaField("field04", "INT64"), + schema.SchemaField("field05", "FLOAT"), + schema.SchemaField("field06", "FLOAT64"), + schema.SchemaField("field07", "NUMERIC"), + schema.SchemaField("field08", "BIGNUMERIC"), + schema.SchemaField("field09", "BOOLEAN"), + schema.SchemaField("field10", "BOOL"), + schema.SchemaField("field11", "TIMESTAMP"), + schema.SchemaField("field12", "DATE"), + schema.SchemaField("field13", "TIME"), + schema.SchemaField("field14", "DATETIME"), + schema.SchemaField("field15", "GEOGRAPHY"), + ) + + field = schema.SchemaField("ignored_name", bq_type, mode="REPEATED", fields=fields) + actual = module_under_test.bq_to_arrow_data_type(field) + + expected = ( + pyarrow.field("field01", pyarrow.string()), + pyarrow.field("field02", pyarrow.binary()), + pyarrow.field("field03", pyarrow.int64()), + pyarrow.field("field04", pyarrow.int64()), + pyarrow.field("field05", pyarrow.float64()), + pyarrow.field("field06", pyarrow.float64()), + pyarrow.field("field07", module_under_test.pyarrow_numeric()), + pyarrow.field("field08", module_under_test.pyarrow_bignumeric()), + pyarrow.field("field09", pyarrow.bool_()), + pyarrow.field("field10", pyarrow.bool_()), + pyarrow.field("field11", module_under_test.pyarrow_timestamp()), + pyarrow.field("field12", pyarrow.date32()), + pyarrow.field("field13", module_under_test.pyarrow_time()), + pyarrow.field("field14", module_under_test.pyarrow_datetime()), + pyarrow.field("field15", pyarrow.string()), + ) + expected_value_type = pyarrow.struct(expected) + + assert pyarrow.types.is_list(actual) + assert pyarrow.types.is_struct(actual.value_type) + assert actual.value_type.num_fields == len(fields) + assert actual.value_type.equals(expected_value_type) + + +def test_bq_to_arrow_data_type_w_struct_unknown_subfield(module_under_test): + fields = ( + schema.SchemaField("field1", "STRING"), + schema.SchemaField("field2", "INTEGER"), + # Don't know what to convert UNKNOWN_TYPE to, let type inference work, + # instead. + schema.SchemaField("field3", "UNKNOWN_TYPE"), + ) + field = schema.SchemaField("ignored_name", "RECORD", mode="NULLABLE", fields=fields) + + with warnings.catch_warnings(record=True) as warned: + actual = module_under_test.bq_to_arrow_data_type(field) + + assert actual is None + assert len(warned) == 1 + warning = warned[0] + assert "field3" in str(warning) + + +def test_bq_to_arrow_field_type_override(module_under_test): + # When loading pandas data, we may need to override the type + # decision based on data contents, because GEOGRAPHY data can be + # stored as either text or binary. + + assert ( + module_under_test.bq_to_arrow_field(schema.SchemaField("g", "GEOGRAPHY")).type + == pyarrow.string() + ) + + assert ( + module_under_test.bq_to_arrow_field( + schema.SchemaField("g", "GEOGRAPHY"), + pyarrow.binary(), + ).type + == pyarrow.binary() + ) + + +@pytest.mark.parametrize( + "field_type, metadata", + [ + ("datetime", {b"ARROW:extension:name": b"google:sqlType:datetime"}), + ( + "geography", + { + b"ARROW:extension:name": b"google:sqlType:geography", + b"ARROW:extension:metadata": b'{"encoding": "WKT"}', + }, + ), + ], +) +def test_bq_to_arrow_field_metadata(module_under_test, field_type, metadata): + assert ( + module_under_test.bq_to_arrow_field( + schema.SchemaField("g", field_type) + ).metadata + == metadata + ) From e2788a8c9e53516d467b4d5b15da8838540f371e Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 18 Oct 2023 17:38:14 +0000 Subject: [PATCH 088/197] chore(main): release 0.9.0 (#108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit :robot: I have created a release *beep* *boop* --- ## [0.9.0](https://togithub.com/googleapis/python-bigquery-dataframes/compare/v0.8.0...v0.9.0) (2023-10-18) ### ⚠ BREAKING CHANGES * rename `bigframes.pandas.reset_session` to `close_session` ([#101](https://togithub.com/googleapis/python-bigquery-dataframes/issues/101)) ### Features * Add `bigframes.options.bigquery.application_name` for partner attribution ([#117](https://togithub.com/googleapis/python-bigquery-dataframes/issues/117)) ([52d64ff](https://togithub.com/googleapis/python-bigquery-dataframes/commit/52d64ffdbbab16b1d94974b543ce9080be1ec0d1)) * Add AtIndexer getitems ([#107](https://togithub.com/googleapis/python-bigquery-dataframes/issues/107)) ([752b01f](https://togithub.com/googleapis/python-bigquery-dataframes/commit/752b01ff9df114c54ed58eb96956e9ce34a8ed47)) * Rename `bigframes.pandas.reset_session` to `close_session` ([#101](https://togithub.com/googleapis/python-bigquery-dataframes/issues/101)) ([36693bf](https://togithub.com/googleapis/python-bigquery-dataframes/commit/36693bff398c23e179d9bde95d52cbaddaf85c45)) * Send BigQuery cancel request when canceling bigframes process ([#103](https://togithub.com/googleapis/python-bigquery-dataframes/issues/103)) ([e325fbb](https://togithub.com/googleapis/python-bigquery-dataframes/commit/e325fbb1c91e040d87df10f7d4d5ce53f7c052cb)) * Support external packages in `remote_function` ([#98](https://togithub.com/googleapis/python-bigquery-dataframes/issues/98)) ([ec10c4a](https://togithub.com/googleapis/python-bigquery-dataframes/commit/ec10c4a5a7833c42e28fe9e7b734bc0c4fb84b6e)) * Use ArrowDtype for STRUCT columns in `to_pandas` ([#85](https://togithub.com/googleapis/python-bigquery-dataframes/issues/85)) ([9238fad](https://togithub.com/googleapis/python-bigquery-dataframes/commit/9238fadcfa7e843be6564813ff3131893b79f8b0)) ### Bug Fixes * Support multiindex for three loc getitem overloads ([#113](https://togithub.com/googleapis/python-bigquery-dataframes/issues/113)) ([68e3cd3](https://togithub.com/googleapis/python-bigquery-dataframes/commit/68e3cd37258084d045ea1075e5e61df12c28faac)) ### Performance Improvements * If primary keys are defined, `read_gbq` avoids copying table data ([#112](https://togithub.com/googleapis/python-bigquery-dataframes/issues/112)) ([e6c0cd1](https://togithub.com/googleapis/python-bigquery-dataframes/commit/e6c0cd1777736e0fa7285da59625fbac487573bd)) ### Documentation * Add documentation for `Series.struct.field` and `Series.struct.explode` ([#114](https://togithub.com/googleapis/python-bigquery-dataframes/issues/114)) ([a6dab9c](https://togithub.com/googleapis/python-bigquery-dataframes/commit/a6dab9cdb7dd0e56c93ca96b665ab1be1baac5e5)) * Add open-source link in API doc ([#106](https://togithub.com/googleapis/python-bigquery-dataframes/issues/106)) ([db51fe3](https://togithub.com/googleapis/python-bigquery-dataframes/commit/db51fe340f644a0d7c911c11d92c8299a4be3446)) * Update ML overview API doc ([#105](https://togithub.com/googleapis/python-bigquery-dataframes/issues/105)) ([1b3f3a5](https://togithub.com/googleapis/python-bigquery-dataframes/commit/1b3f3a5374915b2833c6c1ac05670e9708f07bff)) --- 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 | 33 +++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e34b073999a..974e5d3d43d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,39 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.9.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.8.0...v0.9.0) (2023-10-18) + + +### ⚠ BREAKING CHANGES + +* rename `bigframes.pandas.reset_session` to `close_session` ([#101](https://github.com/googleapis/python-bigquery-dataframes/issues/101)) + +### Features + +* Add `bigframes.options.bigquery.application_name` for partner attribution ([#117](https://github.com/googleapis/python-bigquery-dataframes/issues/117)) ([52d64ff](https://github.com/googleapis/python-bigquery-dataframes/commit/52d64ffdbbab16b1d94974b543ce9080be1ec0d1)) +* Add AtIndexer getitems ([#107](https://github.com/googleapis/python-bigquery-dataframes/issues/107)) ([752b01f](https://github.com/googleapis/python-bigquery-dataframes/commit/752b01ff9df114c54ed58eb96956e9ce34a8ed47)) +* Rename `bigframes.pandas.reset_session` to `close_session` ([#101](https://github.com/googleapis/python-bigquery-dataframes/issues/101)) ([36693bf](https://github.com/googleapis/python-bigquery-dataframes/commit/36693bff398c23e179d9bde95d52cbaddaf85c45)) +* Send BigQuery cancel request when canceling bigframes process ([#103](https://github.com/googleapis/python-bigquery-dataframes/issues/103)) ([e325fbb](https://github.com/googleapis/python-bigquery-dataframes/commit/e325fbb1c91e040d87df10f7d4d5ce53f7c052cb)) +* Support external packages in `remote_function` ([#98](https://github.com/googleapis/python-bigquery-dataframes/issues/98)) ([ec10c4a](https://github.com/googleapis/python-bigquery-dataframes/commit/ec10c4a5a7833c42e28fe9e7b734bc0c4fb84b6e)) +* Use ArrowDtype for STRUCT columns in `to_pandas` ([#85](https://github.com/googleapis/python-bigquery-dataframes/issues/85)) ([9238fad](https://github.com/googleapis/python-bigquery-dataframes/commit/9238fadcfa7e843be6564813ff3131893b79f8b0)) + + +### Bug Fixes + +* Support multiindex for three loc getitem overloads ([#113](https://github.com/googleapis/python-bigquery-dataframes/issues/113)) ([68e3cd3](https://github.com/googleapis/python-bigquery-dataframes/commit/68e3cd37258084d045ea1075e5e61df12c28faac)) + + +### Performance Improvements + +* If primary keys are defined, `read_gbq` avoids copying table data ([#112](https://github.com/googleapis/python-bigquery-dataframes/issues/112)) ([e6c0cd1](https://github.com/googleapis/python-bigquery-dataframes/commit/e6c0cd1777736e0fa7285da59625fbac487573bd)) + + +### Documentation + +* Add documentation for `Series.struct.field` and `Series.struct.explode` ([#114](https://github.com/googleapis/python-bigquery-dataframes/issues/114)) ([a6dab9c](https://github.com/googleapis/python-bigquery-dataframes/commit/a6dab9cdb7dd0e56c93ca96b665ab1be1baac5e5)) +* Add open-source link in API doc ([#106](https://github.com/googleapis/python-bigquery-dataframes/issues/106)) ([db51fe3](https://github.com/googleapis/python-bigquery-dataframes/commit/db51fe340f644a0d7c911c11d92c8299a4be3446)) +* Update ML overview API doc ([#105](https://github.com/googleapis/python-bigquery-dataframes/issues/105)) ([1b3f3a5](https://github.com/googleapis/python-bigquery-dataframes/commit/1b3f3a5374915b2833c6c1ac05670e9708f07bff)) + ## [0.8.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.7.0...v0.8.0) (2023-10-12) diff --git a/bigframes/version.py b/bigframes/version.py index 974fbf1ac9b..3f7c8e43998 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.8.0" +__version__ = "0.9.0" From 29dd4144c7e0569de3555a16f916be9c4489bf61 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 19 Oct 2023 06:16:29 +0000 Subject: [PATCH 089/197] feat: Implement DataFrame.dot for matrix multiplication (#67) * feat: Implement DataFrame.dot for matrix multiplication Change-Id: I8baa8c8ee0aeef587be25519a9df549e6db70067 * adjust tests for left column alignment with right rows * refactor pivot to use known unique values * Support DataFrame.dot(Series), remove local-only tests * add explanation for `column_unique_values` param * NotImplementedError for multi-index matrix multiplication * avoid multi-column condition check on series * Add feedback link in the exception messages --- bigframes/core/blocks.py | 25 ++++- bigframes/dataframe.py | 102 +++++++++++++++++- tests/data/matrix_2by3.json | 22 ++++ tests/data/matrix_2by3.jsonl | 2 + tests/data/matrix_3by4.json | 27 +++++ tests/data/matrix_3by4.jsonl | 3 + tests/system/conftest.py | 68 ++++++++++++ tests/system/small/test_dataframe.py | 54 ++++++++++ tests/system/small/test_multiindex.py | 44 ++++++++ .../bigframes_vendored/pandas/core/frame.py | 30 ++++++ 10 files changed, 371 insertions(+), 6 deletions(-) create mode 100644 tests/data/matrix_2by3.json create mode 100644 tests/data/matrix_2by3.jsonl create mode 100644 tests/data/matrix_3by4.json create mode 100644 tests/data/matrix_3by4.jsonl diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 8966b6189bd..046d2b3a444 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -1261,13 +1261,28 @@ def pivot( *, columns: Sequence[str], values: Sequence[str], + columns_unique_values: typing.Optional[ + typing.Union[pd.Index, Sequence[object]] + ] = None, 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) + # We need the unique values from the pivot columns to turn them into + # column ids. It can be deteremined by running a SQL query on the + # underlying data. However, the caller can save that if they know the + # unique values upfront by providing them explicitly. + if columns_unique_values is 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) + else: + columns_values = ( + columns_unique_values + if isinstance(columns_unique_values, pd.Index) + else pd.Index(columns_unique_values) + ) column_index = columns_values column_ids: list[str] = [] diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 5740d2c4dc8..32a2908a421 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1716,10 +1716,13 @@ def kurt(self, *, numeric_only: bool = False): kurtosis = kurt - def pivot( + def _pivot( self, *, columns: typing.Union[blocks.Label, Sequence[blocks.Label]], + columns_unique_values: typing.Optional[ + typing.Union[pandas.Index, Sequence[object]] + ] = None, index: typing.Optional[ typing.Union[blocks.Label, Sequence[blocks.Label]] ] = None, @@ -1743,10 +1746,24 @@ def pivot( pivot_block = block.pivot( columns=column_ids, values=value_col_ids, + columns_unique_values=columns_unique_values, values_in_index=utils.is_list_like(values), ) return DataFrame(pivot_block) + 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: + return self._pivot(columns=columns, index=index, values=values) + def stack(self, level: LevelsType = -1): if not isinstance(self.columns, pandas.MultiIndex): if level not in [0, -1, self.columns.name]: @@ -2578,3 +2595,86 @@ def _get_block(self) -> blocks.Block: def _cached(self) -> DataFrame: return DataFrame(self._block.cached()) + + _DataFrameOrSeries = typing.TypeVar("_DataFrameOrSeries") + + def dot(self, other: _DataFrameOrSeries) -> _DataFrameOrSeries: + if not isinstance(other, (DataFrame, bf_series.Series)): + raise NotImplementedError( + f"Only DataFrame or Series operand is supported. {constants.FEEDBACK_LINK}" + ) + + if len(self.index.names) > 1 or len(other.index.names) > 1: + raise NotImplementedError( + f"Multi-index input is not supported. {constants.FEEDBACK_LINK}" + ) + + if len(self.columns.names) > 1 or ( + isinstance(other, DataFrame) and len(other.columns.names) > 1 + ): + raise NotImplementedError( + f"Multi-level column input is not supported. {constants.FEEDBACK_LINK}" + ) + + # Convert the dataframes into cell-value-decomposed representation, i.e. + # each cell value is present in a separate row + row_id = "row" + col_id = "col" + val_id = "val" + left_suffix = "_left" + right_suffix = "_right" + cvd_columns = [row_id, col_id, val_id] + + def get_left_id(id): + return f"{id}{left_suffix}" + + def get_right_id(id): + return f"{id}{right_suffix}" + + other_frame = other if isinstance(other, DataFrame) else other.to_frame() + + left = self.stack().reset_index() + left.columns = cvd_columns + + right = other_frame.stack().reset_index() + right.columns = cvd_columns + + merged = left.merge( + right, + left_on=col_id, + right_on=row_id, + suffixes=(left_suffix, right_suffix), + ) + + left_row_id = get_left_id(row_id) + right_col_id = get_right_id(col_id) + + aggregated = ( + merged.assign( + val=merged[get_left_id(val_id)] * merged[get_right_id(val_id)] + )[[left_row_id, right_col_id, val_id]] + .groupby([left_row_id, right_col_id]) + .sum(numeric_only=True) + ) + aggregated_noindex = aggregated.reset_index() + aggregated_noindex.columns = cvd_columns + result = aggregated_noindex._pivot( + columns=col_id, columns_unique_values=other_frame.columns, index=row_id + ) + + # Set the index names to match the left side matrix + result.index.names = self.index.names + + # Pivot has the result columns ordered alphabetically. It should still + # match the columns in the right sided matrix. Let's reorder them as per + # the right side matrix + if not result.columns.difference(other_frame.columns).empty: + raise RuntimeError( + f"Could not construct all columns. {constants.FEEDBACK_LINK}" + ) + result = result[other_frame.columns] + + if isinstance(other, bf_series.Series): + result = result[other.name].rename() + + return result diff --git a/tests/data/matrix_2by3.json b/tests/data/matrix_2by3.json new file mode 100644 index 00000000000..a0cf0c14da0 --- /dev/null +++ b/tests/data/matrix_2by3.json @@ -0,0 +1,22 @@ +[ + { + "mode": "REQUIRED", + "name": "rowindex", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "a", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "b", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "c", + "type": "INTEGER" + } + ] diff --git a/tests/data/matrix_2by3.jsonl b/tests/data/matrix_2by3.jsonl new file mode 100644 index 00000000000..c2c96a5423a --- /dev/null +++ b/tests/data/matrix_2by3.jsonl @@ -0,0 +1,2 @@ +{"rowindex": 0, "a": 1, "b": 2, "c": 3} +{"rowindex": 1, "a": 2, "b": 5, "c": 7} diff --git a/tests/data/matrix_3by4.json b/tests/data/matrix_3by4.json new file mode 100644 index 00000000000..6b272ee6001 --- /dev/null +++ b/tests/data/matrix_3by4.json @@ -0,0 +1,27 @@ +[ + { + "mode": "REQUIRED", + "name": "rowindex", + "type": "STRING" + }, + { + "mode": "NULLABLE", + "name": "w", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "x", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "y", + "type": "INTEGER" + }, + { + "mode": "NULLABLE", + "name": "z", + "type": "INTEGER" + } + ] diff --git a/tests/data/matrix_3by4.jsonl b/tests/data/matrix_3by4.jsonl new file mode 100644 index 00000000000..5a081dafc33 --- /dev/null +++ b/tests/data/matrix_3by4.jsonl @@ -0,0 +1,3 @@ +{"rowindex": "a", "w": 2, "x": 4, "y": 8, "z": 21} +{"rowindex": "b", "w": 1, "x": 5, "y": 10, "z": -11} +{"rowindex": "c", "w": 3, "x": 6, "y": 9, "z": 0} diff --git a/tests/system/conftest.py b/tests/system/conftest.py index f36a29b0ab1..cb664302a8f 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -235,6 +235,8 @@ def load_test_data_tables( ("penguins", "penguins_schema.json", "penguins.jsonl"), ("time_series", "time_series_schema.json", "time_series.jsonl"), ("hockey_players", "hockey_players.json", "hockey_players.jsonl"), + ("matrix_2by3", "matrix_2by3.json", "matrix_2by3.jsonl"), + ("matrix_3by4", "matrix_3by4.json", "matrix_3by4.jsonl"), ]: test_data_hash = hashlib.md5() _hash_digest_file(test_data_hash, DATA_DIR / schema_filename) @@ -304,6 +306,16 @@ def time_series_table_id(test_data_tables) -> str: return test_data_tables["time_series"] +@pytest.fixture(scope="session") +def matrix_2by3_table_id(test_data_tables) -> str: + return test_data_tables["matrix_2by3"] + + +@pytest.fixture(scope="session") +def matrix_3by4_table_id(test_data_tables) -> str: + return test_data_tables["matrix_3by4"] + + @pytest.fixture(scope="session") def scalars_df_default_index( scalars_df_index: bigframes.dataframe.DataFrame, @@ -411,6 +423,62 @@ def hockey_pandas_df() -> pd.DataFrame: return df +@pytest.fixture(scope="session") +def matrix_2by3_df( + matrix_2by3_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at a test 2-by-3 matrix data.""" + df = session.read_gbq(matrix_2by3_table_id) + df = df.set_index("rowindex").sort_index() + return df + + +@pytest.fixture(scope="session") +def matrix_2by3_pandas_df() -> pd.DataFrame: + """pd.DataFrame pointing at a test 2-by-3 matrix data.""" + df = pd.read_json( + DATA_DIR / "matrix_2by3.jsonl", + lines=True, + dtype={ + "rowindex": pd.Int64Dtype(), + "a": pd.Int64Dtype(), + "b": pd.Int64Dtype(), + "c": pd.Int64Dtype(), + }, + ) + df = df.set_index("rowindex").sort_index() + df.index = df.index.astype("Int64") + return df + + +@pytest.fixture(scope="session") +def matrix_3by4_df( + matrix_3by4_table_id: str, session: bigframes.Session +) -> bigframes.dataframe.DataFrame: + """DataFrame pointing at a test 3-by-4 matrix data.""" + df = session.read_gbq(matrix_3by4_table_id) + df = df.set_index("rowindex").sort_index() + return df + + +@pytest.fixture(scope="session") +def matrix_3by4_pandas_df() -> pd.DataFrame: + """pd.DataFrame pointing at a test 3-by-4 matrix data.""" + df = pd.read_json( + DATA_DIR / "matrix_3by4.jsonl", + lines=True, + dtype={ + "rowindex": pd.StringDtype(storage="pyarrow"), + "w": pd.Int64Dtype(), + "x": pd.Int64Dtype(), + "y": pd.Int64Dtype(), + "z": pd.Int64Dtype(), + }, + ) + df = df.set_index("rowindex").sort_index() + return df + + @pytest.fixture(scope="session") def penguins_df_default_index( penguins_table_id: str, session: bigframes.Session diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 309e8df4f0d..19e50eb06de 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3167,3 +3167,57 @@ def test_df_cached(scalars_df_index): df_cached_copy = df._cached() pandas.testing.assert_frame_equal(df.to_pandas(), df_cached_copy.to_pandas()) + + +def test_df_dot_inline(session): + df1 = pd.DataFrame([[1, 2, 3], [2, 5, 7]]) + df2 = pd.DataFrame([[2, 4, 8], [1, 5, 10], [3, 6, 9]]) + + bf1 = session.read_pandas(df1) + bf2 = session.read_pandas(df2) + bf_result = bf1.dot(bf2).to_pandas() + pd_result = df1.dot(df2) + + # Patch pandas dtypes for testing parity + # Pandas uses int64 instead of Int64 (nullable) dtype. + for name in pd_result.columns: + pd_result[name] = pd_result[name].astype(pd.Int64Dtype()) + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_dot( + matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df +): + bf_result = matrix_2by3_df.dot(matrix_3by4_df).to_pandas() + pd_result = matrix_2by3_pandas_df.dot(matrix_3by4_pandas_df) + + # Patch pandas dtypes for testing parity + # Pandas result is object instead of Int64 (nullable) dtype. + for name in pd_result.columns: + pd_result[name] = pd_result[name].astype(pd.Int64Dtype()) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + +def test_df_dot_series( + matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df +): + bf_result = matrix_2by3_df.dot(matrix_3by4_df["x"]).to_pandas() + pd_result = matrix_2by3_pandas_df.dot(matrix_3by4_pandas_df["x"]) + + # Patch pandas dtypes for testing parity + # Pandas result is object instead of Int64 (nullable) dtype. + pd_result = pd_result.astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index a1326767700..b5c78de69c2 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -934,3 +934,47 @@ def test_column_multi_index_swaplevel(scalars_df_index, scalars_pandas_df_index) pd_result = pd_df.swaplevel(-3, -1, axis=1) pandas.testing.assert_frame_equal(bf_result, pd_result) + + +def test_df_multi_index_dot_not_supported(): + left_matrix = [[1, 2, 3], [2, 5, 7]] + right_matrix = [[2, 4, 8], [1, 5, 10], [3, 6, 9]] + + # Left multi-index + left_index = pandas.MultiIndex.from_tuples([("a", "aa"), ("a", "ab")]) + bf1 = bpd.DataFrame(left_matrix, index=left_index) + bf2 = bpd.DataFrame(right_matrix) + with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): + bf1.dot(bf2) + + # right multi-index + right_index = pandas.MultiIndex.from_tuples([("a", "aa"), ("a", "ab"), ("b", "bb")]) + bf1 = bpd.DataFrame(left_matrix) + bf2 = bpd.DataFrame(right_matrix, index=right_index) + with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): + bf1.dot(bf2) + + +def test_column_multi_index_dot_not_supported(): + left_matrix = [[1, 2, 3], [2, 5, 7]] + right_matrix = [[2, 4, 8], [1, 5, 10], [3, 6, 9]] + + multi_level_columns = pandas.MultiIndex.from_arrays( + [["col0", "col0", "col1"], ["col00", "col01", "col11"]] + ) + + # Left multi-columns + bf1 = bpd.DataFrame(left_matrix, columns=multi_level_columns) + bf2 = bpd.DataFrame(right_matrix) + with pytest.raises( + NotImplementedError, match="Multi-level column input is not supported" + ): + bf1.dot(bf2) + + # right multi-columns + bf1 = bpd.DataFrame(left_matrix) + bf2 = bpd.DataFrame(right_matrix, columns=multi_level_columns) + with pytest.raises( + NotImplementedError, match="Multi-level column input is not supported" + ): + bf1.dot(bf2) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 621d052cb8e..a5c12d7b32a 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2117,3 +2117,33 @@ def iat(self): def at(self): """Access a single value for a row/column label pair.""" raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def dot(self, other): + """ + Compute the matrix multiplication between the DataFrame and other. + + This method computes the matrix product between the DataFrame and the + values of an other Series or DataFrame. + + It can also be called using `self @ other`. + + .. note:: + The dimensions of DataFrame and other must be compatible in order to + compute the matrix multiplication. In addition, the column names of + DataFrame and the index of other must contain the same values, as they + will be aligned prior to the multiplication. + + The dot method for Series computes the inner product, instead of the + matrix product here. + + Args: + other (Series or DataFrame): + The other object to compute the matrix product with. + + Returns: + Series or DataFrame + If `other` is a Series, return the matrix product between self and + other as a Series. If other is a DataFrame, return + the matrix product of self and other in a DataFrame. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 782a215421c743999782f94eba76c82db93cc85c Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Thu, 19 Oct 2023 10:43:10 -0500 Subject: [PATCH 090/197] chore(main): release 0.10.0 (#122) 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 974e5d3d43d..4d9f63d4c66 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.10.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.9.0...v0.10.0) (2023-10-19) + + +### Features + +* Implement DataFrame.dot for matrix multiplication ([#67](https://github.com/googleapis/python-bigquery-dataframes/issues/67)) ([29dd414](https://github.com/googleapis/python-bigquery-dataframes/commit/29dd4144c7e0569de3555a16f916be9c4489bf61)) + ## [0.9.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.8.0...v0.9.0) (2023-10-18) diff --git a/bigframes/version.py b/bigframes/version.py index 3f7c8e43998..7a37ebd2205 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.9.0" +__version__ = "0.10.0" From 59bc7be9891ac12926ac44531d9dbd4a88bcfc44 Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Thu, 19 Oct 2023 17:02:14 +0000 Subject: [PATCH 091/197] chore: rename rst files to avoid conflict with service names (#120) Source-Link: https://togithub.com/googleapis/synthtool/commit/d52e638b37b091054c869bfa6f5a9fedaba9e0dd Post-Processor: gcr.io/cloud-devrel-public-resources/owlbot-python:latest@sha256:4f9b3b106ad0beafc2c8a415e3f62c1a0cc23cabea115dbe841b848f581cfe99 --- .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 dd98abbdeeb..7f291dbd5f9 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:08e34975760f002746b1d8c86fdc90660be45945ee6d9db914d1508acdf9a547 -# created: 2023-10-09T14:06:13.397766266Z + digest: sha256:4f9b3b106ad0beafc2c8a415e3f62c1a0cc23cabea115dbe841b848f581cfe99 +# created: 2023-10-18T20:26:37.410353675Z diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 0332d3267e1..16170d0ca7b 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.17 \ - --hash=sha256:24d6a242c28d29af46c3fae832c36db3bbebcc533dd1bb549172cd739c82df21 \ - --hash=sha256:94a757d178c9be92ef5539b8840d48dc9cf1b2709c9d6b588232a055c524458b +urllib3==1.26.18 \ + --hash=sha256:34b97092d7e0a3a8cf7cd10e386f401b3737364026c45e622aa02903dffe0f07 \ + --hash=sha256:f8ecc1bba5667413457c529ab955bf8c67b45db799d159066261719e328580a0 # via # requests # twine From c9c46d45eca536d1fc2dfed0cce25bfb41a41554 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 19 Oct 2023 13:08:14 -0500 Subject: [PATCH 092/197] chore: remove unneeded scripts, no longer need THIRD_PARTY_NOTICES (#121) 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 internal issue 306233684 🦕 --- .kokoro/release-nightly.sh | 44 +--- scripts/generate_third_party_notices.py | 332 ------------------------ scripts/update_firebase_docs_site.sh | 105 -------- scripts/update_x20_docs_site.sh | 106 -------- 4 files changed, 1 insertion(+), 586 deletions(-) delete mode 100644 scripts/generate_third_party_notices.py delete mode 100644 scripts/update_firebase_docs_site.sh delete mode 100644 scripts/update_x20_docs_site.sh diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh index 488dbb9e130..5cc1275308a 100755 --- a/.kokoro/release-nightly.sh +++ b/.kokoro/release-nightly.sh @@ -63,26 +63,6 @@ export PYTHONUNBUFFERED=1 # Install dependencies, as the following steps depend on it python3.10 -m pip install -e .[all] -# 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) @@ -101,33 +81,13 @@ if [ $num_wheel_files -ne 1 ] ; then 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/ \ +for gcs_path in gs://vertex_sdk_private_releases/bigframe/ \ gs://dl-platform-colab/bigframes/ \ gs://bigframes-wheels/; do @@ -155,8 +115,6 @@ fi # 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} if ! [ ${DRY_RUN} ]; then # Copy docs and wheels to Google Drive diff --git a/scripts/generate_third_party_notices.py b/scripts/generate_third_party_notices.py deleted file mode 100644 index 7040bb2e5f8..00000000000 --- a/scripts/generate_third_party_notices.py +++ /dev/null @@ -1,332 +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 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/update_firebase_docs_site.sh b/scripts/update_firebase_docs_site.sh deleted file mode 100644 index f0ef866c90b..00000000000 --- a/scripts/update_firebase_docs_site.sh +++ /dev/null @@ -1,105 +0,0 @@ -#!/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 deleted file mode 100644 index 31da116bdd4..00000000000 --- a/scripts/update_x20_docs_site.sh +++ /dev/null @@ -1,106 +0,0 @@ -#!/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} From 694a85a0ef90d838700014a204d72b23362db1d8 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 20 Oct 2023 10:52:49 -0500 Subject: [PATCH 093/197] feat: add back `reset_session` as an alias for `close_session` (#124) --- bigframes/pandas/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 1b9144fb625..24b19fa70a2 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -452,6 +452,7 @@ def read_gbq_function(function_name: str): # Session management APIs get_global_session = global_session.get_global_session close_session = global_session.close_session +reset_session = global_session.close_session # Use __all__ to let type checkers know what is part of the public API. From f9bb3c4bc88c5ba2be6f17e12a0ec4f482ce161f Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 20 Oct 2023 13:48:13 -0700 Subject: [PATCH 094/197] feat: change `query` parameter to `query_or_table` in `read_gbq` (#127) 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/pandas/__init__.py | 6 +++--- bigframes/session/__init__.py | 8 ++++---- third_party/bigframes_vendored/pandas/io/gbq.py | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 24b19fa70a2..971d40f801e 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -290,16 +290,16 @@ def read_json( def read_gbq( - query: str, + query_or_table: 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) + _set_default_session_location_if_possible(query_or_table) return global_session.with_default_session( bigframes.session.Session.read_gbq, - query, + query_or_table, index_col=index_col, col_order=col_order, max_results=max_results, diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 1031fde9b5b..2f001d7d49a 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -262,7 +262,7 @@ def close(self): def read_gbq( self, - query: str, + query_or_table: str, *, index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), @@ -270,9 +270,9 @@ def read_gbq( # 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): + if _is_query(query_or_table): return self._read_gbq_query( - query, + query_or_table, index_col=index_col, col_order=col_order, max_results=max_results, @@ -283,7 +283,7 @@ def read_gbq( # deterministic query so we can avoid serializing if we have a # unique index. return self._read_gbq_table( - query, + query_or_table, index_col=index_col, col_order=col_order, max_results=max_results, diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 95531ff5e82..8919f4ed16d 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -11,7 +11,7 @@ class GBQIOMixin: def read_gbq( self, - query: str, + query_or_table: str, *, index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), @@ -69,7 +69,7 @@ def read_gbq( [5 rows x 3 columns] Args: - query (str): + query_or_table (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`. From ba2e824f3abbea30010d845e53d645db52504361 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Sat, 21 Oct 2023 01:18:20 +0000 Subject: [PATCH 095/197] refactor: Tweak notebooks to test minimal IAM permissions (#123) * refactor: Tweak notebooks to test minimal IAM permissions * Make IAM permission comment more helpful --- notebooks/dataframes/dataframe.ipynb | 313 ++- .../getting_started/ml_fundamentals.ipynb | 2005 +++++++++-------- .../regression/easy_linear_regression.ipynb | 190 +- 3 files changed, 1425 insertions(+), 1083 deletions(-) diff --git a/notebooks/dataframes/dataframe.ipynb b/notebooks/dataframes/dataframe.ipynb index 85ea61d2810..c6b276af877 100644 --- a/notebooks/dataframes/dataframe.ipynb +++ b/notebooks/dataframes/dataframe.ipynb @@ -35,12 +35,26 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "944f0e4417154e81b6496302fe756465", + "model_id": "11c27813da5c4d2e8108bf4bd9e7e55d", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "HTML(value='Query job ac4d1f2b-e9f3-4d95-b78d-57e40eee93fa is RUNNING. \n", "
241156Biscoe46.214.9221.0MALEGentoo penguin (Pygoscelis papua)
121Dream48.116.4199.014.5209.0FEMALEChinstrap penguin (Pygoscelis antarctica)Gentoo penguin (Pygoscelis papua)
209189Biscoe42.718.3196.0MALE35.318.9187.0FEMALEAdelie Penguin (Pygoscelis adeliae)
270279Biscoe37.716.0183.045.114.5215.0FEMALEAdelie Penguin (Pygoscelis adeliae)Gentoo penguin (Pygoscelis papua)
187245Biscoe43.414.4218.0FEMALE49.516.2229.0MALEGentoo penguin (Pygoscelis papua)
343Torgersen37.320.5199.0MALEAdelie Penguin (Pygoscelis adeliae)
\n", "

5 rows × 6 columns

\n", "[5 rows x 6 columns in total]" ], "text/plain": [ - " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", - "penguin_id \n", - "241 Biscoe 46.2 14.9 221.0 \n", - "121 Dream 48.1 16.4 199.0 \n", - "209 Biscoe 42.7 18.3 196.0 \n", - "270 Biscoe 37.7 16.0 183.0 \n", - "187 Biscoe 43.4 14.4 218.0 \n", + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "penguin_id \n", + "156 Biscoe 46.2 14.5 209.0 \n", + "189 Biscoe 35.3 18.9 187.0 \n", + "279 Biscoe 45.1 14.5 215.0 \n", + "245 Biscoe 49.5 16.2 229.0 \n", + "343 Torgersen 37.3 20.5 199.0 \n", "\n", - " sex species \n", - "penguin_id \n", - "241 MALE Gentoo penguin (Pygoscelis papua) \n", - "121 FEMALE Chinstrap penguin (Pygoscelis antarctica) \n", - "209 MALE Adelie Penguin (Pygoscelis adeliae) \n", - "270 FEMALE Adelie Penguin (Pygoscelis adeliae) \n", - "187 FEMALE Gentoo penguin (Pygoscelis papua) \n", + " sex species \n", + "penguin_id \n", + "156 FEMALE Gentoo penguin (Pygoscelis papua) \n", + "189 FEMALE Adelie Penguin (Pygoscelis adeliae) \n", + "279 FEMALE Gentoo penguin (Pygoscelis papua) \n", + "245 MALE Gentoo penguin (Pygoscelis papua) \n", + "343 MALE Adelie Penguin (Pygoscelis adeliae) \n", "\n", "[5 rows x 6 columns]" ] }, - "execution_count": 3, + "execution_count": 20, "metadata": {}, "output_type": "execute_result" } @@ -763,18 +749,18 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 21, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "f09b707c5c9540cdae59af3765339b6e", + "model_id": "d6dd794f89724099950dcc927d63d0f5", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "HTML(value='Query job 9f0cf977-a895-41fa-9f33-8b5346c62786 is DONE. 31.7 kB processed. \n", "
2415300.01564800.0
1213325.01893800.0
2094075.02795000.0
2703075.02455800.0
1874600.03433775.0
\n", @@ -865,16 +837,16 @@ "text/plain": [ " body_mass_g\n", "penguin_id \n", - "241 5300.0\n", - "121 3325.0\n", - "209 4075.0\n", - "270 3075.0\n", - "187 4600.0\n", + "156 4800.0\n", + "189 3800.0\n", + "279 5000.0\n", + "245 5800.0\n", + "343 3775.0\n", "\n", "[5 rows x 1 columns]" ] }, - "execution_count": 4, + "execution_count": 21, "metadata": {}, "output_type": "execute_result" } @@ -908,18 +880,18 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 22, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "299c0c895e3d4a83a8495924a1966ce0", + "model_id": "380c57dc3fe54fbd8ad2fb23f1e66e37", "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", - " scaled_culmen_length_mm\n", - " scaled_culmen_depth_mm\n", - " scaled_flipper_length_mm\n", + " standard_scaled_culmen_length_mm\n", + " standard_scaled_culmen_depth_mm\n", + " standard_scaled_flipper_length_mm\n", " \n", " \n", " penguin_id\n", @@ -1016,153 +1002,153 @@ " \n", " \n", " 0\n", - " -1.364965\n", - " 0.629892\n", - " -1.226537\n", + " -1.344188\n", + " 0.642519\n", + " -1.193942\n", " \n", " \n", " 1\n", - " -0.771824\n", - " 0.984275\n", - " -1.226537\n", + " -0.750047\n", + " 1.005876\n", + " -1.193942\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", + " -0.545811\n", + " 0.90206\n", + " -1.193942\n", " \n", " \n", " 4\n", - " -1.235216\n", - " -0.180128\n", - " -0.652517\n", + " -1.214219\n", + " -0.188011\n", + " -0.619171\n", " \n", " \n", " 5\n", - " -0.141612\n", - " 0.680518\n", - " -0.652517\n", + " -0.118772\n", + " 0.694427\n", + " -0.619171\n", " \n", " \n", " 6\n", - " 0.544207\n", - " -0.281381\n", - " -0.652517\n", + " 0.568203\n", + " -0.291828\n", + " -0.619171\n", " \n", " \n", " 7\n", - " 1.21149\n", - " 0.629892\n", - " -0.078497\n", + " 1.236611\n", + " 0.642519\n", + " -0.044401\n", " \n", " \n", - " 8\n", - " 1.026133\n", - " 0.933649\n", - " -0.078497\n", + " 9\n", + " -0.675779\n", + " 1.524957\n", + " -0.044401\n", " \n", " \n", " 10\n", - " -0.586468\n", - " 0.883023\n", - " 0.495523\n", + " -0.564378\n", + " 0.90206\n", + " 0.530369\n", " \n", " \n", " 11\n", - " -0.92011\n", - " 0.781771\n", - " -1.154784\n", + " -0.898582\n", + " 0.798243\n", + " -1.122096\n", + " \n", + " \n", + " 12\n", + " -1.26992\n", + " -0.136103\n", + " -1.122096\n", + " \n", + " \n", + " 13\n", + " 0.58677\n", + " 0.071529\n", + " -1.122096\n", " \n", " \n", " 14\n", - " -1.846892\n", - " -0.028249\n", - " -1.154784\n", + " -1.826927\n", + " -0.032287\n", + " -1.122096\n", " \n", " \n", " 15\n", - " -1.290822\n", - " -0.332007\n", - " -1.154784\n", + " -1.26992\n", + " -0.343736\n", + " -1.122096\n", " \n", " \n", " 16\n", - " 0.321779\n", - " 0.073003\n", - " -0.580765\n", + " 0.3454\n", + " 0.071529\n", + " -0.547325\n", " \n", " \n", - " 17\n", - " 1.230026\n", - " 1.136154\n", - " -0.580765\n", + " 18\n", + " -0.768614\n", + " 0.382978\n", + " -0.547325\n", " \n", " \n", - " 18\n", - " -0.79036\n", - " 0.376761\n", - " -0.580765\n", + " 19\n", + " -1.121385\n", + " 0.486795\n", + " -0.547325\n", " \n", " \n", " 20\n", - " 0.4886\n", - " 0.326134\n", - " -0.580765\n", + " 0.512502\n", + " 0.33107\n", + " -0.547325\n", " \n", " \n", " 21\n", - " 1.359775\n", - " 1.034902\n", - " -0.580765\n", + " 1.385146\n", + " 1.057784\n", + " -0.547325\n", + " \n", + " \n", + " 22\n", + " -0.675779\n", + " -0.032287\n", + " -0.547325\n", " \n", " \n", " 24\n", - " 1.044669\n", - " 0.528639\n", - " -0.580765\n", + " 1.069509\n", + " 0.538703\n", + " -0.547325\n", " \n", " \n", " 26\n", - " -0.456718\n", - " 0.680518\n", - " -0.006745\n", + " -0.43441\n", + " 0.694427\n", + " 0.027445\n", " \n", " \n", - " 27\n", - " 1.21149\n", - " 1.237407\n", - " -0.006745\n", + " 28\n", + " 1.923586\n", + " 1.888314\n", + " 0.027445\n", " \n", " \n", - " 29\n", - " 1.378311\n", - " 0.933649\n", - " -0.006745\n", + " 30\n", + " 1.292312\n", + " 0.694427\n", + " 0.027445\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", + " -1.994029\n", + " -0.551368\n", + " -1.62502\n", " \n", " \n", "\n", @@ -1170,67 +1156,67 @@ "[267 rows x 3 columns in total]" ], "text/plain": [ - " scaled_culmen_length_mm scaled_culmen_depth_mm \\\n", - "penguin_id \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", + " standard_scaled_culmen_length_mm standard_scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "0 -1.344188 0.642519 \n", + "1 -0.750047 1.005876 \n", + "2 -0.545811 0.90206 \n", + "4 -1.214219 -0.188011 \n", + "5 -0.118772 0.694427 \n", + "6 0.568203 -0.291828 \n", + "7 1.236611 0.642519 \n", + "9 -0.675779 1.524957 \n", + "10 -0.564378 0.90206 \n", + "11 -0.898582 0.798243 \n", + "12 -1.26992 -0.136103 \n", + "13 0.58677 0.071529 \n", + "14 -1.826927 -0.032287 \n", + "15 -1.26992 -0.343736 \n", + "16 0.3454 0.071529 \n", + "18 -0.768614 0.382978 \n", + "19 -1.121385 0.486795 \n", + "20 0.512502 0.33107 \n", + "21 1.385146 1.057784 \n", + "22 -0.675779 -0.032287 \n", + "24 1.069509 0.538703 \n", + "26 -0.43441 0.694427 \n", + "28 1.923586 1.888314 \n", + "30 1.292312 0.694427 \n", + "31 -1.994029 -0.551368 \n", "\n", - " scaled_flipper_length_mm \n", - "penguin_id \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", + " standard_scaled_flipper_length_mm \n", + "penguin_id \n", + "0 -1.193942 \n", + "1 -1.193942 \n", + "2 -1.193942 \n", + "4 -0.619171 \n", + "5 -0.619171 \n", + "6 -0.619171 \n", + "7 -0.044401 \n", + "9 -0.044401 \n", + "10 0.530369 \n", + "11 -1.122096 \n", + "12 -1.122096 \n", + "13 -1.122096 \n", + "14 -1.122096 \n", + "15 -1.122096 \n", + "16 -0.547325 \n", + "18 -0.547325 \n", + "19 -0.547325 \n", + "20 -0.547325 \n", + "21 -0.547325 \n", + "22 -0.547325 \n", + "24 -0.547325 \n", + "26 0.027445 \n", + "28 0.027445 \n", + "30 0.027445 \n", + "31 -1.62502 \n", "...\n", "\n", "[267 rows x 3 columns]" ] }, - "execution_count": 5, + "execution_count": 22, "metadata": {}, "output_type": "execute_result" } @@ -1251,18 +1237,18 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 23, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "6ec4ab3e60844e989dbebad89e7665ca", + "model_id": "74f3c24c0a434e12bf6a56dc4809b501", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "HTML(value='Query job a8029b51-2ef1-4acd-9759-d808db954298 is DONE. 31.7 kB processed. \n", - " -1.364965\n", - " 0.629892\n", - " -1.226537\n", + " -1.344188\n", + " 0.642519\n", + " -1.193942\n", " [{'index': 2, 'value': 1.0}]\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", + " -0.750047\n", + " 1.005876\n", + " -1.193942\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 2\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.567932\n", - " 0.883023\n", - " -1.226537\n", + " -0.545811\n", + " 0.90206\n", + " -1.193942\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", " 4\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.235216\n", - " -0.180128\n", - " -0.652517\n", + " -1.214219\n", + " -0.188011\n", + " -0.619171\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 5\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.141612\n", - " 0.680518\n", - " -0.652517\n", + " -0.118772\n", + " 0.694427\n", + " -0.619171\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 6\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.544207\n", - " -0.281381\n", - " -0.652517\n", + " 0.568203\n", + " -0.291828\n", + " -0.619171\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", " 7\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.21149\n", - " 0.629892\n", - " -0.078497\n", + " 1.236611\n", + " 0.642519\n", + " -0.044401\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 8\n", + " 9\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.026133\n", - " 0.933649\n", - " -0.078497\n", + " -0.675779\n", + " 1.524957\n", + " -0.044401\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 10\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.586468\n", - " 0.883023\n", - " 0.495523\n", + " -0.564378\n", + " 0.90206\n", + " 0.530369\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 11\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.92011\n", - " 0.781771\n", - " -1.154784\n", + " -0.898582\n", + " 0.798243\n", + " -1.122096\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", + " 12\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.26992\n", + " -0.136103\n", + " -1.122096\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", + " \n", + " \n", + " 13\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.58677\n", + " 0.071529\n", + " -1.122096\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", + " \n", + " \n", " 14\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.846892\n", - " -0.028249\n", - " -1.154784\n", + " -1.826927\n", + " -0.032287\n", + " -1.122096\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 15\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.290822\n", - " -0.332007\n", - " -1.154784\n", + " -1.26992\n", + " -0.343736\n", + " -1.122096\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 16\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.321779\n", - " 0.073003\n", - " -0.580765\n", + " 0.3454\n", + " 0.071529\n", + " -0.547325\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 17\n", + " 18\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.230026\n", - " 1.136154\n", - " -0.580765\n", + " -0.768614\n", + " 0.382978\n", + " -0.547325\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 18\n", + " 19\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.79036\n", - " 0.376761\n", - " -0.580765\n", + " -1.121385\n", + " 0.486795\n", + " -0.547325\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 20\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.4886\n", - " 0.326134\n", - " -0.580765\n", + " 0.512502\n", + " 0.33107\n", + " -0.547325\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", + " 1.385146\n", + " 1.057784\n", + " -0.547325\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", + " 22\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.675779\n", + " -0.032287\n", + " -0.547325\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", + " \n", + " \n", " 24\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.044669\n", - " 0.528639\n", - " -0.580765\n", + " 1.069509\n", + " 0.538703\n", + " -0.547325\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", " 26\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.456718\n", - " 0.680518\n", - " -0.006745\n", + " -0.43441\n", + " 0.694427\n", + " 0.027445\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 27\n", + " 28\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.21149\n", - " 1.237407\n", - " -0.006745\n", + " 1.923586\n", + " 1.888314\n", + " 0.027445\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 29\n", + " 30\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.378311\n", - " 0.933649\n", - " -0.006745\n", + " 1.292312\n", + " 0.694427\n", + " 0.027445\n", " [{'index': 3, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \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", - " 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", - " 33\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.827431\n", - " -0.230754\n", - " -1.657052\n", + " -1.994029\n", + " -0.551368\n", + " -1.62502\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", @@ -1966,95 +1980,123 @@ "[267 rows x 6 columns in total]" ], "text/plain": [ - " onehotencoded_island scaled_culmen_length_mm \\\n", - "penguin_id \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", + " onehotencoded_island standard_scaled_culmen_length_mm \\\n", + "penguin_id \n", + "0 [{'index': 2, 'value': 1.0}] -1.344188 \n", + "1 [{'index': 2, 'value': 1.0}] -0.750047 \n", + "2 [{'index': 2, 'value': 1.0}] -0.545811 \n", + "4 [{'index': 2, 'value': 1.0}] -1.214219 \n", + "5 [{'index': 2, 'value': 1.0}] -0.118772 \n", + "6 [{'index': 2, 'value': 1.0}] 0.568203 \n", + "7 [{'index': 2, 'value': 1.0}] 1.236611 \n", + "9 [{'index': 2, 'value': 1.0}] -0.675779 \n", + "10 [{'index': 2, 'value': 1.0}] -0.564378 \n", + "11 [{'index': 2, 'value': 1.0}] -0.898582 \n", + "12 [{'index': 2, 'value': 1.0}] -1.26992 \n", + "13 [{'index': 2, 'value': 1.0}] 0.58677 \n", + "14 [{'index': 2, 'value': 1.0}] -1.826927 \n", + "15 [{'index': 2, 'value': 1.0}] -1.26992 \n", + "16 [{'index': 2, 'value': 1.0}] 0.3454 \n", + "18 [{'index': 2, 'value': 1.0}] -0.768614 \n", + "19 [{'index': 2, 'value': 1.0}] -1.121385 \n", + "20 [{'index': 2, 'value': 1.0}] 0.512502 \n", + "21 [{'index': 2, 'value': 1.0}] 1.385146 \n", + "22 [{'index': 2, 'value': 1.0}] -0.675779 \n", + "24 [{'index': 2, 'value': 1.0}] 1.069509 \n", + "26 [{'index': 2, 'value': 1.0}] -0.43441 \n", + "28 [{'index': 2, 'value': 1.0}] 1.923586 \n", + "30 [{'index': 2, 'value': 1.0}] 1.292312 \n", + "31 [{'index': 2, 'value': 1.0}] -1.994029 \n", "\n", - " scaled_culmen_depth_mm scaled_flipper_length_mm \\\n", - "penguin_id \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", + " standard_scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "0 0.642519 \n", + "1 1.005876 \n", + "2 0.90206 \n", + "4 -0.188011 \n", + "5 0.694427 \n", + "6 -0.291828 \n", + "7 0.642519 \n", + "9 1.524957 \n", + "10 0.90206 \n", + "11 0.798243 \n", + "12 -0.136103 \n", + "13 0.071529 \n", + "14 -0.032287 \n", + "15 -0.343736 \n", + "16 0.071529 \n", + "18 0.382978 \n", + "19 0.486795 \n", + "20 0.33107 \n", + "21 1.057784 \n", + "22 -0.032287 \n", + "24 0.538703 \n", + "26 0.694427 \n", + "28 1.888314 \n", + "30 0.694427 \n", + "31 -0.551368 \n", "\n", - " onehotencoded_sex onehotencoded_species \n", - "penguin_id \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", + " standard_scaled_flipper_length_mm onehotencoded_sex \\\n", + "penguin_id \n", + "0 -1.193942 [{'index': 2, 'value': 1.0}] \n", + "1 -1.193942 [{'index': 3, 'value': 1.0}] \n", + "2 -1.193942 [{'index': 3, 'value': 1.0}] \n", + "4 -0.619171 [{'index': 2, 'value': 1.0}] \n", + "5 -0.619171 [{'index': 3, 'value': 1.0}] \n", + "6 -0.619171 [{'index': 2, 'value': 1.0}] \n", + "7 -0.044401 [{'index': 2, 'value': 1.0}] \n", + "9 -0.044401 [{'index': 3, 'value': 1.0}] \n", + "10 0.530369 [{'index': 3, 'value': 1.0}] \n", + "11 -1.122096 [{'index': 3, 'value': 1.0}] \n", + "12 -1.122096 [{'index': 2, 'value': 1.0}] \n", + "13 -1.122096 [{'index': 2, 'value': 1.0}] \n", + "14 -1.122096 [{'index': 2, 'value': 1.0}] \n", + "15 -1.122096 [{'index': 2, 'value': 1.0}] \n", + "16 -0.547325 [{'index': 2, 'value': 1.0}] \n", + "18 -0.547325 [{'index': 3, 'value': 1.0}] \n", + "19 -0.547325 [{'index': 3, 'value': 1.0}] \n", + "20 -0.547325 [{'index': 2, 'value': 1.0}] \n", + "21 -0.547325 [{'index': 3, 'value': 1.0}] \n", + "22 -0.547325 [{'index': 2, 'value': 1.0}] \n", + "24 -0.547325 [{'index': 3, 'value': 1.0}] \n", + "26 0.027445 [{'index': 3, 'value': 1.0}] \n", + "28 0.027445 [{'index': 3, 'value': 1.0}] \n", + "30 0.027445 [{'index': 3, 'value': 1.0}] \n", + "31 -1.62502 [{'index': 2, 'value': 1.0}] \n", + "\n", + " onehotencoded_species \n", + "penguin_id \n", + "0 [{'index': 1, 'value': 1.0}] \n", + "1 [{'index': 1, 'value': 1.0}] \n", + "2 [{'index': 1, 'value': 1.0}] \n", + "4 [{'index': 1, 'value': 1.0}] \n", + "5 [{'index': 1, 'value': 1.0}] \n", + "6 [{'index': 2, 'value': 1.0}] \n", + "7 [{'index': 2, 'value': 1.0}] \n", + "9 [{'index': 1, 'value': 1.0}] \n", + "10 [{'index': 1, 'value': 1.0}] \n", + "11 [{'index': 1, 'value': 1.0}] \n", + "12 [{'index': 1, 'value': 1.0}] \n", + "13 [{'index': 2, 'value': 1.0}] \n", + "14 [{'index': 1, 'value': 1.0}] \n", + "15 [{'index': 1, 'value': 1.0}] \n", + "16 [{'index': 2, 'value': 1.0}] \n", + "18 [{'index': 1, 'value': 1.0}] \n", + "19 [{'index': 1, 'value': 1.0}] \n", + "20 [{'index': 2, 'value': 1.0}] \n", + "21 [{'index': 2, 'value': 1.0}] \n", + "22 [{'index': 1, 'value': 1.0}] \n", + "24 [{'index': 2, 'value': 1.0}] \n", + "26 [{'index': 1, 'value': 1.0}] \n", + "28 [{'index': 2, 'value': 1.0}] \n", + "30 [{'index': 2, 'value': 1.0}] \n", + "31 [{'index': 1, 'value': 1.0}] \n", "...\n", "\n", "[267 rows x 6 columns]" ] }, - "execution_count": 7, + "execution_count": 24, "metadata": {}, "output_type": "execute_result" } @@ -2096,18 +2138,32 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 25, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "b7c22c0858914b46951adde174b43e25", + "model_id": "5db4c5c80ba4417db151aa561dab5ee7", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job ceced0cc-13a7-4b14-b42c-4d5f69e7e49a is RUNNING. \n", " \n", " \n", - " 9\n", - " 4295.335461\n", + " 3\n", + " 3394.118128\n", " \n", " \n", - " 12\n", - " 3338.44131\n", + " 8\n", + " 4048.685642\n", " \n", " \n", - " 13\n", - " 3201.820204\n", + " 17\n", + " 3976.454093\n", " \n", " \n", - " 19\n", - " 3982.814079\n", + " 23\n", + " 3541.582194\n", " \n", " \n", - " 22\n", - " 3538.610664\n", + " 25\n", + " 4032.844186\n", " \n", " \n", - " 23\n", - " 3613.50305\n", + " 27\n", + " 4118.351772\n", " \n", " \n", - " 25\n", - " 4009.759444\n", + " 29\n", + " 4087.767826\n", " \n", " \n", - " 28\n", - " 4240.515635\n", + " 34\n", + " 3183.755249\n", " \n", " \n", - " 30\n", - " 4028.904195\n", + " 35\n", + " 3418.802274\n", " \n", " \n", - " 38\n", - " 4206.810346\n", + " 39\n", + " 3519.186468\n", " \n", " \n", - " 41\n", - " 3736.225488\n", + " 51\n", + " 3398.135365\n", " \n", " \n", - " 45\n", - " 4383.368544\n", + " 52\n", + " 3223.615957\n", " \n", " \n", - " 49\n", - " 3125.772789\n", + " 60\n", + " 3445.014718\n", " \n", " \n", - " 50\n", - " 3149.28765\n", + " 61\n", + " 3505.638864\n", " \n", " \n", - " 62\n", - " 3531.69488\n", + " 64\n", + " 3515.905786\n", " \n", " \n", " 65\n", - " 4073.900616\n", + " 4028.363185\n", " \n", " \n", - " 66\n", - " 4160.810162\n", + " 67\n", + " 4159.993943\n", " \n", " \n", " 83\n", - " 3272.396279\n", + " 3348.16883\n", " \n", " \n", - " 87\n", - " 3989.364493\n", + " 85\n", + " 3485.050273\n", " \n", " \n", - " 92\n", - " 4240.495294\n", + " 93\n", + " 4172.874548\n", " \n", " \n", - " 98\n", - " 3911.455384\n", + " 104\n", + " 3299.302424\n", " \n", " \n", - " 104\n", - " 3271.202866\n", + " 105\n", + " 3515.687917\n", " \n", " \n", - " 114\n", - " 3244.728549\n", + " 108\n", + " 3405.224618\n", " \n", " \n", - " 115\n", - " 3737.374636\n", + " 113\n", + " 4209.140425\n", " \n", " \n", - " 118\n", - " 3485.95604\n", + " 130\n", + " 4197.905737\n", " \n", " \n", "\n", @@ -2306,37 +2362,37 @@ "text/plain": [ " predicted_body_mass_g\n", "penguin_id \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", + "3 3394.118128\n", + "8 4048.685642\n", + "17 3976.454093\n", + "23 3541.582194\n", + "25 4032.844186\n", + "27 4118.351772\n", + "29 4087.767826\n", + "34 3183.755249\n", + "35 3418.802274\n", + "39 3519.186468\n", + "51 3398.135365\n", + "52 3223.615957\n", + "60 3445.014718\n", + "61 3505.638864\n", + "64 3515.905786\n", + "65 4028.363185\n", + "67 4159.993943\n", + "83 3348.16883\n", + "85 3485.050273\n", + "93 4172.874548\n", + "104 3299.302424\n", + "105 3515.687917\n", + "108 3405.224618\n", + "113 4209.140425\n", + "130 4197.905737\n", "...\n", "\n", "[67 rows x 1 columns]" ] }, - "execution_count": 8, + "execution_count": 25, "metadata": {}, "output_type": "execute_result" } @@ -2367,18 +2423,32 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 26, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "fa6bd24b44cb42ec946e262ac2f25d09", + "model_id": "d7a16e04253a42b7a5ce247d8f63b656", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "HTML(value='Query job 6f19614c-82c0-4f8b-b74b-9d91a894efdd is RUNNING. \n", " \n", " \n", - " 9\n", - " 4\n", + " 3\n", + " 3\n", " \n", " \n", - " 12\n", - " 4\n", + " 8\n", + " 3\n", " \n", " \n", - " 13\n", - " 2\n", + " 17\n", + " 3\n", " \n", " \n", - " 19\n", - " 4\n", + " 23\n", + " 1\n", " \n", " \n", - " 22\n", - " 4\n", + " 25\n", + " 3\n", " \n", " \n", - " 23\n", - " 4\n", + " 27\n", + " 3\n", " \n", " \n", - " 25\n", - " 2\n", + " 29\n", + " 3\n", " \n", " \n", - " 28\n", - " 2\n", + " 34\n", + " 3\n", " \n", " \n", - " 30\n", - " 2\n", + " 35\n", + " 1\n", " \n", " \n", - " 38\n", - " 4\n", + " 39\n", + " 3\n", " \n", " \n", - " 41\n", - " 4\n", + " 51\n", + " 1\n", " \n", " \n", - " 45\n", - " 2\n", + " 52\n", + " 3\n", " \n", " \n", - " 49\n", - " 4\n", + " 60\n", + " 3\n", " \n", " \n", - " 50\n", - " 2\n", + " 61\n", + " 3\n", " \n", " \n", - " 62\n", - " 4\n", + " 64\n", + " 1\n", " \n", " \n", " 65\n", - " 4\n", + " 1\n", " \n", " \n", - " 66\n", - " 2\n", + " 67\n", + " 3\n", " \n", " \n", " 83\n", - " 2\n", + " 3\n", " \n", " \n", - " 87\n", - " 2\n", + " 85\n", + " 1\n", " \n", " \n", - " 92\n", - " 2\n", + " 93\n", + " 1\n", " \n", " \n", - " 98\n", - " 4\n", + " 104\n", + " 3\n", " \n", " \n", - " 104\n", - " 2\n", + " 105\n", + " 1\n", " \n", " \n", - " 114\n", - " 2\n", + " 108\n", + " 3\n", " \n", " \n", - " 115\n", - " 2\n", + " 113\n", + " 3\n", " \n", " \n", - " 118\n", - " 4\n", + " 130\n", + " 1\n", " \n", " \n", "\n", @@ -2577,37 +2647,37 @@ "text/plain": [ " CENTROID_ID\n", "penguin_id \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", + "3 3\n", + "8 3\n", + "17 3\n", + "23 1\n", + "25 3\n", + "27 3\n", + "29 3\n", + "34 3\n", + "35 1\n", + "39 3\n", + "51 1\n", + "52 3\n", + "60 3\n", + "61 3\n", + "64 1\n", + "65 1\n", + "67 3\n", + "83 3\n", + "85 1\n", + "93 1\n", + "104 3\n", + "105 1\n", + "108 3\n", + "113 3\n", + "130 1\n", "...\n", "\n", "[67 rows x 1 columns]" ] }, - "execution_count": 9, + "execution_count": 26, "metadata": {}, "output_type": "execute_result" } @@ -2634,7 +2704,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 27, "metadata": {}, "outputs": [ { @@ -2651,7 +2721,7 @@ " ('linreg', LinearRegression())])" ] }, - "execution_count": 10, + "execution_count": 27, "metadata": {}, "output_type": "execute_result" } @@ -2678,18 +2748,18 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 28, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "f0465986682742af92759c3f5fce96e0", + "model_id": "887bf58cebf14bdba95db828390fd33d", "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", - " 9\n", - " 4295.328991\n", + " 3\n", + " 3394.116212\n", " \n", " \n", - " 12\n", - " 3338.434943\n", + " 8\n", + " 4048.683645\n", " \n", " \n", - " 13\n", - " 3201.813783\n", + " 17\n", + " 3976.452358\n", " \n", " \n", - " 19\n", - " 3982.807707\n", + " 23\n", + " 3541.580346\n", " \n", " \n", - " 22\n", - " 3538.604385\n", + " 25\n", + " 4032.842027\n", " \n", " \n", - " 23\n", - " 3613.496641\n", + " 27\n", + " 4118.34983\n", " \n", " \n", - " 25\n", - " 4009.753161\n", + " 29\n", + " 4087.765797\n", " \n", " \n", - " 28\n", - " 4240.509087\n", + " 34\n", + " 3183.75379\n", " \n", " \n", - " 30\n", - " 4028.897875\n", + " 35\n", + " 3418.800633\n", " \n", " \n", - " 38\n", - " 4206.80377\n", + " 39\n", + " 3519.18471\n", " \n", " \n", - " 41\n", - " 3736.219256\n", + " 51\n", + " 3398.133564\n", " \n", " \n", - " 45\n", - " 4383.362136\n", + " 52\n", + " 3223.614107\n", " \n", " \n", - " 49\n", - " 3125.766474\n", + " 60\n", + " 3445.012713\n", " \n", " \n", - " 50\n", - " 3149.281322\n", + " 61\n", + " 3505.637004\n", " \n", " \n", - " 62\n", - " 3531.688645\n", + " 64\n", + " 3515.903779\n", " \n", " \n", " 65\n", - " 4073.894238\n", + " 4028.361259\n", " \n", " \n", - " 66\n", - " 4160.803738\n", + " 67\n", + " 4159.991956\n", " \n", " \n", " 83\n", - " 3272.389735\n", + " 3348.167212\n", " \n", " \n", - " 87\n", - " 3989.358086\n", + " 85\n", + " 3485.048557\n", " \n", " \n", - " 92\n", - " 4240.488891\n", + " 93\n", + " 4172.872284\n", " \n", " \n", - " 98\n", - " 3911.449023\n", + " 104\n", + " 3299.300454\n", " \n", " \n", - " 104\n", - " 3271.196535\n", + " 105\n", + " 3515.68617\n", " \n", " \n", - " 114\n", - " 3244.722283\n", + " 108\n", + " 3405.222757\n", " \n", " \n", - " 115\n", - " 3737.368277\n", + " 113\n", + " 4209.13832\n", " \n", " \n", - " 118\n", - " 3485.949702\n", + " 130\n", + " 4197.90382\n", " \n", " \n", "\n", @@ -2888,37 +2972,37 @@ "text/plain": [ " predicted_body_mass_g\n", "penguin_id \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", + "3 3394.116212\n", + "8 4048.683645\n", + "17 3976.452358\n", + "23 3541.580346\n", + "25 4032.842027\n", + "27 4118.34983\n", + "29 4087.765797\n", + "34 3183.75379\n", + "35 3418.800633\n", + "39 3519.18471\n", + "51 3398.133564\n", + "52 3223.614107\n", + "60 3445.012713\n", + "61 3505.637004\n", + "64 3515.903779\n", + "65 4028.361259\n", + "67 4159.991956\n", + "83 3348.167212\n", + "85 3485.048557\n", + "93 4172.872284\n", + "104 3299.300454\n", + "105 3515.68617\n", + "108 3405.222757\n", + "113 4209.13832\n", + "130 4197.90382\n", "...\n", "\n", "[67 rows x 1 columns]" ] }, - "execution_count": 11, + "execution_count": 28, "metadata": {}, "output_type": "execute_result" } @@ -2950,18 +3034,18 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 29, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "e6409bce719940f4ae326a8b18871d9b", + "model_id": "2d32081be31f44abb8de67e2209d76cd", "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", - " 241.640738\n", - " 90117.84266\n", - " 0.005652\n", - " 200.718678\n", - " 0.8727\n", - " 0.878359\n", + " 229.48269\n", + " 82962.794947\n", + " 0.004248\n", + " 206.728384\n", + " 0.88633\n", + " 0.892953\n", " \n", " \n", "\n", @@ -3069,15 +3167,15 @@ ], "text/plain": [ " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", - "0 241.640738 90117.84266 0.005652 \n", + "0 229.48269 82962.794947 0.004248 \n", "\n", " median_absolute_error r2_score explained_variance \n", - "0 200.718678 0.8727 0.878359 \n", + "0 206.728384 0.88633 0.892953 \n", "\n", "[1 rows x 6 columns]" ] }, - "execution_count": 12, + "execution_count": 29, "metadata": {}, "output_type": "execute_result" } @@ -3097,18 +3195,18 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 30, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "e56e2cc197894ee9acc448a8c12e8a30", + "model_id": "f32692d89f00406499f4ea5aa55268fb", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "HTML(value='Query job 929c826c-1051-47fc-9256-546f4ef11c32 is DONE. 31.7 kB processed. Date: Fri, 20 Oct 2023 22:27:22 -0700 Subject: [PATCH 096/197] docs: add runnable code samples for reading methods (#125) * docs: add runnable and testable I/O code samples * docs: add runnable and testable reading methods code snippets * fix: assign a df and show the first 2 rows * address comments --- bigframes/session/__init__.py | 100 ++++++++++++++++++ .../bigframes_vendored/pandas/io/gbq.py | 24 +++-- .../bigframes_vendored/pandas/io/parquet.py | 14 +++ .../pandas/io/parsers/readers.py | 30 +++++- .../bigframes_vendored/pandas/io/pickle.py | 18 ++++ 5 files changed, 179 insertions(+), 7 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 2f001d7d49a..5ec3da1a5a4 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -342,6 +342,51 @@ def read_gbq_query( ``row_number() over ()`` if there is no natural unique index or you want to preserve ordering. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Simple query input: + + >>> df = bpd.read_gbq_query(''' + ... SELECT + ... pitcherFirstName, + ... pitcherLastName, + ... pitchSpeed, + ... FROM `bigquery-public-data.baseball.games_wide` + ... ''') + >>> df.head(2) + pitcherFirstName pitcherLastName pitchSpeed + 0 0 + 1 0 + + [2 rows x 3 columns] + + Preserve ordering in a query input. + + >>> df = bpd.read_gbq_query(''' + ... 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") + >>> df.head(2) + pitcherFirstName pitcherLastName averagePitchSpeed + rowindex + 1 Albertin Chapman 96.514113 + 2 Zachary Britton 94.591039 + + [2 rows x 3 columns] + See also: :meth:`Session.read_gbq`. """ # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so @@ -405,6 +450,25 @@ def read_gbq_table( ) -> dataframe.DataFrame: """Turn a BigQuery table into a DataFrame. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Read a whole table, with arbitrary ordering or ordering corresponding to the primary key(s). + + >>> df = bpd.read_gbq_table("bigquery-public-data.ml_datasets.penguins") + >>> df.head(2) + species island culmen_length_mm \\ + 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 + 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 + + culmen_depth_mm flipper_length_mm body_mass_g sex + 0 18.4 184.0 3475.0 FEMALE + 1 19.1 184.0 4650.0 MALE + + [2 rows x 7 columns] + See also: :meth:`Session.read_gbq`. """ # NOTE: This method doesn't (yet) exist in pandas or pandas-gbq, so @@ -792,6 +856,16 @@ def _read_ibis( def read_gbq_model(self, model_name: str): """Loads a BigQuery ML model from BigQuery. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Read an existing BigQuery ML model. + + >>> model_name = "bigframes-dev.bqml_tutorial.penguins_model" + >>> model = bpd.read_gbq_model(model_name) + Args: model_name (str): the model's name in BigQuery in the format @@ -815,6 +889,22 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame The pandas DataFrame will be persisted as a temporary BigQuery table, which can be automatically recycled after the Session is closed. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import pandas as pd + >>> bpd.options.display.progress_bar = None + + >>> d = {'col1': [1, 2], 'col2': [3, 4]} + >>> pandas_df = pd.DataFrame(data=d) + >>> df = bpd.read_pandas(pandas_df) + >>> df + col1 col2 + 0 1 3 + 1 2 4 + + [2 rows x 2 columns] + Args: pandas_dataframe (pandas.DataFrame): a pandas DataFrame object to be loaded. @@ -1365,6 +1455,16 @@ def read_gbq_function( The return type of the function must be explicitly specified in the function's original definition even if not otherwise required. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> function_name = "bqutil.fn.cw_lower_case_ascii_only" + >>> func = bpd.read_gbq_function(function_name=function_name) + >>> func.bigframes_remote_function + 'bqutil.fn.cw_lower_case_ascii_only' + Args: function_name (str): the function's name in BigQuery in the format diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 8919f4ed16d..575c5016187 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -42,9 +42,23 @@ def read_gbq( >>> import bigframes.pandas as bpd >>> bpd.options.display.progress_bar = None + If the input is a table ID: + + >>> df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + >>> df.head(2) + species island culmen_length_mm \\ + 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 + 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 + + culmen_depth_mm flipper_length_mm body_mass_g sex + 0 18.4 184.0 3475.0 FEMALE + 1 19.1 184.0 4650.0 MALE + + [2 rows x 7 columns] + Preserve ordering in a query input. - >>> bpd.read_gbq(''' + >>> df = bpd.read_gbq(''' ... SELECT ... -- Instead of an ORDER BY clause on the query, use ... -- ROW_NUMBER() to create an ordered DataFrame. @@ -57,16 +71,14 @@ def read_gbq( ... FROM `bigquery-public-data.baseball.games_wide` ... WHERE year = 2016 ... GROUP BY pitcherFirstName, pitcherLastName - ... ''', index_col="rowindex").head(n=5) + ... ''', index_col="rowindex") + >>> df.head(2) 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] + [2 rows x 3 columns] Args: query_or_table (str): diff --git a/third_party/bigframes_vendored/pandas/io/parquet.py b/third_party/bigframes_vendored/pandas/io/parquet.py index 9aed9af5a82..f97bd386a40 100644 --- a/third_party/bigframes_vendored/pandas/io/parquet.py +++ b/third_party/bigframes_vendored/pandas/io/parquet.py @@ -17,6 +17,20 @@ def read_parquet( Instead, set a serialized index column as the index and sort by that in the resulting DataFrame. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> gcs_path = "gs://cloud-samples-data/bigquery/us-states/us-states.parquet" + >>> df = bpd.read_parquet(path=gcs_path) + >>> df.head(2) + name post_abbr + 0 Alabama AL + 1 Alaska AK + + [2 rows x 2 columns] + Args: path (str): Local or Cloud Storage path to Parquet file. diff --git a/third_party/bigframes_vendored/pandas/io/parsers/readers.py b/third_party/bigframes_vendored/pandas/io/parsers/readers.py index d19a92ecdf8..e8ed6182a6f 100644 --- a/third_party/bigframes_vendored/pandas/io/parsers/readers.py +++ b/third_party/bigframes_vendored/pandas/io/parsers/readers.py @@ -55,6 +55,20 @@ def read_csv( file. Instead, set a serialized index column as the index and sort by that in the resulting DataFrame. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> gcs_path = "gs://cloud-samples-data/bigquery/us-states/us-states.csv" + >>> df = bpd.read_csv(filepath_or_buffer=gcs_path) + >>> df.head(2) + name post_abbr + 0 Alabama AL + 1 Alaska AK + + [2 rows x 2 columns] + Args: filepath_or_buffer (str): A local or Google Cloud Storage (`gs://`) path with `engine="bigquery"` @@ -64,7 +78,7 @@ def read_csv( 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 + 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 (Optional[int], default 0): @@ -146,6 +160,20 @@ def read_json( file. Instead, set a serialized index column as the index and sort by that in the resulting DataFrame. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> gcs_path = "gs://bigframes-dev-testing/sample1.json" + >>> df = bpd.read_json(path_or_buf=gcs_path, lines=True, orient="records") + >>> df.head(2) + id name + 0 1 Alice + 1 2 Bob + + [2 rows x 2 columns] + 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"` diff --git a/third_party/bigframes_vendored/pandas/io/pickle.py b/third_party/bigframes_vendored/pandas/io/pickle.py index 71b31956a07..053ba4871c1 100644 --- a/third_party/bigframes_vendored/pandas/io/pickle.py +++ b/third_party/bigframes_vendored/pandas/io/pickle.py @@ -25,6 +25,24 @@ def read_pickle( 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. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> gcs_path = "gs://bigframes-dev-testing/test_pickle.pkl" + >>> df = bpd.read_pickle(filepath_or_buffer=gcs_path) + >>> df.head(2) + species island culmen_length_mm \\ + 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 + 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 + + culmen_depth_mm flipper_length_mm body_mass_g sex + 0 18.4 184.0 3475.0 FEMALE + 1 19.1 184.0 4650.0 MALE + + [2 rows x 7 columns] + Args: filepath_or_buffer (str, path object, or file-like object): String, path object (implementing os.PathLike[str]), or file-like object From 02984a4530d63e00dd628705149a1fe788cc263e Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Sat, 21 Oct 2023 10:33:03 -0700 Subject: [PATCH 097/197] test: add sample code for use BigFrames developer guide (#118) * test: add sample code for use BigFrames developer guide --- .../snippets/load_data_from_bigquery_test.py | 24 +++++++++++++ samples/snippets/load_data_from_csv_test.py | 25 ++++++++++++++ samples/snippets/pandas_methods_test.py | 34 +++++++++++++++++++ samples/snippets/set_options_test.py | 34 +++++++++++++++++++ 4 files changed, 117 insertions(+) create mode 100644 samples/snippets/load_data_from_bigquery_test.py create mode 100644 samples/snippets/load_data_from_csv_test.py create mode 100644 samples/snippets/pandas_methods_test.py create mode 100644 samples/snippets/set_options_test.py diff --git a/samples/snippets/load_data_from_bigquery_test.py b/samples/snippets/load_data_from_bigquery_test.py new file mode 100644 index 00000000000..e4c65688bdd --- /dev/null +++ b/samples/snippets/load_data_from_bigquery_test.py @@ -0,0 +1,24 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_bigquery(): + # [START bigquery_dataframes_load_data_from_bigquery] + # Create a DataFrame from a BigQuery table: + import bigframes.pandas as bpd + + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + # [END bigquery_dataframes_load_data_from_bigquery] + assert bq_df is not None diff --git a/samples/snippets/load_data_from_csv_test.py b/samples/snippets/load_data_from_csv_test.py new file mode 100644 index 00000000000..31ab9255bf4 --- /dev/null +++ b/samples/snippets/load_data_from_csv_test.py @@ -0,0 +1,25 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_csv(): + # [START bigquery_dataframes_load_data_from_csv] + import bigframes.pandas as bpd + + filepath_or_buffer = "gs://cloud-samples-data/bigquery/us-states/us-states.csv" + df_from_gcs = bpd.read_csv(filepath_or_buffer) + # Display the first few rows of the DataFrame: + df_from_gcs.head() + # [END bigquery_dataframes_load_data_from_csv] + assert df_from_gcs is not None diff --git a/samples/snippets/pandas_methods_test.py b/samples/snippets/pandas_methods_test.py new file mode 100644 index 00000000000..1f472d63466 --- /dev/null +++ b/samples/snippets/pandas_methods_test.py @@ -0,0 +1,34 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_pandas_methods(): + # [START bigquery_dataframes_pandas_methods] + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Inspect one of the columns (or series) of the DataFrame: + bq_df["body_mass_g"].head(10) + + # Compute the mean of this series: + average_body_mass = bq_df["body_mass_g"].mean() + print(f"average_body_mass: {average_body_mass}") + + # Calculate the mean body_mass_g by species using the groupby operation: + bq_df["body_mass_g"].groupby(by=bq_df["species"]).mean().head() + # [END bigquery_dataframes_pandas_methods] + assert average_body_mass is not None diff --git a/samples/snippets/set_options_test.py b/samples/snippets/set_options_test.py new file mode 100644 index 00000000000..ef6f41ce541 --- /dev/null +++ b/samples/snippets/set_options_test.py @@ -0,0 +1,34 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_set_options(): + # Close the session before resetting the options + import bigframes.pandas as bpd + + bpd.close_session() + + # [START bigquery_dataframes_set_options] + import bigframes.pandas as bpd + + PROJECT_ID = "bigframes-dec" # @param {type:"string"} + REGION = "US" # @param {type:"string"} + + # Set BigQuery DataFrames options + bpd.options.bigquery.project = PROJECT_ID + bpd.options.bigquery.location = REGION + + # [END bigquery_dataframes_set_options] + assert bpd.options.bigquery.project == PROJECT_ID + assert bpd.options.bigquery.location == REGION From b17e1f43cd0f7567bc5b59b0e916cd20528312b3 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 23 Oct 2023 12:05:58 -0500 Subject: [PATCH 098/197] fix: expose `bigframes.pandas.reset_session` as a public API (#128) --- bigframes/pandas/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 971d40f801e..5c1928e6f04 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -484,4 +484,5 @@ def read_gbq_function(function_name: str): # Session management APIs "get_global_session", "close_session", + "reset_session", ] From f9ba28c6a6ab1ceaeecd70f7b5a87ec7c404ed13 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 23 Oct 2023 15:30:43 -0500 Subject: [PATCH 099/197] chore: remove unused reference to `THIRD_PARTY_NOTICES` from "nightly" (#130) --- .kokoro/release-nightly.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh index 5cc1275308a..0751cf2502c 100755 --- a/.kokoro/release-nightly.sh +++ b/.kokoro/release-nightly.sh @@ -93,7 +93,6 @@ for gcs_path in gs://vertex_sdk_private_releases/bigframe/ \ 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 -r -v "notebooks/" ${gcs_path}notebooks/ done From 386f35d2840ab677bfb83f9a0b6f3c8de06e78e8 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 24 Oct 2023 10:44:13 -0500 Subject: [PATCH 100/197] refactor: move DDL gen to `bigframes.session._io`, add missing test (#131) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Towards internal issue 280662868 🦕 --- bigframes/dataframe.py | 14 +++++++------- bigframes/session/__init__.py | 2 +- bigframes/session/_io/__init__.py | 13 +++++++++++++ bigframes/{core/io.py => session/_io/bigquery.py} | 0 .../test_io.py => session/test_io_bigquery.py} | 9 +++++---- 5 files changed, 26 insertions(+), 12 deletions(-) create mode 100644 bigframes/session/_io/__init__.py rename bigframes/{core/io.py => session/_io/bigquery.py} (100%) rename tests/unit/{core/test_io.py => session/test_io_bigquery.py} (93%) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 32a2908a421..5c0d9b78e11 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -46,7 +46,6 @@ import bigframes.core.guid import bigframes.core.indexers as indexers import bigframes.core.indexes as indexes -import bigframes.core.io import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.core.window @@ -56,6 +55,7 @@ import bigframes.operations.aggregations as agg_ops import bigframes.series import bigframes.series as bf_series +import bigframes.session._io.bigquery import third_party.bigframes_vendored.pandas.core.frame as vendored_pandas_frame import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing @@ -2201,9 +2201,9 @@ def to_csv( raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_csv_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_csv_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path_or_buf, field_delimiter=sep, @@ -2243,9 +2243,9 @@ def to_json( ) result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_data_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_data_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path_or_buf, format="JSON", @@ -2319,9 +2319,9 @@ def to_parquet( export_options["compression"] = compression.upper() result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_data_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_data_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path, format="PARQUET", diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 5ec3da1a5a4..db9c5a353ca 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -67,13 +67,13 @@ import bigframes.core as core import bigframes.core.blocks as blocks 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 from bigframes.remote_function import remote_function as bigframes_rf +import bigframes.session._io.bigquery as bigframes_io import bigframes.session.clients import bigframes.version diff --git a/bigframes/session/_io/__init__.py b/bigframes/session/_io/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/bigframes/session/_io/__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/bigframes/core/io.py b/bigframes/session/_io/bigquery.py similarity index 100% rename from bigframes/core/io.py rename to bigframes/session/_io/bigquery.py diff --git a/tests/unit/core/test_io.py b/tests/unit/session/test_io_bigquery.py similarity index 93% rename from tests/unit/core/test_io.py rename to tests/unit/session/test_io_bigquery.py index afb38a5f753..d2255d5edf6 100644 --- a/tests/unit/core/test_io.py +++ b/tests/unit/session/test_io_bigquery.py @@ -18,7 +18,7 @@ import google.cloud.bigquery as bigquery import pytest -import bigframes.core.io +import bigframes.session._io.bigquery def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): @@ -26,7 +26,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): "my-test-project._e8166e0cdb.anonbb92cd" ) - sql = bigframes.core.io.create_snapshot_sql( + sql = bigframes.session._io.bigquery.create_snapshot_sql( table_ref, datetime.datetime.now(datetime.timezone.utc) ) @@ -40,7 +40,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): 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( + sql = bigframes.session._io.bigquery.create_snapshot_sql( table_ref, datetime.datetime.now(datetime.timezone.utc) ) @@ -101,4 +101,5 @@ def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): ), ) def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str): - pass + sql = bigframes.session._io.bigquery.bq_schema_to_sql(schema) + assert sql == expected From 95bff3f1902bc09dc3310798a42df8ffd31ed8ee Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 24 Oct 2023 15:08:14 -0700 Subject: [PATCH 101/197] fix: use series's own session in series.reindex listlike case (#135) 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/series.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bigframes/series.py b/bigframes/series.py index 84ca2a578f8..49df8ab61ed 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1242,7 +1242,9 @@ def reindex(self, index=None, *, validate: typing.Optional[bool] = None): raise NotImplementedError( "Cannot reindex with index with different nlevels" ) - new_indexer = bigframes.dataframe.DataFrame(index=index)[[]] + new_indexer = bigframes.dataframe.DataFrame( + index=index, session=self._get_block().expr._session + )[[]] # 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)), From 6fea8efac35871985677ebeb948a576e64a1ffa4 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 24 Oct 2023 21:07:15 -0700 Subject: [PATCH 102/197] docs: add runnable code samples for DataFrames I/O methods and property (#129) * docs: add runnable code samples for DataFrames I/O methods and property * fix: expose `bigframes.pandas.reset_session` as a public API (#128) * fix: address the comment * Empty commit * fix: address comments for better visualization of the output * Empty commit --------- Co-authored-by: Tim Swast --- .../bigframes_vendored/pandas/core/frame.py | 172 +++++++++++++++++- 1 file changed, 164 insertions(+), 8 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index a5c12d7b32a..13a81b46454 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -33,7 +33,19 @@ class DataFrame(NDFrame): @property def shape(self) -> tuple[int, int]: - """Return a tuple representing the dimensionality of the DataFrame.""" + """ + Return a tuple representing the dimensionality of the DataFrame. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2, 3], + ... 'col2': [4, 5, 6]}) + >>> df.shape + (3, 2) + """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property @@ -44,14 +56,14 @@ def axes(self) -> list: It has the row axis labels and column axis labels as the only members. They are returned in that order. - Examples + **Examples:** - .. code-block:: + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None - df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) - df.axes - [RangeIndex(start=0, stop=2, step=1), Index(['col1', 'col2'], - dtype='object')] + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.axes[1:] + [Index(['col1', 'col2'], dtype='object')] """ return [self.index, self.columns] @@ -59,6 +71,16 @@ def axes(self) -> list: def values(self) -> np.ndarray: """Return the values of DataFrame in the form of a NumPy array. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.values + array([[1, 3], + [2, 4]], dtype=object) + Args: dytype (default None): The dtype to pass to `numpy.asarray()`. @@ -76,6 +98,16 @@ def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarra """ Convert the DataFrame to a NumPy array. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.to_numpy() + array([[1, 3], + [2, 4]], dtype=object) + Args: dtype (None): The dtype to pass to `numpy.asarray()`. @@ -101,6 +133,15 @@ def to_gbq( ) -> None: """Write a DataFrame to a BigQuery table. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> # destination_table = PROJECT_ID + "." + DATASET_ID + "." + TABLE_NAME + >>> df.to_gbq("bigframes-dev.birds.test-numbers", if_exists="replace") + Args: destination_table (str): Name of table to be written, in the form ``dataset.tablename`` @@ -137,6 +178,15 @@ def to_parquet( This function writes the dataframe as a `parquet file `_ to Cloud Storage. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> gcs_bucket = "gs://bigframes-dev-testing/sample_parquet*.parquet" + >>> df.to_parquet(path=gcs_bucket) + Args: path (str): Destination URI(s) of Cloud Storage files(s) to store the extracted dataframe @@ -171,6 +221,35 @@ def to_dict( The type of the key-value pairs can be customized with the parameters (see below). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.to_dict() + {'col1': {0: 1, 1: 2}, 'col2': {0: 3, 1: 4}} + + You can specify the return orientation. + + >>> df.to_dict('series') + {'col1': 0 1 + 1 2 + Name: col1, dtype: Int64, + 'col2': 0 3 + 1 4 + Name: col2, dtype: Int64} + + >>> df.to_dict('split') + {'index': [0, 1], 'columns': ['col1', 'col2'], 'data': [[1, 3], [2, 4]]} + + >>> df.to_dict("tight") + {'index': [0, 1], + 'columns': ['col1', 'col2'], + 'data': [[1, 3], [2, 4]], + 'index_names': [None], + 'column_names': [None]} + Args: orient (str {'dict', 'list', 'series', 'split', 'tight', 'records', 'index'}): Determines the type of the values of the dictionary. @@ -213,6 +292,15 @@ def to_excel(self, excel_writer, sheet_name: str = "Sheet1", **kwargs) -> None: Note that creating an `ExcelWriter` object with a file name that already exists will result in the contents of the existing file being erased. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import tempfile + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.to_excel(tempfile.TemporaryFile()) + Args: excel_writer (path-like, file-like, or ExcelWriter object): File path or existing ExcelWriter. @@ -231,6 +319,23 @@ def to_latex( into a main LaTeX document or read from an external file with ``\input{{table.tex}}``. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> print(df.to_latex()) + \begin{tabular}{lrr} + \toprule + & col1 & col2 \\ + \midrule + 0 & 1 & 3 \\ + 1 & 2 & 4 \\ + \bottomrule + \end{tabular} + + Args: buf (str, Path or StringIO-like, optional, default None): Buffer to write to. If None, the output is returned as a string. @@ -253,6 +358,16 @@ def to_records( Index will be included as the first field of the record array if requested. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.to_records() + rec.array([(0, 1, 3), (1, 2, 4)], + dtype=[('index', 'O'), ('col1', 'O'), ('col2', 'O')]) + Args: index (bool, default True): Include index in resulting record array, stored in 'index' @@ -298,6 +413,17 @@ def to_string( ): """Render a DataFrame to a console-friendly tabular output. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> print(df.to_string()) + col1 col2 + 0 1 3 + 1 2 4 + Args: buf (str, Path or StringIO-like, optional, default None): Buffer to write to. If None, the output is returned as a string. @@ -363,6 +489,18 @@ def to_markdown( ): """Print DataFrame in Markdown-friendly format. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> print(df.to_markdown()) + | | col1 | col2 | + |---:|-------:|-------:| + | 0 | 1 | 3 | + | 1 | 2 | 4 | + Args: buf (str, Path or StringIO-like, optional, default None): Buffer to write to. If None, the output is returned as a string. @@ -371,7 +509,7 @@ def to_markdown( index (bool, optional, default True): Add index (row) labels. **kwargs - These parameters will be passed to `tabulate `_. + These parameters will be passed to `tabulate `_. Returns: DataFrame in Markdown-friendly format. @@ -381,6 +519,15 @@ def to_markdown( def to_pickle(self, path, **kwargs) -> None: """Pickle (serialize) object to file. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> gcs_bucket = "gs://bigframes-dev-testing/sample_pickle_gcs.pkl" + >>> df.to_pickle(path=gcs_bucket) + Args: path (str): File path where the pickled object will be stored. @@ -391,6 +538,15 @@ def to_orc(self, path=None, **kwargs) -> bytes | None: """ Write a DataFrame to the ORC format. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> import tempfile + >>> df.to_orc(tempfile.TemporaryFile()) + Args: path (str, file-like object or None, default None): If a string, it will be used as Root Directory path From 05d7618c50acf7b7f9d73e02a8870c8eac910aab Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Wed, 25 Oct 2023 15:52:16 -0700 Subject: [PATCH 103/197] test: allow for alternative PCA solutions in tests (#143) 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 # 🦕 --- tests/system/large/ml/test_decomposition.py | 4 ++-- tests/system/large/ml/test_pipeline.py | 14 ++++++++++---- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/system/large/ml/test_decomposition.py b/tests/system/large/ml/test_decomposition.py index 460f07b8168..a7049d4c18e 100644 --- a/tests/system/large/ml/test_decomposition.py +++ b/tests/system/large/ml/test_decomposition.py @@ -67,8 +67,8 @@ def test_decomposition_configure_fit_score_predict( index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"), ) pd.testing.assert_frame_equal( - result.sort_index(), - expected, + abs(result.sort_index()), # results may differ by a minus sign + abs(expected), check_exact=False, rtol=0.1, ) diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 9294740dd67..6874a9f301a 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -431,10 +431,16 @@ def test_pipeline_PCA_fit_score_predict(session, penguins_df_default_index): 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, + abs( # results may differ by a minus sign + predictions[ + [ + "principal_component_1", + "principal_component_2", + "principal_component_3", + ] + ] + ), + abs(expected), check_exact=False, rtol=0.1, ) From 1641aff37d601b47e0bc4f25ff148be4f718bd1a Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 26 Oct 2023 00:16:15 +0000 Subject: [PATCH 104/197] ci: Disable presubmit LLM tests temporarily (#144) 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 # 🦕 --- tests/system/small/ml/test_llm.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index b7257dde1be..a801c36c830 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -26,6 +26,9 @@ def test_create_text_generator_model(palm2_text_generator_model): assert palm2_text_generator_model._bqml_model is not None +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @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 @@ -48,6 +51,9 @@ def test_create_text_generator_model_default_session(bq_connection, llm_text_pan assert all(series.str.len() > 20) +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_create_text_generator_model_default_connection(llm_text_pandas_df): from bigframes import _config @@ -74,6 +80,9 @@ def test_create_text_generator_model_default_connection(llm_text_pandas_df): # Marked as flaky only because BQML LLM is in preview, the service only has limited capacity, not stable enough. +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_default_params_success( palm2_text_generator_model, llm_text_df @@ -85,6 +94,9 @@ def test_text_generator_predict_default_params_success( assert all(series.str.len() > 20) +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_series_default_params_success( palm2_text_generator_model, llm_text_df @@ -96,6 +108,9 @@ def test_text_generator_predict_series_default_params_success( assert all(series.str.len() > 20) +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_arbitrary_col_label_success( palm2_text_generator_model, llm_text_df @@ -108,6 +123,9 @@ def test_text_generator_predict_arbitrary_col_label_success( assert all(series.str.len() > 20) +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_with_params_success( palm2_text_generator_model, llm_text_df @@ -139,6 +157,9 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df @@ -152,6 +173,9 @@ def test_embedding_generator_predict_success( assert value.size == 768 +@pytest.mark.skip( + reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." +) @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_series_success( palm2_embedding_generator_model, llm_text_df From c3b24b59cd02eeef4fab46761faf1699daa90252 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 25 Oct 2023 19:39:22 -0700 Subject: [PATCH 105/197] chore(main): release 0.11.0 (#126) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 20 ++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d9f63d4c66..93ebadb56fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,26 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.11.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.10.0...v0.11.0) (2023-10-26) + + +### Features + +* Add back `reset_session` as an alias for `close_session` ([#124](https://github.com/googleapis/python-bigquery-dataframes/issues/124)) ([694a85a](https://github.com/googleapis/python-bigquery-dataframes/commit/694a85a0ef90d838700014a204d72b23362db1d8)) +* Change `query` parameter to `query_or_table` in `read_gbq` ([#127](https://github.com/googleapis/python-bigquery-dataframes/issues/127)) ([f9bb3c4](https://github.com/googleapis/python-bigquery-dataframes/commit/f9bb3c4bc88c5ba2be6f17e12a0ec4f482ce161f)) + + +### Bug Fixes + +* Expose `bigframes.pandas.reset_session` as a public API ([#128](https://github.com/googleapis/python-bigquery-dataframes/issues/128)) ([b17e1f4](https://github.com/googleapis/python-bigquery-dataframes/commit/b17e1f43cd0f7567bc5b59b0e916cd20528312b3)) +* Use series's own session in series.reindex listlike case ([#135](https://github.com/googleapis/python-bigquery-dataframes/issues/135)) ([95bff3f](https://github.com/googleapis/python-bigquery-dataframes/commit/95bff3f1902bc09dc3310798a42df8ffd31ed8ee)) + + +### Documentation + +* Add runnable code samples for DataFrames I/O methods and property ([#129](https://github.com/googleapis/python-bigquery-dataframes/issues/129)) ([6fea8ef](https://github.com/googleapis/python-bigquery-dataframes/commit/6fea8efac35871985677ebeb948a576e64a1ffa4)) +* Add runnable code samples for reading methods ([#125](https://github.com/googleapis/python-bigquery-dataframes/issues/125)) ([a669919](https://github.com/googleapis/python-bigquery-dataframes/commit/a669919ff25b56156bd70ccd816a0bf19adb48aa)) + ## [0.10.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.9.0...v0.10.0) (2023-10-19) diff --git a/bigframes/version.py b/bigframes/version.py index 7a37ebd2205..18edfa56150 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.10.0" +__version__ = "0.11.0" From ade505cb9fac96985e62f0c546f0df9b65f301c9 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 26 Oct 2023 05:12:13 +0000 Subject: [PATCH 106/197] Revert "ci: Disable presubmit LLM tests temporarily (#144)" (#148) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 1641aff37d601b47e0bc4f25ff148be4f718bd1a, which was merged due to automerge label while still being discussed. 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/ml/test_llm.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index a801c36c830..b7257dde1be 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -26,9 +26,6 @@ def test_create_text_generator_model(palm2_text_generator_model): assert palm2_text_generator_model._bqml_model is not None -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @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 @@ -51,9 +48,6 @@ def test_create_text_generator_model_default_session(bq_connection, llm_text_pan assert all(series.str.len() > 20) -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_create_text_generator_model_default_connection(llm_text_pandas_df): from bigframes import _config @@ -80,9 +74,6 @@ def test_create_text_generator_model_default_connection(llm_text_pandas_df): # Marked as flaky only because BQML LLM is in preview, the service only has limited capacity, not stable enough. -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_default_params_success( palm2_text_generator_model, llm_text_df @@ -94,9 +85,6 @@ def test_text_generator_predict_default_params_success( assert all(series.str.len() > 20) -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_series_default_params_success( palm2_text_generator_model, llm_text_df @@ -108,9 +96,6 @@ def test_text_generator_predict_series_default_params_success( assert all(series.str.len() > 20) -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_arbitrary_col_label_success( palm2_text_generator_model, llm_text_df @@ -123,9 +108,6 @@ def test_text_generator_predict_arbitrary_col_label_success( assert all(series.str.len() > 20) -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_text_generator_predict_with_params_success( palm2_text_generator_model, llm_text_df @@ -157,9 +139,6 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df @@ -173,9 +152,6 @@ def test_embedding_generator_predict_success( assert value.size == 768 -@pytest.mark.skip( - reason="Temporarily disable to validate the hypothesis that LLM capacity is causing the presubmit tests to take long to run." -) @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_series_success( palm2_embedding_generator_model, llm_text_df From bfd49a54e52e2d502345aa8f9b54457c902abf5c Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 26 Oct 2023 11:48:14 -0500 Subject: [PATCH 107/197] refactor: make `to_pandas()` call `to_arrow()` and use local dtypes in DataFrame construction (#132) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Towards internal issue 280662868 🦕 --- bigframes/core/blocks.py | 41 +--- bigframes/core/indexes/index.py | 3 +- bigframes/dtypes.py | 6 + bigframes/session/__init__.py | 10 +- bigframes/session/_io/pandas.py | 77 +++++++ tests/system/small/test_dataframe.py | 10 - tests/system/small/test_series.py | 48 ++++- tests/unit/session/test_io_pandas.py | 296 +++++++++++++++++++++++++++ tests/unit/test_dtypes.py | 57 +++--- 9 files changed, 457 insertions(+), 91 deletions(-) create mode 100644 bigframes/session/_io/pandas.py create mode 100644 tests/unit/session/test_io_pandas.py diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 046d2b3a444..eab46454779 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -28,11 +28,8 @@ from typing import Iterable, List, Optional, Sequence, Tuple import warnings -import geopandas as gpd # type: ignore import google.cloud.bigquery as bigquery -import numpy import pandas as pd -import pyarrow as pa # type: ignore import bigframes.constants as constants import bigframes.core as core @@ -46,6 +43,7 @@ import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops +import bigframes.session._io.pandas import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common # Type constraint for wherever column labels are used @@ -372,34 +370,11 @@ 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) - @classmethod - def _to_dataframe( - cls, result, schema: typing.Mapping[str, bigframes.dtypes.Dtype] - ) -> pd.DataFrame: + def _to_dataframe(self, result) -> pd.DataFrame: """Convert BigQuery data to pandas DataFrame with specific dtypes.""" - dtypes = bigframes.dtypes.to_pandas_dtypes_overrides(result.schema) - df = result.to_dataframe( - dtypes=dtypes, - 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, 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}), - # BigQuery geography type is based on the WGS84 reference ellipsoid. - crs="EPSG:4326", - ) - return df + dtypes = dict(zip(self.index_columns, self.index_dtypes)) + dtypes.update(zip(self.value_columns, self.dtypes)) + return self._expr._session._rows_to_dataframe(result, dtypes) def to_pandas( self, @@ -480,8 +455,7 @@ def _compute_and_count( if sampling_method == _HEAD: total_rows = int(results_iterator.total_rows * fraction) results_iterator.max_results = total_rows - schema = dict(zip(self.value_columns, self.dtypes)) - df = self._to_dataframe(results_iterator, schema) + df = self._to_dataframe(results_iterator) if self.index_columns: df.set_index(list(self.index_columns), inplace=True) @@ -510,8 +484,7 @@ def _compute_and_count( ) else: total_rows = results_iterator.total_rows - schema = dict(zip(self.value_columns, self.dtypes)) - df = self._to_dataframe(results_iterator, schema) + df = self._to_dataframe(results_iterator) if self.index_columns: df.set_index(list(self.index_columns), inplace=True) diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 677bb8529c2..b9ffdff21ec 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -399,9 +399,10 @@ 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) + dtypes = dict(zip(index_columns, self.dtypes)) expr = self._expr.select_columns(index_columns) results, _ = expr.start_query() - df = expr._session._rows_to_dataframe(results) + df = expr._session._rows_to_dataframe(results, dtypes) df = df.set_index(index_columns) index = df.index index.names = list(self._block._index_labels) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index da221a95aca..079f0cc27a2 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -169,6 +169,10 @@ def ibis_dtype_to_bigframes_dtype( if isinstance(ibis_dtype, ibis_dtypes.Struct): return pd.ArrowDtype(ibis_dtype_to_arrow_dtype(ibis_dtype)) + # BigQuery only supports integers of size 64 bits. + if isinstance(ibis_dtype, ibis_dtypes.Integer): + return pd.Int64Dtype() + if ibis_dtype in IBIS_TO_BIGFRAMES: return IBIS_TO_BIGFRAMES[ibis_dtype] elif isinstance(ibis_dtype, ibis_dtypes.Null): @@ -372,6 +376,8 @@ def cast_ibis_value( ibis_dtypes.float64: (ibis_dtypes.string, ibis_dtypes.int64), ibis_dtypes.string: (ibis_dtypes.int64, ibis_dtypes.float64), ibis_dtypes.date: (), + ibis_dtypes.Decimal(precision=38, scale=9): (ibis_dtypes.float64,), + ibis_dtypes.Decimal(precision=76, scale=38): (ibis_dtypes.float64,), ibis_dtypes.time: (), ibis_dtypes.timestamp: (ibis_dtypes.Timestamp(timezone="UTC"),), ibis_dtypes.Timestamp(timezone="UTC"): (ibis_dtypes.timestamp,), diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index db9c5a353ca..af1f70d54d5 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1515,14 +1515,10 @@ def _get_table_size(self, destination_table): return table.num_bytes def _rows_to_dataframe( - self, row_iterator: bigquery.table.RowIterator + self, row_iterator: bigquery.table.RowIterator, dtypes: Dict ) -> 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"), - ) + arrow_table = row_iterator.to_arrow() + return bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) def _start_generic_job(self, job: formatting_helpers.GenericJob): if bigframes.options.display.progress_bar is not None: diff --git a/bigframes/session/_io/pandas.py b/bigframes/session/_io/pandas.py new file mode 100644 index 00000000000..163127b546d --- /dev/null +++ b/bigframes/session/_io/pandas.py @@ -0,0 +1,77 @@ +# 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, Union + +import geopandas # type: ignore +import pandas +import pandas.arrays +import pyarrow # type: ignore +import pyarrow.compute # type: ignore + +import bigframes.constants + + +def arrow_to_pandas( + arrow_table: Union[pyarrow.Table, pyarrow.RecordBatch], dtypes: Dict +): + if len(dtypes) != arrow_table.num_columns: + raise ValueError( + f"Number of types {len(dtypes)} doesn't match number of columns " + f"{arrow_table.num_columns}. {bigframes.constants.FEEDBACK_LINK}" + ) + + serieses = {} + for field, column in zip(arrow_table.schema, arrow_table): + dtype = dtypes[field.name] + + if dtype == geopandas.array.GeometryDtype(): + series = geopandas.GeoSeries.from_wkt( + column, + # BigQuery geography type is based on the WGS84 reference ellipsoid. + crs="EPSG:4326", + ) + elif dtype == pandas.Float64Dtype(): + # Preserve NA/NaN distinction. Note: This is currently needed, even if we use + # nullable Float64Dtype in the types_mapper. See: + # https://github.com/pandas-dev/pandas/issues/55668 + # Regarding type: ignore, this class has been public at this + # location since pandas 1.2.0. See: + # https://pandas.pydata.org/docs/dev/reference/api/pandas.arrays.FloatingArray.html + pd_array = pandas.arrays.FloatingArray( # type: ignore + column.to_numpy(), + pyarrow.compute.is_null(column).to_numpy(), + ) + series = pandas.Series(pd_array, dtype=dtype) + elif dtype == pandas.Int64Dtype(): + # Avoid out-of-bounds errors in Pandas 1.5.x, which incorrectly + # casts to float64 in an intermediate step. + pd_array = pandas.arrays.IntegerArray( + pyarrow.compute.fill_null(column, 0).to_numpy(), + pyarrow.compute.is_null(column).to_numpy(), + ) + series = pandas.Series(pd_array, dtype=dtype) + elif isinstance(dtype, pandas.ArrowDtype): + # Avoid conversion logic if we are backing the pandas Series by the + # arrow array. + series = pandas.Series( + pandas.arrays.ArrowExtensionArray(column), # type: ignore + dtype=dtype, + ) + else: + series = column.to_pandas(types_mapper=lambda _: dtype) + + serieses[field.name] = series + + return pandas.DataFrame(serieses) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 19e50eb06de..84e8def83b3 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2046,16 +2046,6 @@ 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].to_pandas() 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, diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index bd9edbb1ca0..c9510290b66 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -575,7 +575,15 @@ def test_series_int_int_operators_series(scalars_dfs, operator): ) 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() + x_bf = scalars_df[col_x] + y_bf = scalars_df[col_y] + bf_series = getattr(x_bf, method)(y_bf) + # BigQuery's mod functions return [BIG]NUMERIC values unless both arguments are integers. + # https://cloud.google.com/bigquery/docs/reference/standard-sql/mathematical_functions#mod + if x_bf.dtype == pd.Int64Dtype() and y_bf.dtype == pd.Int64Dtype(): + bf_result = bf_series.to_pandas() + else: + bf_result = bf_series.astype("Float64").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) @@ -620,8 +628,20 @@ def test_divmods_series(scalars_dfs, col_x, col_y, method): 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()) + # BigQuery's mod functions return NUMERIC values for non-INT64 inputs. + if bf_div_result.dtype == pd.Int64Dtype(): + pd.testing.assert_series_equal(pd_div_result, bf_div_result.to_pandas()) + else: + pd.testing.assert_series_equal( + pd_div_result, bf_div_result.astype("Float64").to_pandas() + ) + + if bf_mod_result.dtype == pd.Int64Dtype(): + pd.testing.assert_series_equal(pd_mod_result, bf_mod_result.to_pandas()) + else: + pd.testing.assert_series_equal( + pd_mod_result, bf_mod_result.astype("Float64").to_pandas() + ) @pytest.mark.parametrize( @@ -649,8 +669,20 @@ 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()) + # BigQuery's mod functions return NUMERIC values for non-INT64 inputs. + if bf_div_result.dtype == pd.Int64Dtype(): + pd.testing.assert_series_equal(pd_div_result, bf_div_result.to_pandas()) + else: + pd.testing.assert_series_equal( + pd_div_result, bf_div_result.astype("Float64").to_pandas() + ) + + if bf_mod_result.dtype == pd.Int64Dtype(): + pd.testing.assert_series_equal(pd_mod_result, bf_mod_result.to_pandas()) + else: + pd.testing.assert_series_equal( + pd_mod_result, bf_mod_result.astype("Float64").to_pandas() + ) @pytest.mark.parametrize( @@ -1941,12 +1973,6 @@ 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].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 - 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, diff --git a/tests/unit/session/test_io_pandas.py b/tests/unit/session/test_io_pandas.py new file mode 100644 index 00000000000..8b95977ec3d --- /dev/null +++ b/tests/unit/session/test_io_pandas.py @@ -0,0 +1,296 @@ +# 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 +from typing import Dict, Union + +import geopandas # type: ignore +import numpy +import pandas +import pandas.arrays +import pandas.testing +import pyarrow # type: ignore +import pytest + +import bigframes.session._io.pandas + + +@pytest.mark.parametrize( + ("arrow_table", "dtypes", "expected"), + ( + pytest.param( + pyarrow.Table.from_pydict({}), + {}, + pandas.DataFrame(), + id="empty-df", + ), + pytest.param( + pyarrow.Table.from_pydict( + { + "bool": pyarrow.array([None, None, None], type=pyarrow.bool_()), + "float": pyarrow.array([None, None, None], type=pyarrow.float64()), + "int": pyarrow.array([None, None, None], type=pyarrow.int64()), + "string": pyarrow.array([None, None, None], type=pyarrow.string()), + "time": pyarrow.array( + [None, None, None], type=pyarrow.time64("us") + ), + } + ), + { + "bool": "boolean", + "float": pandas.Float64Dtype(), + "int": pandas.Int64Dtype(), + "string": "string[pyarrow]", + "time": pandas.ArrowDtype(pyarrow.time64("us")), + }, + pandas.DataFrame( + { + "bool": pandas.Series([None, None, None], dtype="boolean"), + "float": pandas.Series( + pandas.arrays.FloatingArray( # type: ignore + numpy.array( + [float("nan"), float("nan"), float("nan")], + dtype="float64", + ), + numpy.array([True, True, True], dtype="bool"), + ), + dtype=pandas.Float64Dtype(), + ), + "int": pandas.Series( + [None, None, None], + dtype=pandas.Int64Dtype(), + ), + "string": pandas.Series( + [None, None, None], dtype="string[pyarrow]" + ), + "time": pandas.Series( + [ + None, + None, + None, + ], + dtype=pandas.ArrowDtype(pyarrow.time64("us")), + ), + } + ), + id="nulls-df", + ), + pytest.param( + pyarrow.Table.from_pydict( + { + "date": pyarrow.array( + [ + datetime.date(2023, 8, 29), + None, + datetime.date(2024, 4, 9), + datetime.date(1, 1, 1), + ], + type=pyarrow.date32(), + ), + "datetime": pyarrow.array( + [ + datetime.datetime(2023, 8, 29), + None, + datetime.datetime(2024, 4, 9, 23, 59, 59), + datetime.datetime(1, 1, 1, 0, 0, 0, 1), + ], + type=pyarrow.timestamp("us"), + ), + "string": ["123", None, "abc", "xyz"], + "time": pyarrow.array( + [ + datetime.time(0, 0, 0, 1), + datetime.time(12, 0, 0), + None, + datetime.time(23, 59, 59, 999999), + ], + type=pyarrow.time64("us"), + ), + "timestamp": pyarrow.array( + [ + datetime.datetime(2023, 8, 29), + datetime.datetime(1, 1, 1, 0, 0, 0, 1), + None, + datetime.datetime(2024, 4, 9, 23, 59, 59), + ], + type=pyarrow.timestamp("us", datetime.timezone.utc), + ), + } + ), + { + "date": pandas.ArrowDtype(pyarrow.date32()), + "datetime": pandas.ArrowDtype(pyarrow.timestamp("us")), + "string": "string[pyarrow]", + "time": pandas.ArrowDtype(pyarrow.time64("us")), + "timestamp": pandas.ArrowDtype( + pyarrow.timestamp("us", datetime.timezone.utc) + ), + }, + pandas.DataFrame( + { + "date": pandas.Series( + [ + datetime.date(2023, 8, 29), + None, + datetime.date(2024, 4, 9), + datetime.date(1, 1, 1), + ], + dtype=pandas.ArrowDtype(pyarrow.date32()), + ), + "datetime": pandas.Series( + [ + datetime.datetime(2023, 8, 29), + None, + datetime.datetime(2024, 4, 9, 23, 59, 59), + datetime.datetime(1, 1, 1, 0, 0, 0, 1), + ], + dtype=pandas.ArrowDtype(pyarrow.timestamp("us")), + ), + "string": pandas.Series( + ["123", None, "abc", "xyz"], dtype="string[pyarrow]" + ), + "time": pandas.Series( + [ + datetime.time(0, 0, 0, 1), + datetime.time(12, 0, 0), + None, + datetime.time(23, 59, 59, 999999), + ], + dtype=pandas.ArrowDtype(pyarrow.time64("us")), + ), + "timestamp": pandas.Series( + [ + datetime.datetime(2023, 8, 29), + datetime.datetime(1, 1, 1, 0, 0, 0, 1), + None, + datetime.datetime(2024, 4, 9, 23, 59, 59), + ], + dtype=pandas.ArrowDtype( + pyarrow.timestamp("us", datetime.timezone.utc) + ), + ), + } + ), + id="arrow-dtypes", + ), + pytest.param( + pyarrow.Table.from_pydict( + { + "bool": [True, None, True, False], + "bytes": [b"123", None, b"abc", b"xyz"], + "float": pyarrow.array( + [1.0, None, float("nan"), -1.0], + type=pyarrow.float64(), + ), + "int": pyarrow.array( + [1, None, -1, 2**63 - 1], + type=pyarrow.int64(), + ), + "string": ["123", None, "abc", "xyz"], + } + ), + { + "bool": "boolean", + "bytes": "object", + "float": pandas.Float64Dtype(), + "int": pandas.Int64Dtype(), + "string": "string[pyarrow]", + }, + pandas.DataFrame( + { + "bool": pandas.Series([True, None, True, False], dtype="boolean"), + "bytes": [b"123", None, b"abc", b"xyz"], + "float": pandas.Series( + pandas.arrays.FloatingArray( # type: ignore + numpy.array( + [1.0, float("nan"), float("nan"), -1.0], dtype="float64" + ), + numpy.array([False, True, False, False], dtype="bool"), + ), + dtype=pandas.Float64Dtype(), + ), + "int": pandas.Series( + [1, None, -1, 2**63 - 1], + dtype=pandas.Int64Dtype(), + ), + "string": pandas.Series( + ["123", None, "abc", "xyz"], dtype="string[pyarrow]" + ), + } + ), + id="scalar-dtypes", + ), + pytest.param( + pyarrow.Table.from_pydict( + { + "geocol": [ + "POINT(32 210)", + None, + "LINESTRING(1 1, 2 1, 3.1 2.88, 3 -3)", + ] + } + ), + {"geocol": geopandas.array.GeometryDtype()}, + pandas.DataFrame( + { + "geocol": geopandas.GeoSeries.from_wkt( + ["POINT(32 210)", None, "LINESTRING(1 1, 2 1, 3.1 2.88, 3 -3)"], + crs="EPSG:4326", + ), + } + ), + id="geography-dtype", + ), + ), +) +def test_arrow_to_pandas( + arrow_table: Union[pyarrow.Table, pyarrow.RecordBatch], + dtypes: Dict, + expected: pandas.DataFrame, +): + actual = bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) + pandas.testing.assert_series_equal(actual.dtypes, expected.dtypes) + + # assert_frame_equal is converting to numpy internally, which causes some + # loss of precision with the extreme values in this test. + for column in actual.columns: + assert tuple( + (index, value) if (value is pandas.NA or value == value) else (index, "nan") + for index, value in actual[column].items() + ) == tuple( + (index, value) if (value is pandas.NA or value == value) else (index, "nan") + for index, value in expected[column].items() + ) + + +@pytest.mark.parametrize( + ("arrow_table", "dtypes"), + ( + pytest.param( + pyarrow.Table.from_pydict({"col1": [1], "col2": [2]}), + {"col1": "Int64"}, + id="too-few-dtypes", + ), + pytest.param( + pyarrow.RecordBatch.from_pydict({"col1": [1]}), + {"col1": "Int64", "col2": "string[pyarrow]"}, + id="too-many-dtypes", + ), + ), +) +def test_arrow_to_pandas_wrong_size_dtypes( + arrow_table: Union[pyarrow.Table, pyarrow.RecordBatch], dtypes: Dict +): + with pytest.raises(ValueError, match=f"Number of types {len(dtypes)}"): + bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) diff --git a/tests/unit/test_dtypes.py b/tests/unit/test_dtypes.py index 3baff2e1f58..6ceaaf911b9 100644 --- a/tests/unit/test_dtypes.py +++ b/tests/unit/test_dtypes.py @@ -29,41 +29,42 @@ # 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()), - ( + pytest.param( + ibis_dtypes.Decimal(precision=76, scale=38, nullable=True), + np.dtype("O"), + id="bignumeric", + ), + pytest.param(ibis_dtypes.boolean, pd.BooleanDtype(), id="bool"), + pytest.param(ibis_dtypes.binary, np.dtype("O"), id="bytes"), + pytest.param(ibis_dtypes.date, pd.ArrowDtype(pa.date32()), id="date"), + pytest.param( + ibis_dtypes.Timestamp(), pd.ArrowDtype(pa.timestamp("us")), id="datetime" + ), + pytest.param(ibis_dtypes.float64, pd.Float64Dtype(), id="float"), + pytest.param( ibis_dtypes.GeoSpatial(geotype="geography", srid=4326, nullable=True), gpd.array.GeometryDtype(), + id="geography", ), - (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"))), - ( + pytest.param(ibis_dtypes.int8, pd.Int64Dtype(), id="int8-as-int64"), + pytest.param(ibis_dtypes.int64, pd.Int64Dtype(), id="int64"), + # TODO(tswast): custom dtype (or at least string dtype) for JSON objects + pytest.param(ibis_dtypes.json, np.dtype("O"), id="json"), + pytest.param( + ibis_dtypes.Decimal(precision=38, scale=9, nullable=True), + np.dtype("O"), + id="numeric", + ), + pytest.param( + ibis_dtypes.string, pd.StringDtype(storage="pyarrow"), id="string" + ), + pytest.param(ibis_dtypes.time, pd.ArrowDtype(pa.time64("us")), id="time"), + pytest.param( ibis_dtypes.Timestamp(timezone="UTC"), pd.ArrowDtype(pa.timestamp("us", tz="UTC")), # type: ignore + id="timestamp", ), ], - 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""" From d423e102453c070af64aa37741fb9ff6fb9a6d25 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 26 Oct 2023 17:44:14 +0000 Subject: [PATCH 108/197] test: Log slowest tests durations (#146) 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 # 🦕 --- noxfile.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index 1864da9fe7d..d0bbda80fd4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -305,8 +305,10 @@ def run_system( "py.test", "--quiet", "-n=20", - # Any individual test taking longer than 10 mins will be terminated. + # Any individual test taking longer than 15 mins will be terminated. "--timeout=900", + # Log 20 slowest tests + "--durations=20", f"--junitxml={prefix_name}_{session.python}_sponge_log.xml", ] if print_duration: From 45c617fee7becc42f1c129246ffdc32f3a963f12 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 26 Oct 2023 11:36:14 -0700 Subject: [PATCH 109/197] docs: link to ML.EVALUATE BQML page for score() methods (#137) 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/ml/ensemble.py | 12 ++++++++++++ bigframes/ml/forecasting.py | 6 ++++++ third_party/bigframes_vendored/sklearn/base.py | 14 +++++++++++++- .../bigframes_vendored/sklearn/cluster/_kmeans.py | 9 +++++++-- .../sklearn/decomposition/_pca.py | 8 +++++++- 5 files changed, 45 insertions(+), 4 deletions(-) diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index 113ad872b5c..19ca8608ffb 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -507,6 +507,12 @@ def score( ): """Calculate evaluation metrics of the model. + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#regression_models + for the outputs relevant to this model type. + Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): A BigQuery DataFrame as evaluation data. @@ -676,6 +682,12 @@ def score( ): """Calculate evaluation metrics of the model. + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#classification_models + for the outputs relevant to this model type. + Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): A BigQuery DataFrame as evaluation data. diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index 8a6de1dd81f..8e309d5e736 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -112,6 +112,12 @@ def score( ) -> bpd.DataFrame: """Calculate evaluation metrics of the model. + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#time_series_models + for the outputs relevant to this model type. + Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): A BigQuery DataFrame only contains 1 column as diff --git a/third_party/bigframes_vendored/sklearn/base.py b/third_party/bigframes_vendored/sklearn/base.py index 42868ce51f2..768328e5529 100644 --- a/third_party/bigframes_vendored/sklearn/base.py +++ b/third_party/bigframes_vendored/sklearn/base.py @@ -85,6 +85,12 @@ def score(self, X, y): which is a harsh metric since you require for each sample that each label set be correctly predicted. + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#classification_models + for the outputs relevant to this model type. + Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): DataFrame of shape (n_samples, n_features). Test samples. @@ -105,7 +111,13 @@ class RegressorMixin: _estimator_type = "regressor" def score(self, X, y): - """Return the evaluation metrics of the model. + """Calculate evaluation metrics of the model. + + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#regression_models + for the outputs relevant to this model type. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index ece62dc1470..5369d3662d7 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -12,7 +12,6 @@ # License: BSD 3 clause from abc import ABC -from typing import List, Optional from bigframes import constants from third_party.bigframes_vendored.sklearn.base import BaseEstimator @@ -83,7 +82,13 @@ def score( X, y=None, ): - """Metrics of the model. + """Calculate evaluation metrics of the model. + + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#k-means_models + for the outputs relevant to this model type. Args: X (bigframes.dataframe.DataFrame or bigframes.series.Series): diff --git a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py index 97fee5a501e..011ecc06dd3 100644 --- a/third_party/bigframes_vendored/sklearn/decomposition/_pca.py +++ b/third_party/bigframes_vendored/sklearn/decomposition/_pca.py @@ -55,7 +55,13 @@ def fit(self, X, y=None): raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def score(self, X=None, y=None): - """Return the metrics of the model. + """Calculate evaluation metrics of the model. + + .. note:: + + Output matches that of the BigQuery ML.EVALUTE function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#pca_models + for the outputs relevant to this model type. Args: X (default None): From c639a3657465e2b68a3b93c363bd3ae1e969d2cc Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 26 Oct 2023 12:30:15 -0700 Subject: [PATCH 110/197] feat: populate ibis version in user agent (#140) 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/session/clients.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py index 544f74265fe..e33413002fe 100644 --- a/bigframes/session/clients.py +++ b/bigframes/session/clients.py @@ -28,12 +28,13 @@ import google.cloud.bigquery_storage_v1 import google.cloud.functions_v2 import google.cloud.resourcemanager_v3 +import ibis import pydata_google_auth import bigframes.version _ENV_DEFAULT_PROJECT = "GOOGLE_CLOUD_PROJECT" -_APPLICATION_NAME = f"bigframes/{bigframes.version.__version__}" +_APPLICATION_NAME = f"bigframes/{bigframes.version.__version__} ibis/{ibis.__version__}" _SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] # BigQuery is a REST API, which requires the protocol as part of the URL. From 2ddbf743efc2fd8ffb61ae8d3333fc4b98ce4b55 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 26 Oct 2023 15:12:14 -0500 Subject: [PATCH 111/197] fix: don't override the global logging config (#138) 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/clients.py | 3 --- bigframes/remote_function.py | 5 ----- 2 files changed, 8 deletions(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index 4ba9d93d694..de2421e499f 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -24,9 +24,6 @@ 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__) diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index c82ba840564..a39cd033f69 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -53,11 +53,6 @@ from bigframes import clients import bigframes.constants as constants -# 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 From 27c57255c7fe11e1ef9b9826d988d80fc17442a6 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Thu, 26 Oct 2023 14:01:04 -0700 Subject: [PATCH 112/197] fix: use indexee's session for loc listlike cases (#152) 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/core/indexers.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 4f5a9471b99..d18a0a38efe 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -310,7 +310,9 @@ def _loc_getitem_series_or_dataframe( index_name = temporary_index_names[i] values = [entry[i] for entry in key] index_cols_dict[index_name] = values - keys_df = bigframes.dataframe.DataFrame(index_cols_dict) + keys_df = bigframes.dataframe.DataFrame( + index_cols_dict, session=series_or_dataframe._get_block().expr._session + ) keys_df = keys_df.set_index(temporary_index_names, drop=True) keys_df = keys_df.rename_axis(original_index_names) else: @@ -320,7 +322,10 @@ def _loc_getitem_series_or_dataframe( 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}) + 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 From 8e4451841ba09099b0ed5433f9102511741dfbed Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 26 Oct 2023 14:40:13 -0700 Subject: [PATCH 113/197] feat: add pandas.qcut (#104) 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/core/reshape/__init__.py | 33 ++++++++++++ bigframes/operations/aggregations.py | 51 +++++++++++++++++++ bigframes/pandas/__init__.py | 13 +++++ tests/system/small/test_pandas.py | 25 +++++++++ .../pandas/core/reshape/tile.py | 30 +++++++++++ 5 files changed, 152 insertions(+) diff --git a/bigframes/core/reshape/__init__.py b/bigframes/core/reshape/__init__.py index 339ce7466a1..dc61c3baada 100644 --- a/bigframes/core/reshape/__init__.py +++ b/bigframes/core/reshape/__init__.py @@ -20,6 +20,7 @@ import bigframes.core as core import bigframes.core.utils as utils import bigframes.dataframe +import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops import bigframes.series @@ -118,3 +119,35 @@ def cut( 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()) + + +def qcut( + x: bigframes.series.Series, + q: typing.Union[int, typing.Sequence[float]], + *, + labels: Optional[bool] = None, + duplicates: typing.Literal["drop", "error"] = "error", +) -> bigframes.series.Series: + if isinstance(q, int) and q <= 0: + raise ValueError("`q` 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}" + ) + if duplicates != "drop": + raise NotImplementedError( + f"Only duplicates='drop' is supported in BigQuery DataFrames so far. {constants.FEEDBACK_LINK}" + ) + block = x._block + label = block.col_id_to_label[x._value_column] + block, nullity_id = block.apply_unary_op(x._value_column, ops.notnull_op) + block, result = block.apply_window_op( + x._value_column, + agg_ops.QcutOp(q), + window_spec=core.WindowSpec(grouping_keys=(nullity_id,)), + ) + block, result = block.apply_binary_op( + result, nullity_id, ops.partial_arg3(ops.where_op, None), result_label=label + ) + return bigframes.series.Series(block.select_column(result)) diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py index 23271e82205..465d1887247 100644 --- a/bigframes/operations/aggregations.py +++ b/bigframes/operations/aggregations.py @@ -254,6 +254,53 @@ def handles_ties(self): return True +class QcutOp(WindowOp): + def __init__(self, quantiles: typing.Union[int, typing.Sequence[float]]): + self.name = f"qcut-{quantiles}" + self._quantiles = quantiles + + @numeric_op + def _as_ibis( + self, column: ibis_types.Column, window=None + ) -> ibis_types.IntegerValue: + if isinstance(self._quantiles, int): + quantiles_ibis = dtypes.literal_to_ibis_scalar(self._quantiles) + percent_ranks = typing.cast( + ibis_types.FloatingColumn, + _apply_window_if_present(column.percent_rank(), window), + ) + float_bucket = typing.cast( + ibis_types.FloatingColumn, (percent_ranks * quantiles_ibis) + ) + return float_bucket.ceil().clip(lower=_ibis_num(1)) - _ibis_num(1) + else: + percent_ranks = typing.cast( + ibis_types.FloatingColumn, + _apply_window_if_present(column.percent_rank(), window), + ) + out = ibis.case() + first_ibis_quantile = dtypes.literal_to_ibis_scalar(self._quantiles[0]) + out = out.when(percent_ranks < first_ibis_quantile, None) + for bucket_n in range(len(self._quantiles) - 1): + ibis_quantile = dtypes.literal_to_ibis_scalar( + self._quantiles[bucket_n + 1] + ) + out = out.when( + percent_ranks <= ibis_quantile, + dtypes.literal_to_ibis_scalar(bucket_n, force_dtype=Int64Dtype()), + ) + out = out.else_(None) + return out.end() + + @property + def skips_nulls(self): + return False + + @property + def handles_ties(self): + return True + + class NuniqueOp(AggregateOp): name = "nunique" @@ -491,3 +538,7 @@ def lookup_agg_func(key: str) -> AggregateOp: return _AGGREGATIONS_LOOKUP[key] else: raise ValueError(f"Unrecognize aggregate function: {key}") + + +def _ibis_num(number: float): + return typing.cast(ibis_types.NumericValue, ibis_types.literal(number)) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 5c1928e6f04..8d9726312fd 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -134,6 +134,19 @@ def cut( cut.__doc__ = vendored_pandas_tile.cut.__doc__ +def qcut( + x: bigframes.series.Series, + q: int, + *, + labels: Optional[bool] = None, + duplicates: typing.Literal["drop", "error"] = "error", +) -> bigframes.series.Series: + return bigframes.core.reshape.qcut(x, q, labels=labels, duplicates=duplicates) + + +qcut.__doc__ = vendored_pandas_tile.qcut.__doc__ + + def merge( left: DataFrame, right: DataFrame, diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index a429c6551d6..f8fa78587f9 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -223,3 +223,28 @@ def test_cut(scalars_dfs): bf_result = bf_result.to_pandas() pd_result = pd_result.astype("Int64") pd.testing.assert_series_equal(bf_result, pd_result) + + +@pytest.mark.parametrize( + ("q",), + [ + (1,), + (2,), + (7,), + (32,), + ([0, 0.1, 0.3, 0.4, 0.9, 1.0],), + ([0.5, 0.9],), + ], +) +def test_qcut(scalars_dfs, q): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = pd.qcut( + scalars_pandas_df["float64_col"], q, labels=False, duplicates="drop" + ) + bf_result = bpd.qcut(scalars_df["float64_col"], q, labels=False, duplicates="drop") + + bf_result = bf_result.to_pandas() + pd_result = pd_result.astype("Int64") + + pd.testing.assert_series_equal(bf_result, pd_result) diff --git a/third_party/bigframes_vendored/pandas/core/reshape/tile.py b/third_party/bigframes_vendored/pandas/core/reshape/tile.py index 4f5f2efef0a..24ea655a5fe 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/tile.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/tile.py @@ -65,3 +65,33 @@ def cut( False : returns an ndarray of integers. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + +def qcut(x, q, *, labels=None, duplicates="error"): + """ + Quantile-based discretization function. + + Discretize variable into equal-sized buckets based on rank or based + on sample quantiles. For example 1000 values for 10 quantiles would + produce a Categorical object indicating quantile membership for each data point. + + Args: + x (Series): + The input Series to be binned. Must be 1-dimensional. + q (int or list-like of float): + Number of quantiles. 10 for deciles, 4 for quartiles, etc. Alternately + array of quantiles, e.g. [0, .25, .5, .75, 1.] for quartiles. + labels (None): + Used as labels for the resulting bins. Must be of the same length as + the resulting bins. If False, return only integer indicators of the + bins. If True, raises an error. + duplicates ({default 'raise', 'drop'}, optional): + If bin edges are not unique, raise ValueError or drop non-uniques. + + Returns: + Series: Categorical or Series of integers if labels is False + The return type (Categorical or Series) depends on the input: a Series + of type category if input is a Series else Categorical. Bins are + represented as categories when categorical data is returned. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 5edcd19e6200db9b9ebe3d4945816b3ebf1f7bcd Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 26 Oct 2023 15:20:15 -0700 Subject: [PATCH 114/197] feat: add unstack to series, add level param (#115) 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/core/blocks.py | 26 ++++++++++-- bigframes/dataframe.py | 32 +++++---------- bigframes/series.py | 40 +++++++++++-------- tests/system/conftest.py | 8 +++- tests/system/small/test_dataframe.py | 10 ++++- tests/system/small/test_multiindex.py | 31 ++++++++++++-- .../bigframes_vendored/pandas/core/series.py | 13 ++++++ 7 files changed, 112 insertions(+), 48 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index eab46454779..e8a3968b3db 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -67,6 +67,10 @@ _MONOTONIC_DECREASING = "monotonic_decreasing" +LevelType = typing.Union[str, int] +LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]] + + class BlockHolder(typing.Protocol): """Interface for mutable objects with state represented by a block value object.""" @@ -1423,9 +1427,7 @@ def _get_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") - ) + return pd.MultiIndex.from_frame(pd_values) else: return pd.Index(pd_values.squeeze(axis=1).sort_values(na_position="first")) @@ -1611,6 +1613,24 @@ def cached(self) -> Block: index_labels=self.index_labels, ) + def resolve_index_level(self, level: LevelsType) -> typing.Sequence[str]: + if utils.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.index_columns[level_ref]) + elif isinstance(level_ref, typing.Hashable): + matching_ids = self.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 _is_monotonic( self, column_ids: typing.Union[str, Sequence[str]], increasing: bool ) -> bool: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 5c0d9b78e11..869075a9706 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1038,22 +1038,7 @@ def reorder_levels(self, order: LevelsType, axis: int | str = 0): 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): - 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, 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") - resolved_level_ids.append(matching_ids[0]) - else: - raise ValueError(f"Unexpected level: {level_ref}") - return resolved_level_ids + return self._block.resolve_index_level(level) def rename(self, *, columns: Mapping[blocks.Label, blocks.Label]) -> DataFrame: block = self._block.rename(columns=columns) @@ -1802,20 +1787,25 @@ def _stack_multi(self, level: LevelsType = -1): block = block.stack(levels=len(level)) return DataFrame(block) - def unstack(self): + def unstack(self, level: LevelsType = -1): + if isinstance(level, int) or isinstance(level, str): + level = [level] + block = self._block # Special case, unstack with mono-index transpose into a series if self.index.nlevels == 1: block = block.stack(how="right", levels=self.columns.nlevels) return bigframes.series.Series(block) - # Pivot by last level of index - index_ids = block.index_columns + # Pivot by index levels + unstack_ids = self._resolve_levels(level) block = block.reset_index(drop=False) - block = block.set_index(index_ids[:-1]) + block = block.set_index( + [col for col in self._block.index_columns if col not in unstack_ids] + ) pivot_block = block.pivot( - columns=[index_ids[-1]], + columns=unstack_ids, values=self._block.value_columns, values_in_index=True, ) diff --git a/bigframes/series.py b/bigframes/series.py index 49df8ab61ed..c1914527834 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -352,22 +352,7 @@ def reorder_levels(self, order: LevelsType, axis: int | str = 0): 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, 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") - resolved_level_ids.append(matching_ids[0]) - else: - raise ValueError(f"Unexpected level: {level_ref}") - return resolved_level_ids + return self._block.resolve_index_level(level) def between(self, left, right, inclusive="both"): if inclusive not in ["both", "neither", "left", "right"]: @@ -918,6 +903,29 @@ def argmin(self) -> int: scalars.Scalar, Series(block.select_column(row_nums)).iloc[0] ) + def unstack(self, level: LevelsType = -1): + if isinstance(level, int) or isinstance(level, str): + level = [level] + + block = self._block + + if self.index.nlevels == 1: + raise ValueError("Series must have multi-index to unstack") + + # Pivot by index levels + unstack_ids = self._resolve_levels(level) + block = block.reset_index(drop=False) + block = block.set_index( + [col for col in self._block.index_columns if col not in unstack_ids] + ) + + pivot_block = block.pivot( + columns=unstack_ids, + values=self._block.value_columns, + values_in_index=False, + ) + return bigframes.dataframe.DataFrame(pivot_block) + def idxmax(self) -> blocks.Label: block = self._block.order_by( [ diff --git a/tests/system/conftest.py b/tests/system/conftest.py index cb664302a8f..8885b03d34c 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -400,7 +400,11 @@ 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) + return ( + session.read_gbq(hockey_table_id) + .set_index(["player_name", "season"]) + .sort_index() + ) @pytest.fixture(scope="session") @@ -419,7 +423,7 @@ def hockey_pandas_df() -> pd.DataFrame: "season": pd.Int64Dtype(), }, ) - df.index = df.index.astype("Int64") + df = df.set_index(["player_name", "season"]).sort_index() return df diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 84e8def83b3..a746a1867cf 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1949,8 +1949,14 @@ def test_df_pivot(scalars_dfs, values, index, columns): ], ) 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) + bf_result = ( + hockey_df.reset_index() + .pivot(values=values, index=index, columns=columns) + .to_pandas() + ) + pd_result = hockey_pandas_df.reset_index().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) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index b5c78de69c2..a87dacae044 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -909,13 +909,36 @@ 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): +@pytest.mark.parametrize( + ("level",), + [(["position", "team_name"],), ([-2, -1],), (["position"],), ("season",), (-3,)], +) +def test_df_multi_index_unstack(hockey_df, hockey_pandas_df, level): bf_result = ( - hockey_df.set_index(["team_name", "season", "position"]).unstack().to_pandas() + hockey_df.set_index(["team_name", "position"], append=True) + .unstack(level=level) + .to_pandas() ) pd_result = hockey_pandas_df.set_index( - ["team_name", "season", "position"] - ).unstack() + ["team_name", "position"], append=True + ).unstack(level=level) + + pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +@pytest.mark.parametrize( + ("level",), + [(["position", "team_name"],), ([-2, -1],), (["position"],), ("season",), (-3,)], +) +def test_series_multi_index_unstack(hockey_df, hockey_pandas_df, level): + bf_result = ( + hockey_df.set_index(["team_name", "position"], append=True)["number"] + .unstack(level=level) + .to_pandas() + ) + pd_result = hockey_pandas_df.set_index(["team_name", "position"], append=True)[ + "number" + ].unstack(level=level) pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index bd1f9a9a185..f0e13e16f5a 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1654,6 +1654,19 @@ def clip(self): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def unstack(self, level): + """ + Unstack, also known as pivot, Series with MultiIndex to produce DataFrame. + + Args: + level (int, str, or list of these, default last level): + Level(s) to unstack, can pass level name. + + Returns: + DataFrame: Unstacked Series. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def argmax(self): """ Return int position of the smallest value in the Series. From 3afd4a35f4c38dad86dab17ff62444cd418cab88 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 26 Oct 2023 18:02:14 -0500 Subject: [PATCH 115/197] feat: add `DataFrame.to_pandas_batches()` to download large `DataFrame` objects (#136) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on https://togithub.com/googleapis/python-bigquery-dataframes/pull/132 Towards internal issue 280662868 🦕 --- bigframes/core/blocks.py | 29 +++++++++++-- bigframes/dataframe.py | 4 ++ bigframes/session/_io/pandas.py | 20 +++++++-- tests/system/small/test_dataframe_io.py | 8 ++++ tests/unit/session/test_io_pandas.py | 56 +++++++++++++++++++++++++ 5 files changed, 109 insertions(+), 8 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index e8a3968b3db..9db193a04e5 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -416,6 +416,30 @@ def to_pandas( ) return df, query_job + def to_pandas_batches(self): + """Download results one message at a time.""" + dtypes = dict(zip(self.index_columns, self.index_dtypes)) + dtypes.update(zip(self.value_columns, self.dtypes)) + results_iterator, _ = self._expr.start_query() + for arrow_table in results_iterator.to_arrow_iterable( + bqstorage_client=self._expr._session.bqstoragereadclient + ): + df = bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) + self._copy_index_to_pandas(df) + yield df + + def _copy_index_to_pandas(self, df: pd.DataFrame): + """Set the index on pandas DataFrame to match this block. + + Warning: This method modifies ``df`` inplace. + """ + if self.index_columns: + df.set_index(list(self.index_columns), inplace=True) + # Pandas names is annotated as list[str] rather than the more + # general Sequence[Label] that BigQuery DataFrames has. + # See: https://github.com/pandas-dev/pandas-stubs/issues/804 + df.index.names = self.index.names # type: ignore + def _compute_and_count( self, value_keys: Optional[Iterable[str]] = None, @@ -489,10 +513,7 @@ def _compute_and_count( else: total_rows = results_iterator.total_rows df = self._to_dataframe(results_iterator) - - if self.index_columns: - df.set_index(list(self.index_columns), inplace=True) - df.index.names = self.index.names # type: ignore + self._copy_index_to_pandas(df) return df, total_rows, query_job diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 869075a9706..3fd83198768 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -893,6 +893,10 @@ def to_pandas( self._set_internal_query_job(query_job) return df.set_axis(self._block.column_labels, axis=1, copy=False) + def to_pandas_batches(self) -> Iterable[pandas.DataFrame]: + """Stream DataFrame results to an iterable of pandas DataFrame""" + return self._block.to_pandas_batches() + def _compute_dry_run(self) -> bigquery.QueryJob: return self._block._compute_dry_run() diff --git a/bigframes/session/_io/pandas.py b/bigframes/session/_io/pandas.py index 163127b546d..1af00a2d015 100644 --- a/bigframes/session/_io/pandas.py +++ b/bigframes/session/_io/pandas.py @@ -46,20 +46,32 @@ def arrow_to_pandas( # Preserve NA/NaN distinction. Note: This is currently needed, even if we use # nullable Float64Dtype in the types_mapper. See: # https://github.com/pandas-dev/pandas/issues/55668 + mask = pyarrow.compute.is_null(column) + nonnull = pyarrow.compute.fill_null(column, float("nan")) # Regarding type: ignore, this class has been public at this # location since pandas 1.2.0. See: # https://pandas.pydata.org/docs/dev/reference/api/pandas.arrays.FloatingArray.html pd_array = pandas.arrays.FloatingArray( # type: ignore - column.to_numpy(), - pyarrow.compute.is_null(column).to_numpy(), + nonnull.to_numpy() + if isinstance(nonnull, pyarrow.ChunkedArray) + else nonnull.to_numpy(zero_copy_only=False), + mask.to_numpy() + if isinstance(mask, pyarrow.ChunkedArray) + else mask.to_numpy(zero_copy_only=False), ) series = pandas.Series(pd_array, dtype=dtype) elif dtype == pandas.Int64Dtype(): # Avoid out-of-bounds errors in Pandas 1.5.x, which incorrectly # casts to float64 in an intermediate step. + mask = pyarrow.compute.is_null(column) + nonnull = pyarrow.compute.fill_null(column, 0) pd_array = pandas.arrays.IntegerArray( - pyarrow.compute.fill_null(column, 0).to_numpy(), - pyarrow.compute.is_null(column).to_numpy(), + nonnull.to_numpy() + if isinstance(nonnull, pyarrow.ChunkedArray) + else nonnull.to_numpy(zero_copy_only=False), + mask.to_numpy() + if isinstance(mask, pyarrow.ChunkedArray) + else mask.to_numpy(zero_copy_only=False), ) series = pandas.Series(pd_array, dtype=dtype) elif isinstance(dtype, pandas.ArrowDtype): diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index d60083a837e..8f5d706f621 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -83,6 +83,14 @@ def test_to_pandas_array_struct_correct_result(session): ) +def test_to_pandas_batches_w_correct_dtypes(scalars_df_default_index): + """Verify to_pandas_batches() APIs returns the expected dtypes.""" + expected = scalars_df_default_index.dtypes + for df in scalars_df_default_index.to_pandas_batches(): + actual = df.dtypes + pd.testing.assert_series_equal(actual, expected) + + @pytest.mark.parametrize( ("index"), [True, False], diff --git a/tests/unit/session/test_io_pandas.py b/tests/unit/session/test_io_pandas.py index 8b95977ec3d..0f6f5dae03b 100644 --- a/tests/unit/session/test_io_pandas.py +++ b/tests/unit/session/test_io_pandas.py @@ -231,6 +231,62 @@ ), id="scalar-dtypes", ), + pytest.param( + pyarrow.Table.from_pydict( + { + "bool": pyarrow.chunked_array( + [[True, None], [True, False]], + type=pyarrow.bool_(), + ), + "bytes": pyarrow.chunked_array( + [[b"123", None], [b"abc", b"xyz"]], + type=pyarrow.binary(), + ), + "float": pyarrow.chunked_array( + [[1.0, None], [float("nan"), -1.0]], + type=pyarrow.float64(), + ), + "int": pyarrow.chunked_array( + [[1, None], [-1, 2**63 - 1]], + type=pyarrow.int64(), + ), + "string": pyarrow.chunked_array( + [["123", None], ["abc", "xyz"]], + type=pyarrow.string(), + ), + } + ), + { + "bool": "boolean", + "bytes": "object", + "float": pandas.Float64Dtype(), + "int": pandas.Int64Dtype(), + "string": "string[pyarrow]", + }, + pandas.DataFrame( + { + "bool": pandas.Series([True, None, True, False], dtype="boolean"), + "bytes": [b"123", None, b"abc", b"xyz"], + "float": pandas.Series( + pandas.arrays.FloatingArray( # type: ignore + numpy.array( + [1.0, float("nan"), float("nan"), -1.0], dtype="float64" + ), + numpy.array([False, True, False, False], dtype="bool"), + ), + dtype=pandas.Float64Dtype(), + ), + "int": pandas.Series( + [1, None, -1, 2**63 - 1], + dtype=pandas.Int64Dtype(), + ), + "string": pandas.Series( + ["123", None, "abc", "xyz"], dtype="string[pyarrow]" + ), + } + ), + id="scalar-dtypes-chunked_array", + ), pytest.param( pyarrow.Table.from_pydict( { From 39df43e243ac0374d1a1eb2a75779324825afbe9 Mon Sep 17 00:00:00 2001 From: Bradford Orr <15842009+orrbradford@users.noreply.github.com> Date: Thu, 26 Oct 2023 16:46:14 -0700 Subject: [PATCH 116/197] =?UTF-8?q?fix:=20resolve=20plotly=20rendering=20i?= =?UTF-8?q?ssue=20by=20using=20ipython=20html=20for=20job=20pro=E2=80=A6?= =?UTF-8?q?=20(#134)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …gress messages Fixes bug that was preventing plotly rendering to show after the progress bar. Original ipywidgets implementation isn't necessary for basic opening of urls Screen recording: https://togithub.com/googleapis/python-bigquery-dataframes/assets/15842009/5225ce05-117a-4808-9ff0-cb2c3aaf3a40 Internal bug: b/297062404 --- bigframes/formatting_helpers.py | 24 ++++--- tests/system/small/test_progress_bar.py | 83 +++++++++++-------------- 2 files changed, 53 insertions(+), 54 deletions(-) diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py index 6851bdd2bd4..752aeb7a10b 100644 --- a/bigframes/formatting_helpers.py +++ b/bigframes/formatting_helpers.py @@ -16,6 +16,7 @@ # TODO(orrbradford): cleanup up typings and documenttion in this file import datetime +import random from typing import Any, Optional, Union import google.api_core.exceptions as api_core_exceptions @@ -57,9 +58,9 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]): Pywidget html table. """ if query_job is None: - return widgets.HTML("No job information available") + return display.HTML("No job information available") if query_job.dry_run: - return widgets.HTML( + return display.HTML( f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}" ) table_html = "" @@ -125,16 +126,20 @@ def wait_for_query_job( 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" try: if progress_bar == "notebook": - display.display(loading_bar) + display_id = str(random.random()) + loading_bar = display.HTML(get_query_job_loading_html(query_job)) + display.display(loading_bar, display_id=display_id) query_result = query_job.result(max_results=max_results) query_job.reload() - loading_bar.value = get_query_job_loading_html(query_job) + display.update_display( + display.HTML(get_query_job_loading_html(query_job)), + display_id=display_id, + ) elif progress_bar == "terminal": initial_loading_bar = get_query_job_loading_string(query_job) print(initial_loading_bar) @@ -171,16 +176,19 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None): progress_bar (str, Optional): Which progress bar to show. """ - 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) + display_id = str(random.random()) + loading_bar = display.HTML(get_base_job_loading_html(job)) + display.display(loading_bar, display_id=display_id) job.result() job.reload() - loading_bar.value = get_base_job_loading_html(job) + display.update_display( + display.HTML(get_base_job_loading_html(job)), display_id=display_id + ) elif progress_bar == "terminal": inital_loading_bar = get_base_job_loading_string(job) print(inital_loading_bar) diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index 00380c26390..f7fc4eaa8f6 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re import tempfile import pandas as pd @@ -19,94 +20,84 @@ import bigframes as bf import bigframes.formatting_helpers as formatting_helpers +job_load_message_regex = r"\w+ job [\w-]+ is \w+\." + def test_progress_bar_dataframe( penguins_df_default_index: bf.dataframe.DataFrame, capsys ): - bf.options.display.progress_bar = "notebook" + bf.options.display.progress_bar = "terminal" + capsys.readouterr() # clear output penguins_df_default_index.to_pandas() - 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 + + assert_loading_msg_exist(capsys.readouterr().out) 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" + bf.options.display.progress_bar = "terminal" series = penguins_df_default_index["body_mass_g"].head(10) + capsys.readouterr() # clear output series.to_pandas() - 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 + + assert_loading_msg_exist(capsys.readouterr().out) 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" + bf.options.display.progress_bar = "terminal" + capsys.readouterr() # clear output 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 = [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 + + assert_loading_msg_exist(capsys.readouterr().out) def test_progress_bar_read_gbq(session: bf.Session, penguins_table_id: str, capsys): - bf.options.display.progress_bar = "notebook" + bf.options.display.progress_bar = "terminal" + capsys.readouterr() # clear output 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 + + assert_loading_msg_exist(capsys.readouterr().out) def test_progress_bar_extract_jobs( penguins_df_default_index: bf.dataframe.DataFrame, gcs_folder, capsys ): - bf.options.display.progress_bar = "notebook" + bf.options.display.progress_bar = "terminal" path = gcs_folder + "test_read_csv_progress_bar*.csv" + capsys.readouterr() # clear output 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 + + assert_loading_msg_exist(capsys.readouterr().out) def test_progress_bar_load_jobs( session: bf.Session, penguins_pandas_df_default_index: pd.DataFrame, capsys ): - bf.options.display.progress_bar = "notebook" + bf.options.display.progress_bar = "terminal" with tempfile.TemporaryDirectory() as dir: path = dir + "/test_read_csv_progress_bar*.csv" penguins_pandas_df_default_index.to_csv(path, index=False) + capsys.readouterr() # clear output session.read_csv(path) - html_check = "HTML(value=" - open_job_check = "Open Job" - lines = capsys.readouterr().out.split("\n") + + assert_loading_msg_exist(capsys.readouterr().out) + + +def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex): + numLoadingMsg = 0 + lines = capystOut.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 + if re.match(pattern, line) is not None: + numLoadingMsg += 1 + assert numLoadingMsg > 0 def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): - bf.options.display.progress_bar = "notebook" + bf.options.display.progress_bar = "terminal" penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = ( False ) From eceeb221f553644411b954ae2db0f0ae5a505687 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 26 Oct 2023 17:28:13 -0700 Subject: [PATCH 117/197] refactor: ArrayValue is now a tree that defers conversion to ibis (#110) 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/core/__init__.py | 1221 +++-------------- bigframes/core/block_transforms.py | 25 +- bigframes/core/blocks.py | 34 +- bigframes/core/compile/__init__.py | 21 + bigframes/core/compile/compiled.py | 1121 +++++++++++++++ bigframes/core/compile/compiler.py | 185 +++ .../core/{joins => compile}/row_identity.py | 14 +- .../core/{joins => compile}/single_column.py | 35 +- bigframes/core/groupby/__init__.py | 20 +- bigframes/core/indexers.py | 6 +- bigframes/core/indexes/index.py | 15 +- bigframes/core/joins/__init__.py | 9 +- bigframes/core/nodes.py | 245 ++++ bigframes/core/ordering.py | 4 +- bigframes/core/window_spec.py | 35 + bigframes/dataframe.py | 18 +- bigframes/ml/metrics.py | 2 +- bigframes/operations/base.py | 4 +- bigframes/series.py | 34 +- bigframes/session/__init__.py | 44 +- tests/system/small/test_progress_bar.py | 4 +- tests/system/small/test_series.py | 4 +- tests/system/small/test_session.py | 9 +- tests/unit/core/test_blocks.py | 5 +- tests/unit/resources.py | 17 +- tests/unit/test_core.py | 37 +- 26 files changed, 1996 insertions(+), 1172 deletions(-) create mode 100644 bigframes/core/compile/__init__.py create mode 100644 bigframes/core/compile/compiled.py create mode 100644 bigframes/core/compile/compiler.py rename bigframes/core/{joins => compile}/row_identity.py (94%) rename bigframes/core/{joins => compile}/single_column.py (87%) create mode 100644 bigframes/core/nodes.py create mode 100644 bigframes/core/window_spec.py diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 6c78a07f3ba..4653f0ab6aa 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -14,29 +14,21 @@ from __future__ import annotations from dataclasses import dataclass -import functools -import math -import textwrap +import io import typing -from typing import Collection, Iterable, Literal, Optional, Sequence, Tuple +from typing import 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.constants as constants +import bigframes.core.compile as compiled import bigframes.core.guid -from bigframes.core.ordering import ( - encode_order_string, - ExpressionOrdering, - IntegerEncoding, - OrderingColumnReference, - reencode_order_string, - StringEncoding, -) -import bigframes.core.utils as utils +import bigframes.core.nodes as nodes +from bigframes.core.ordering import OrderingColumnReference +import bigframes.core.ordering as orderings +from bigframes.core.window_spec import WindowSpec import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -49,470 +41,190 @@ @dataclass(frozen=True) -class WindowSpec: +class ArrayValue: """ - 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 + ArrayValue is an immutable type representing a 2D array with per-column types. """ - 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. + node: nodes.BigFrameNode - 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, + @classmethod + def from_ibis( + cls, session: Session, table: ibis_types.Table, columns: Sequence[ibis_types.Value], - hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None, - ordering: ExpressionOrdering = ExpressionOrdering(), - predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + hidden_ordering_columns: Sequence[ibis_types.Value], + ordering: orderings.ExpressionOrdering, ): - self._session = session - self._table = table - self._predicates = tuple(predicates) if predicates is not None else () - # TODO: Validate ordering - 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. - # 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 - } - ### 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 + node = nodes.ReadGbqNode( + table=table, + table_session=session, + columns=tuple(columns), + hidden_ordering_columns=tuple(hidden_ordering_columns), + ordering=ordering, ) - 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}") + return cls(node) @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. + def from_pandas(cls, pd_df: pandas.DataFrame): + iobytes = io.BytesIO() + # Discard row labels and use simple string ids for columns + column_ids = tuple(str(label) for label in pd_df.columns) + pd_df.reset_index(drop=True).set_axis(column_ids, axis=1).to_feather(iobytes) + node = nodes.ReadLocalNode(iobytes.getvalue(), column_ids=column_ids) + return cls(node) - Caution: If session is None, only a subset of expr functionality will - be available (null Session is usually not supported). - """ - # 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) + @property + def column_ids(self) -> typing.Sequence[str]: + return self.compile().column_ids - # 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)) + @property + def session(self) -> Session: + required_session = self.node.session + from bigframes import get_global_session - # 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 self.node.session[0] if required_session else get_global_session() - 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]), - ), - hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), - ) - - @property - def columns(self) -> typing.Tuple[ibis_types.Value, ...]: - return self._columns + def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: + return self.compile().get_column_type(key) - @property - def column_ids(self) -> typing.Sequence[str]: - return tuple(self._column_names.keys()) + def compile(self) -> compiled.CompiledArrayValue: + return compiled.compile_node(self.node) - @property - def _hidden_column_ids(self) -> typing.Sequence[str]: - return tuple(self._hidden_ordering_column_names.keys()) + def shape(self) -> typing.Tuple[int, int]: + """Returns dimensions as (length, width) tuple.""" + width = len(self.compile().columns) + count_expr = self.compile()._to_ibis_expr("unordered").count() - @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 + # Support in-memory engines for hermetic unit tests. + if not self.node.session: + try: + length = ibis.pandas.connect({}).execute(count_expr) + return (length, width) + except Exception: + # Not all cases can be handled by pandas engine + pass + + 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) - @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.""" - return _convert_ordering_to_table_values( - {**self._column_names, **self._hidden_ordering_column_names}, - self._ordering.all_ordering_columns, + def to_sql( + self, + offset_column: typing.Optional[str] = None, + col_id_overrides: typing.Mapping[str, str] = {}, + sorted: bool = False, + ) -> str: + return self.compile().to_sql( + offset_column=offset_column, + col_id_overrides=col_id_overrides, + sorted=sorted, ) - 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, - columns=self._columns, - hidden_ordering_columns=self._hidden_ordering_columns, - ordering=self._ordering, - predicates=self._predicates, + def start_query( + self, + job_config: Optional[bigquery.job.QueryJobConfig] = None, + max_results: Optional[int] = None, + *, + 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 + # 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? + sql = self.to_sql(sorted=sorted) # type:ignore + return self.session._start_query( + sql=sql, + job_config=job_config, + max_results=max_results, ) - 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_any_column(key).type() + 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.""" + compiled = self.compile() + ibis_expr = compiled._to_ibis_expr("unordered", expose_hidden_cols=True) + destination = self.session._ibis_to_session_table( + ibis_expr, cluster_cols=cluster_cols, api_name="cache" ) - return typing.cast( - bigframes.dtypes.Dtype, - bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), + table_expression = self.session.ibis_client.table( + f"{destination.project}.{destination.dataset_id}.{destination.table_id}" + ) + new_columns = [table_expression[column] for column in compiled.column_ids] + new_hidden_columns = [ + table_expression[column] + for column in compiled._hidden_ordering_column_names + ] + return ArrayValue.from_ibis( + self.session, + table_expression, + columns=new_columns, + hidden_ordering_columns=new_hidden_columns, + ordering=compiled._ordering, ) - def _get_ibis_column(self, key: str) -> ibis_types.Value: - """Gets the Ibis expression for a given column.""" - if key not in self.column_ids: - raise ValueError( - "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: - """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]) + # Operations - 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 drop_columns(self, columns: Iterable[str]) -> ArrayValue: + return ArrayValue( + nodes.DropColumnsNode(child=self.node, columns=tuple(columns)) + ) 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_ibis_column(predicate_id) - ) - if keep_null: - condition = typing.cast( - ibis_types.BooleanValue, - condition.fillna( - typing.cast(ibis_types.BooleanScalar, ibis_types.literal(True)) - ), + return ArrayValue( + nodes.FilterNode( + child=self.node, predicate_id=predicate_id, keep_null=keep_null ) - 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_value] - 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() - - 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( - "unordered", 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, + nodes.OrderByNode(child=self.node, by=tuple(by), stable=stable) ) - @property - 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." - ) - if not self._ordering.total_order_col: - raise ValueError( - "Ordering is invalid. Marked as sequential but no total order columns." - ) - column = self._get_any_column(self._ordering.total_order_col.column_id) - return typing.cast(ibis_types.IntegerColumn, column) - - 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_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)], - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - integer_encoding=IntegerEncoding(True, 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_ibis_column(column_id).name(new_name), - ] - expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) - return expr_builder.build() + def reversed(self) -> ArrayValue: + return ArrayValue(nodes.ReversedNode(child=self.node)) 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. """ - # Special case: offsets already exist - ordering = self._ordering - - if (not ordering.is_sequential) or (not ordering.total_order_col): - 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() + return ArrayValue(nodes.PromoteOffsetsNode(child=self.node, col_id=col_id)) 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] + return ArrayValue( + nodes.SelectNode(child=self.node, column_ids=tuple(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_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 - 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("unordered").count() - sql = self._session.ibis_client.compile(count_expr) - - # 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: """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.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="string_encoded", 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 - + reencode_order_string( - 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_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, - 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, + nodes.ConcatNode(children=tuple([self.node, *[val.node for val in other]])) ) 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_ibis_column(column_name)).name( - output_name or column_name + return ArrayValue( + nodes.ProjectUnaryOpNode( + child=self.node, input_id=column_name, op=op, output_id=output_name + ) ) - return self._set_or_replace_by_id(output_name or column_name, value) def project_binary_op( self, @@ -522,11 +234,15 @@ def project_binary_op( output_column_id: str, ) -> ArrayValue: """Creates a new expression based on this expression with binary operation applied to two columns.""" - value = op( - 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) + return ArrayValue( + nodes.ProjectBinaryOpNode( + child=self.node, + left_input_id=left_column_id, + right_input_id=right_column_id, + op=op, + output_id=output_column_id, + ) + ) def project_ternary_op( self, @@ -537,12 +253,16 @@ def project_ternary_op( output_column_id: str, ) -> ArrayValue: """Creates a new expression based on this expression with ternary operation applied to three columns.""" - value = op( - 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) + return ArrayValue( + nodes.ProjectTernaryOpNode( + child=self.node, + input_id1=col_id_1, + input_id2=col_id_2, + input_id3=col_id_3, + op=op, + output_id=output_column_id, + ) + ) def aggregate( self, @@ -557,46 +277,14 @@ 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("unordered") - 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 - ], - total_ordering_columns=frozenset(by_column_ids), - ) - 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( - 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() - 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_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, + return ArrayValue( + nodes.AggregateNode( + child=self.node, + aggregations=tuple(aggregations), + by_column_ids=tuple(by_column_ids), + dropna=dropna, ) + ) def corr_aggregate( self, corr_aggregations: typing.Sequence[typing.Tuple[str, str, str]] @@ -607,25 +295,8 @@ def corr_aggregate( Arguments: corr_aggregations: left_column_id, right_column_id, output_column_id tuples """ - 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 - } - 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, + nodes.CorrNode(child=self.node, corr_aggregations=tuple(corr_aggregations)) ) def project_window_op( @@ -647,231 +318,17 @@ 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_ibis_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 and not never_skip_nulls: - clauses.append((column.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( - ( - observation_count < 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_sql( - self, - 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="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["string_encoded", "offset_col", "unordered"], - 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. - - ArrayValue objects are sorted, so the following options are available - to reflect this in the ibis expression. - - * "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. - * "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``. - 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 ( - "string_encoded", - "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}" + return ArrayValue( + nodes.WindowOpNode( + child=self.node, + column_name=column_name, + op=op, + window_spec=window_spec, + output_name=output_name, + never_skip_nulls=never_skip_nulls, + skip_reproject_unsafe=skip_reproject_unsafe, ) - - columns = list(self._columns) - 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) - # 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()) - - order_columns = self._create_order_columns( - ordering_mode, order_col_name, expose_hidden_cols ) - columns.extend(order_columns) - - # Special case for empty tables, since we can't create an empty - # projection. - if not columns: - return ibis.memtable([]) - - # Make sure all dtypes are the "canonical" ones for BigFrames. This is - # important for operations like UNION where the schema must match. - 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(base_table[PREDICATE_COLUMN]) - 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 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, - *, - 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 - # 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? - sql = self.to_sql(sorted=True) # type:ignore - return self._session._start_query( - sql=sql, - job_config=job_config, - 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: """ @@ -881,74 +338,25 @@ def _reproject_to_table(self) -> ArrayValue: some operations such as window operations that cannot be used recursively in projections. """ - table = self._to_ibis_expr( - "unordered", - 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 - if column_name in ordering_col_ids - ] 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_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 window_spec.ordering: - order_by = _convert_ordering_to_table_values( - {**self._column_names, **self._hidden_ordering_column_names}, - window_spec.ordering, + nodes.ReprojectOpNode( + child=self.node, ) - if not allow_ties: - # Most operator need an unambiguous ordering, so the table's total ordering is appended - 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._ibis_order) - 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 unpivot( self, row_labels: typing.Sequence[typing.Hashable], unpivot_columns: typing.Sequence[ - typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + typing.Tuple[str, typing.Tuple[typing.Optional[str], ...]] ], *, passthrough_columns: typing.Sequence[str] = (), index_col_ids: typing.Sequence[str] = ["index"], dtype: typing.Union[ - bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + bigframes.dtypes.Dtype, typing.Tuple[bigframes.dtypes.Dtype, ...] ] = pandas.Float64Dtype(), - how="left", + how: typing.Literal["left", "right"] = "left", ) -> ArrayValue: """ Unpivot ArrayValue columns. @@ -963,133 +371,23 @@ def unpivot( Returns: ArrayValue: The unpivoted ArrayValue """ - if how not in ("left", "right"): - raise ValueError("'how' must be 'left' or 'right'") - 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( - len(source_columns) == row_n for _, source_columns in unpivot_columns - ): - raise ValueError("Columns and row labels must all be same length.") - - unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") - unpivot_table = table.cross_join( - ibis.memtable({unpivot_offset_id: range(row_n)}) - ) - # Use ibis memtable to infer type of rowlabels (if possible) - # TODO: Allow caller to specify dtype - 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 - ] - - 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)): - 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 - ) - unpivot_values.append(unpivot_value.name(result_col)) - - unpivot_table = unpivot_table.select( - 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[col_id] for col_id in index_col_ids], - *value_columns, - *passthrough_values, - ], - hidden_ordering_columns=hidden_ordering_columns, - ordering=new_ordering, + nodes.UnpivotNode( + child=self.node, + row_labels=tuple(row_labels), + unpivot_columns=tuple(unpivot_columns), + passthrough_columns=tuple(passthrough_columns), + index_col_ids=tuple(index_col_ids), + dtype=dtype, + how=how, + ) ) def assign(self, source_id: str, destination_id: str) -> ArrayValue: - return self._set_or_replace_by_id( - destination_id, self._get_ibis_column(source_id) + return ArrayValue( + nodes.AssignNode( + child=self.node, source_id=source_id, destination_id=destination_id + ) ) def assign_constant( @@ -1098,128 +396,41 @@ def assign_constant( 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)}. {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) -> 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_ids: - builder.columns = [ - val if (col_id != id) else new_value.name(id) - for col_id, val in zip(self.column_ids, self._columns) - ] - else: - builder.columns = [*self.columns, new_value.name(id)] - return builder.build() - - 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("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_ids] - 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, + nodes.AssignConstantNode( + child=self.node, destination_id=destination_id, value=value, dtype=dtype + ) ) - -class ArrayValueBuilder: - """Mutable expression class. - Use ArrayValue.builder() to create from a ArrayValue object. - """ - - def __init__( + def join( self, - session: Session, - table: ibis_types.Table, - ordering: ExpressionOrdering, - columns: Collection[ibis_types.Value] = (), - hidden_ordering_columns: Collection[ibis_types.Value] = (), - predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + self_column_ids: typing.Sequence[str], + other: ArrayValue, + other_column_ids: typing.Sequence[str], + *, + how: Literal[ + "inner", + "left", + "outer", + "right", + ], + allow_row_identity_join: bool = True, ): - 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) + nodes.JoinNode( + left_child=self.node, + right_child=other.node, + left_column_ids=tuple(self_column_ids), + right_column_ids=tuple(other_column_ids), + how=how, + allow_row_identity_join=allow_row_identity_join, + ) ) - # 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 _uniform_sampling(self, fraction: float) -> ArrayValue: + """Sampling the table on given fraction. -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 + .. warning:: + The row numbers of result is non-deterministic, avoid to use. + """ + return ArrayValue(nodes.RandomSampleNode(self.node, fraction)) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index b0f05f47984..3706bf1681b 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -21,6 +21,7 @@ import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.ordering as ordering +import bigframes.core.window_spec as windows import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -68,21 +69,21 @@ def indicate_duplicates( 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( + window_spec = windows.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( + window_spec = windows.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)) + window_spec = windows.WindowSpec(grouping_keys=tuple(columns)) block, dummy = block.create_constant(1) block, val_count_col_id = block.apply_window_op( dummy, @@ -131,7 +132,7 @@ def value_counts( ) count_id = agg_ids[0] if normalize: - unbound_window = core.WindowSpec() + unbound_window = windows.WindowSpec() block, total_count_id = block.apply_window_op( count_id, agg_ops.sum_op, unbound_window ) @@ -153,7 +154,7 @@ def value_counts( def pct_change(block: blocks.Block, periods: int = 1) -> blocks.Block: column_labels = block.column_labels - window_spec = core.WindowSpec( + window_spec = windows.WindowSpec( preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -195,7 +196,7 @@ def rank( ops.isnull_op, ) nullity_col_ids.append(nullity_col_id) - window = core.WindowSpec( + window = windows.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( @@ -229,7 +230,7 @@ def rank( block, result_id = block.apply_window_op( rownum_col_ids[i], agg_op, - window_spec=core.WindowSpec(grouping_keys=[columns[i]]), + window_spec=windows.WindowSpec(grouping_keys=(columns[i],)), skip_reproject_unsafe=(i < (len(columns) - 1)), ) post_agg_rownum_col_ids.append(result_id) @@ -311,7 +312,7 @@ def nsmallest( block, counter = block.apply_window_op( column_ids[0], agg_ops.rank_op, - window_spec=core.WindowSpec(ordering=order_refs), + window_spec=windows.WindowSpec(ordering=tuple(order_refs)), ) block, condition = block.apply_unary_op( counter, ops.partial_right(ops.le_op, n) @@ -343,7 +344,7 @@ def nlargest( block, counter = block.apply_window_op( column_ids[0], agg_ops.rank_op, - window_spec=core.WindowSpec(ordering=order_refs), + window_spec=windows.WindowSpec(ordering=tuple(order_refs)), ) block, condition = block.apply_unary_op( counter, ops.partial_right(ops.le_op, n) @@ -440,14 +441,14 @@ def _mean_delta_to_power( 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) + window = windows.WindowSpec(grouping_keys=tuple(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) + block = block.drop_columns([delta_id]) delta_ids.append(delta_power_id) return block, delta_ids @@ -645,7 +646,7 @@ def _idx_extrema( for idx_col in original_block.index_columns ], ] - window_spec = core.WindowSpec(ordering=order_refs) + window_spec = windows.WindowSpec(ordering=tuple(order_refs)) idx_col = original_block.index_columns[0] block, result_col = block.apply_window_op( idx_col, agg_ops.first_op, window_spec diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 9db193a04e5..cc13edeaf96 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -35,7 +35,6 @@ 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.joins.name_resolution as join_names import bigframes.core.ordering as ordering import bigframes.core.utils @@ -378,7 +377,7 @@ def _to_dataframe(self, result) -> pd.DataFrame: """Convert BigQuery data to pandas DataFrame with specific dtypes.""" dtypes = dict(zip(self.index_columns, self.index_dtypes)) dtypes.update(zip(self.value_columns, self.dtypes)) - return self._expr._session._rows_to_dataframe(result, dtypes) + return self._expr.session._rows_to_dataframe(result, dtypes) def to_pandas( self, @@ -422,7 +421,7 @@ def to_pandas_batches(self): dtypes.update(zip(self.value_columns, self.dtypes)) results_iterator, _ = self._expr.start_query() for arrow_table in results_iterator.to_arrow_iterable( - bqstorage_client=self._expr._session.bqstoragereadclient + bqstorage_client=self._expr.session.bqstoragereadclient ): df = bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) self._copy_index_to_pandas(df) @@ -454,7 +453,9 @@ def _compute_and_count( results_iterator, query_job = expr.start_query(max_results=max_results) - table_size = expr._get_table_size(query_job.destination) / _BYTES_TO_MEGABYTES + table_size = ( + expr.session._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) @@ -819,7 +820,9 @@ def aggregate_all_and_stack( axis: int | str = 0, value_col_id: str = "values", dropna: bool = True, - dtype=pd.Float64Dtype(), + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Tuple[bigframes.dtypes.Dtype, ...] + ] = pd.Float64Dtype(), ) -> Block: axis_n = utils.get_axis_number(axis) if axis_n == 0: @@ -829,7 +832,7 @@ def aggregate_all_and_stack( result_expr = self.expr.aggregate(aggregations, dropna=dropna).unpivot( row_labels=self.column_labels.to_list(), index_col_ids=["index"], - unpivot_columns=[(value_col_id, self.value_columns)], + unpivot_columns=tuple([(value_col_id, tuple(self.value_columns))]), dtype=dtype, ) return Block(result_expr, index_columns=["index"], column_labels=[None]) @@ -841,7 +844,7 @@ def aggregate_all_and_stack( stacked_expr = expr_with_offsets.unpivot( row_labels=self.column_labels.to_list(), index_col_ids=[guid.generate_guid()], - unpivot_columns=[(value_col_id, self.value_columns)], + unpivot_columns=[(value_col_id, tuple(self.value_columns))], passthrough_columns=[*self.index_columns, offset_col], dtype=dtype, ) @@ -1029,13 +1032,13 @@ def summarize( for col_id in column_ids ] columns = [ - (col_id, [f"{col_id}-{stat.name}" for stat in stats]) + (col_id, tuple(f"{col_id}-{stat.name}" for stat in stats)) for col_id in column_ids ] expr = self.expr.aggregate(aggregations).unpivot( labels, - unpivot_columns=columns, - index_col_ids=[label_col_id], + unpivot_columns=tuple(columns), + index_col_ids=tuple([label_col_id]), ) labels = self._get_labels_for_columns(column_ids) return Block(expr, column_labels=labels, index_columns=[label_col_id]) @@ -1342,7 +1345,7 @@ def stack(self, how="left", levels: int = 1): passthrough_columns=self.index_columns, unpivot_columns=unpivot_columns, index_col_ids=added_index_columns, - dtype=dtypes, + dtype=tuple(dtypes), how=how, ) new_index_level_names = self.column_labels.names[-levels:] @@ -1382,7 +1385,7 @@ def _create_stack_column( 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() + return tuple(input_columns), dtype or pd.Float64Dtype() def _column_type(self, col_id: str) -> bigframes.dtypes.Dtype: col_offset = self.value_columns.index(col_id) @@ -1497,8 +1500,7 @@ def merge( sort: bool, suffixes: tuple[str, str] = ("_x", "_y"), ) -> Block: - joined_expr = joins.join_by_column( - self.expr, + joined_expr = self.expr.join( left_join_ids, other.expr, right_join_ids, @@ -1708,7 +1710,7 @@ def _is_monotonic( return result -def block_from_local(data, session=None) -> Block: +def block_from_local(data) -> Block: pd_data = pd.DataFrame(data) columns = pd_data.columns @@ -1730,7 +1732,7 @@ def block_from_local(data, session=None) -> Block: ) index_ids = pd_data.columns[: len(index_labels)] - keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session) + keys_expr = core.ArrayValue.from_pandas(pd_data) return Block( keys_expr, column_labels=columns, diff --git a/bigframes/core/compile/__init__.py b/bigframes/core/compile/__init__.py new file mode 100644 index 00000000000..c86f4463dc0 --- /dev/null +++ b/bigframes/core/compile/__init__.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. + +from bigframes.core.compile.compiled import CompiledArrayValue +from bigframes.core.compile.compiler import compile_node + +__all__ = [ + "compile_node", + "CompiledArrayValue", +] diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py new file mode 100644 index 00000000000..1134f1aab01 --- /dev/null +++ b/bigframes/core/compile/compiled.py @@ -0,0 +1,1121 @@ +# 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 math +import textwrap +import typing +from typing import Collection, Iterable, Literal, Optional, Sequence + +import ibis +import ibis.backends.bigquery as ibis_bigquery +import ibis.expr.datatypes as ibis_dtypes +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, + reencode_order_string, + StringEncoding, +) +import bigframes.core.utils as utils +from bigframes.core.window_spec import WindowSpec +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops + +ORDER_ID_COLUMN = "bigframes_ordering_id" +PREDICATE_COLUMN = "bigframes_predicate" + + +class CompiledArrayValue: + """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: + 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, + table: ibis_types.Table, + columns: Sequence[ibis_types.Value], + hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None, + ordering: ExpressionOrdering = ExpressionOrdering(), + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + self._table = table + self._predicates = tuple(predicates) if predicates is not None else () + # TODO: Validate ordering + 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. + # 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 + } + ### 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( + cls, + pd_df: pandas.DataFrame, + ) -> CompiledArrayValue: + """ + Builds an in-memory only (SQL only) expr from a pandas dataframe. + """ + # 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( + keys_memtable, + columns=[ + keys_memtable[f"col_{column_index}"].name(column) + for column_index, column in enumerate(column_names) + ], + ordering=ExpressionOrdering( + ordering_value_columns=tuple( + [OrderingColumnReference(ORDER_ID_COLUMN)] + ), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + ), + hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), + ) + + @property + def columns(self) -> typing.Tuple[ibis_types.Value, ...]: + return self._columns + + @property + def column_ids(self) -> typing.Sequence[str]: + return tuple(self._column_names.keys()) + + @property + 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]: + """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.""" + 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._table, + columns=self._columns, + hidden_ordering_columns=self._hidden_ordering_columns, + ordering=self._ordering, + predicates=self._predicates, + ) + + def drop_columns(self, columns: Iterable[str]) -> CompiledArrayValue: + # 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_any_column(key).type() + ) + return typing.cast( + bigframes.dtypes.Dtype, + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), + ) + + def _get_ibis_column(self, key: str) -> ibis_types.Value: + """Gets the Ibis expression for a given column.""" + if key not in self.column_ids: + raise ValueError( + "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: + """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 filter(self, predicate_id: str, keep_null: bool = False) -> CompiledArrayValue: + """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_ibis_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) -> CompiledArrayValue: + """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_value] + return expr.build() + + def order_by( + self, by: Sequence[OrderingColumnReference], stable: bool = False + ) -> CompiledArrayValue: + expr_builder = self.builder() + expr_builder.ordering = self._ordering.with_ordering_columns(by, stable=stable) + return expr_builder.build() + + def reversed(self) -> CompiledArrayValue: + expr_builder = self.builder() + expr_builder.ordering = self._ordering.with_reverse() + return expr_builder.build() + + def _uniform_sampling(self, fraction: float) -> CompiledArrayValue: + """Sampling the table on given fraction. + + .. warning:: + The row numbers of result is non-deterministic, avoid to use. + """ + table = self._to_ibis_expr( + "unordered", 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 CompiledArrayValue( + table, + columns=columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=self._ordering, + ) + + @property + 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." + ) + if not self._ordering.total_order_col: + raise ValueError( + "Ordering is invalid. Marked as sequential but no total order columns." + ) + column = self._get_any_column(self._ordering.total_order_col.column_id) + return typing.cast(ibis_types.IntegerColumn, column) + + def _project_offsets(self) -> CompiledArrayValue: + """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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(True, is_sequential=True), + ) + return CompiledArrayValue( + table, + columns=columns, + hidden_ordering_columns=[table[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def _hide_column(self, column_id) -> CompiledArrayValue: + """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_ibis_column(column_id).name(new_name), + ] + expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) + return expr_builder.build() + + def promote_offsets(self, col_id: str) -> CompiledArrayValue: + """ + 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.total_order_col): + 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() + + def select_columns(self, column_ids: typing.Sequence[str]) -> CompiledArrayValue: + """Creates a new expression based on this expression with new columns.""" + columns = [self._get_ibis_column(col_id) for col_id in column_ids] + expr = self + 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 + expr = expr._hide_column(ordering_column) + builder = expr.builder() + builder.columns = list(columns) + new_expr = builder.build() + return new_expr + + def concat(self, other: typing.Sequence[CompiledArrayValue]) -> CompiledArrayValue: + """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.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="string_encoded", 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 + + reencode_order_string( + 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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + string_encoding=StringEncoding(True, prefix_size + max_encoding_size), + ) + return CompiledArrayValue( + 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 + ) -> CompiledArrayValue: + """Creates a new expression based on this expression with unary operation applied to one column.""" + 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) + + def project_binary_op( + self, + left_column_id: str, + right_column_id: str, + op: ops.BinaryOp, + output_column_id: str, + ) -> CompiledArrayValue: + """Creates a new expression based on this expression with binary operation applied to two columns.""" + value = op( + 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) + + def project_ternary_op( + self, + col_id_1: str, + col_id_2: str, + col_id_3: str, + op: ops.TernaryOp, + output_column_id: str, + ) -> CompiledArrayValue: + """Creates a new expression based on this expression with ternary operation applied to three columns.""" + value = op( + 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) + + def aggregate( + self, + aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp, str]], + by_column_ids: typing.Sequence[str] = (), + dropna: bool = True, + ) -> CompiledArrayValue: + """ + Apply aggregations to the expression. + Arguments: + 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("unordered") + 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( + tuple( + [ + OrderingColumnReference(column_id=column_id) + for column_id in by_column_ids + ] + ), + total_ordering_columns=frozenset(by_column_ids), + ) + columns = tuple(result[key] for key in result.columns) + expr = CompiledArrayValue(result, columns=columns, ordering=ordering) + if dropna: + for column_id in by_column_ids: + expr = expr._filter( + 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() + 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_value_columns=tuple( + [OrderingColumnReference(ORDER_ID_COLUMN)] + ), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ) + return CompiledArrayValue( + result, + columns=[result[col_id] for col_id in [*stats.keys()]], + hidden_ordering_columns=[result[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def corr_aggregate( + self, corr_aggregations: typing.Sequence[typing.Tuple[str, str, str]] + ) -> CompiledArrayValue: + """ + 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("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=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ) + return CompiledArrayValue( + 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, + *, + never_skip_nulls=False, + skip_reproject_unsafe: bool = False, + ) -> CompiledArrayValue: + """ + 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 + 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_ibis_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 and not never_skip_nulls: + clauses.append((column.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( + ( + observation_count < 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_sql( + self, + 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 = ibis_bigquery.Backend().compile( + self._to_ibis_expr( + 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["string_encoded", "offset_col", "unordered"], + 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. + + ArrayValue objects are sorted, so the following options are available + to reflect this in the ibis expression. + + * "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. + * "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``. + 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 ( + "string_encoded", + "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}" + ) + + columns = list(self._columns) + 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) + # 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()) + + order_columns = self._create_order_columns( + ordering_mode, order_col_name, expose_hidden_cols + ) + columns.extend(order_columns) + + # Special case for empty tables, since we can't create an empty + # projection. + if not columns: + return ibis.memtable([]) + + # Make sure all dtypes are the "canonical" ones for BigFrames. This is + # important for operations like UNION where the schema must match. + 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(base_table[PREDICATE_COLUMN]) + 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 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 _reproject_to_table(self) -> CompiledArrayValue: + """ + 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( + "unordered", + 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 + if column_name in ordering_col_ids + ] + return CompiledArrayValue( + 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_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 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._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._ibis_order) + 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 unpivot( + self, + row_labels: typing.Sequence[typing.Hashable], + unpivot_columns: typing.Sequence[ + typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + ], + *, + passthrough_columns: typing.Sequence[str] = (), + index_col_ids: typing.Sequence[str] = ["index"], + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + ] = pandas.Float64Dtype(), + how="left", + ) -> CompiledArrayValue: + """ + 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 + """ + if how not in ("left", "right"): + raise ValueError("'how' must be 'left' or 'right'") + 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( + len(source_columns) == row_n for _, source_columns in unpivot_columns + ): + raise ValueError("Columns and row labels must all be same length.") + + unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") + unpivot_table = table.cross_join( + ibis.memtable({unpivot_offset_id: range(row_n)}) + ) + # Use ibis memtable to infer type of rowlabels (if possible) + # TODO: Allow caller to specify dtype + 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 + ] + + 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)): + 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 + ) + unpivot_values.append(unpivot_value.name(result_col)) + + unpivot_table = unpivot_table.select( + 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=tuple( + [ + *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=tuple( + [ + 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 CompiledArrayValue( + table=unpivot_table, + 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) -> CompiledArrayValue: + return self._set_or_replace_by_id( + destination_id, self._get_ibis_column(source_id) + ) + + def assign_constant( + self, + destination_id: str, + value: typing.Any, + dtype: typing.Optional[bigframes.dtypes.Dtype], + ) -> CompiledArrayValue: + # 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)}. {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 + ) -> CompiledArrayValue: + """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_ids: + builder.columns = [ + val if (col_id != id) else new_value.name(id) + for col_id, val in zip(self.column_ids, self._columns) + ] + else: + builder.columns = [*self.columns, new_value.name(id)] + return builder.build() + + +class ArrayValueBuilder: + """Mutable expression class. + Use ArrayValue.builder() to create from a ArrayValue object. + """ + + def __init__( + self, + table: ibis_types.Table, + ordering: ExpressionOrdering, + columns: Collection[ibis_types.Value] = (), + hidden_ordering_columns: Collection[ibis_types.Value] = (), + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + 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) -> CompiledArrayValue: + return CompiledArrayValue( + 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/compile/compiler.py b/bigframes/core/compile/compiler.py new file mode 100644 index 00000000000..195d830122a --- /dev/null +++ b/bigframes/core/compile/compiler.py @@ -0,0 +1,185 @@ +# 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 io +import typing + +import pandas as pd + +import bigframes.core.compile as compiled +import bigframes.core.compile.single_column +import bigframes.core.nodes as nodes + +if typing.TYPE_CHECKING: + import bigframes.core + import bigframes.session + + +@functools.cache +def compile_node(node: nodes.BigFrameNode) -> compiled.CompiledArrayValue: + """Compile node into CompileArrayValue. Caches result.""" + return _compile_node(node) + + +@functools.singledispatch +def _compile_node(node: nodes.BigFrameNode) -> compiled.CompiledArrayValue: + """Defines transformation but isn't cached, always use compile_node instead""" + raise ValueError(f"Can't compile unnrecognized node: {node}") + + +@_compile_node.register +def compile_join(node: nodes.JoinNode): + compiled_left = compile_node(node.left_child) + compiled_right = compile_node(node.right_child) + return bigframes.core.compile.single_column.join_by_column( + compiled_left, + node.left_column_ids, + compiled_right, + node.right_column_ids, + how=node.how, + allow_row_identity_join=node.allow_row_identity_join, + ) + + +@_compile_node.register +def compile_select(node: nodes.SelectNode): + return compile_node(node.child).select_columns(node.column_ids) + + +@_compile_node.register +def compile_drop(node: nodes.DropColumnsNode): + return compile_node(node.child).drop_columns(node.columns) + + +@_compile_node.register +def compile_readlocal(node: nodes.ReadLocalNode): + array_as_pd = pd.read_feather(io.BytesIO(node.feather_bytes)) + return compiled.CompiledArrayValue.mem_expr_from_pandas(array_as_pd) + + +@_compile_node.register +def compile_readgbq(node: nodes.ReadGbqNode): + return compiled.CompiledArrayValue( + node.table, + node.columns, + node.hidden_ordering_columns, + node.ordering, + ) + + +@_compile_node.register +def compile_promote_offsets(node: nodes.PromoteOffsetsNode): + return compile_node(node.child).promote_offsets(node.col_id) + + +@_compile_node.register +def compile_filter(node: nodes.FilterNode): + return compile_node(node.child).filter(node.predicate_id, node.keep_null) + + +@_compile_node.register +def compile_orderby(node: nodes.OrderByNode): + return compile_node(node.child).order_by(node.by, node.stable) + + +@_compile_node.register +def compile_reversed(node: nodes.ReversedNode): + return compile_node(node.child).reversed() + + +@_compile_node.register +def compile_project_unary(node: nodes.ProjectUnaryOpNode): + return compile_node(node.child).project_unary_op( + node.input_id, node.op, node.output_id + ) + + +@_compile_node.register +def compile_project_binary(node: nodes.ProjectBinaryOpNode): + return compile_node(node.child).project_binary_op( + node.left_input_id, node.right_input_id, node.op, node.output_id + ) + + +@_compile_node.register +def compile_project_ternary(node: nodes.ProjectTernaryOpNode): + return compile_node(node.child).project_ternary_op( + node.input_id1, node.input_id2, node.input_id3, node.op, node.output_id + ) + + +@_compile_node.register +def compile_concat(node: nodes.ConcatNode): + compiled_nodes = [compile_node(node) for node in node.children] + return compiled_nodes[0].concat(compiled_nodes[1:]) + + +@_compile_node.register +def compile_aggregate(node: nodes.AggregateNode): + return compile_node(node.child).aggregate( + node.aggregations, node.by_column_ids, node.dropna + ) + + +@_compile_node.register +def compile_corr(node: nodes.CorrNode): + return compile_node(node.child).corr_aggregate(node.corr_aggregations) + + +@_compile_node.register +def compile_window(node: nodes.WindowOpNode): + return compile_node(node.child).project_window_op( + node.column_name, + node.op, + node.window_spec, + node.output_name, + never_skip_nulls=node.never_skip_nulls, + skip_reproject_unsafe=node.skip_reproject_unsafe, + ) + + +@_compile_node.register +def compile_reproject(node: nodes.ReprojectOpNode): + return compile_node(node.child)._reproject_to_table() + + +@_compile_node.register +def compile_unpivot(node: nodes.UnpivotNode): + return compile_node(node.child).unpivot( + node.row_labels, + node.unpivot_columns, + passthrough_columns=node.passthrough_columns, + index_col_ids=node.index_col_ids, + dtype=node.dtype, + how=node.how, + ) + + +@_compile_node.register +def compile_assign(node: nodes.AssignNode): + return compile_node(node.child).assign(node.source_id, node.destination_id) + + +@_compile_node.register +def compile_assign_constant(node: nodes.AssignConstantNode): + return compile_node(node.child).assign_constant( + node.destination_id, node.value, node.dtype + ) + + +@_compile_node.register +def compiler_random_sample(node: nodes.RandomSampleNode): + return compile_node(node.child)._uniform_sampling(node.fraction) diff --git a/bigframes/core/joins/row_identity.py b/bigframes/core/compile/row_identity.py similarity index 94% rename from bigframes/core/joins/row_identity.py rename to bigframes/core/compile/row_identity.py index 76e456ec942..2e9bc0527ca 100644 --- a/bigframes/core/joins/row_identity.py +++ b/bigframes/core/compile/row_identity.py @@ -23,15 +23,16 @@ import ibis.expr.types as ibis_types import bigframes.constants as constants -import bigframes.core as core +import bigframes.core.compile as compiled import bigframes.core.joins.name_resolution as naming +import bigframes.core.ordering as orderings SUPPORTED_ROW_IDENTITY_HOW = {"outer", "left", "inner"} def join_by_row_identity( - left: core.ArrayValue, right: core.ArrayValue, *, how: str -) -> core.ArrayValue: + left: compiled.CompiledArrayValue, right: compiled.CompiledArrayValue, *, how: str +) -> compiled.CompiledArrayValue: """Compute join when we are joining by row identity not a specific column.""" if how not in SUPPORTED_ROW_IDENTITY_HOW: raise NotImplementedError( @@ -101,8 +102,8 @@ def join_by_row_identity( ) # 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 + new_ordering = orderings.ExpressionOrdering( + tuple(ordering_columns), total_ordering_columns=join_total_order_cols ) hidden_ordering_columns = [ @@ -117,8 +118,7 @@ def join_by_row_identity( if key.column_id in right._hidden_ordering_column_names.keys() ] - joined_expr = core.ArrayValue( - left._session, + joined_expr = compiled.CompiledArrayValue( left._table, columns=joined_columns, hidden_ordering_columns=hidden_ordering_columns, diff --git a/bigframes/core/joins/single_column.py b/bigframes/core/compile/single_column.py similarity index 87% rename from bigframes/core/joins/single_column.py rename to bigframes/core/compile/single_column.py index 0c0e2008b58..b992aa1d1d9 100644 --- a/bigframes/core/joins/single_column.py +++ b/bigframes/core/compile/single_column.py @@ -23,16 +23,16 @@ import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types -import bigframes.core as core -import bigframes.core.joins.name_resolution as naming -import bigframes.core.joins.row_identity -import bigframes.core.ordering +import bigframes.core.compile as compiled +import bigframes.core.compile.row_identity +import bigframes.core.joins as joining +import bigframes.core.ordering as orderings def join_by_column( - left: core.ArrayValue, + left: compiled.CompiledArrayValue, left_column_ids: typing.Sequence[str], - right: core.ArrayValue, + right: compiled.CompiledArrayValue, right_column_ids: typing.Sequence[str], *, how: Literal[ @@ -42,7 +42,7 @@ def join_by_column( "right", ], allow_row_identity_join: bool = True, -) -> core.ArrayValue: +) -> compiled.CompiledArrayValue: """Join two expressions by column equality. Arguments: @@ -61,7 +61,7 @@ def join_by_column( """ if ( allow_row_identity_join - and how in bigframes.core.joins.row_identity.SUPPORTED_ROW_IDENTITY_HOW + and how in bigframes.core.compile.row_identity.SUPPORTED_ROW_IDENTITY_HOW 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 @@ -73,15 +73,15 @@ def join_by_column( for lcol, rcol in zip(left_column_ids, right_column_ids) ) ): - return bigframes.core.joins.row_identity.join_by_row_identity( + return bigframes.core.compile.row_identity.join_by_row_identity( left, right, how=how ) else: # 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( + l_public_mapping, r_public_mapping = joining.JOIN_NAME_REMAPPER( left.column_ids, right.column_ids ) - l_hidden_mapping, r_hidden_mapping = naming.JoinNameRemapper( + l_hidden_mapping, r_hidden_mapping = joining.JoinNameRemapper( namespace="hidden" )(left._hidden_column_ids, right._hidden_column_ids) l_mapping = {**l_public_mapping, **l_hidden_mapping} @@ -134,8 +134,7 @@ def join_by_column( for col in right._hidden_ordering_columns ], ] - return core.ArrayValue( - left._session, + return compiled.CompiledArrayValue( combined_table, columns=columns, hidden_ordering_columns=hidden_ordering_columns, @@ -151,12 +150,12 @@ def value_to_join_key(value: ibis_types.Value): def join_orderings( - left: core.ExpressionOrdering, - right: core.ExpressionOrdering, + left: orderings.ExpressionOrdering, + right: orderings.ExpressionOrdering, left_id_mapping: Mapping[str, str], right_id_mapping: Mapping[str, str], left_order_dominates: bool = True, -) -> core.ExpressionOrdering: +) -> orderings.ExpressionOrdering: left_ordering_refs = [ ref.with_name(left_id_mapping[ref.column_id]) for ref in left.all_ordering_columns @@ -176,7 +175,7 @@ def join_orderings( right_total_order_cols = frozenset( [right_id_mapping[id] for id in right.total_ordering_columns] ) - return core.ExpressionOrdering( - ordering_value_columns=joined_refs, + return orderings.ExpressionOrdering( + ordering_value_columns=tuple(joined_refs), total_ordering_columns=left_total_order_cols | right_total_order_cols, ) diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index db0843fcbc0..2a19a83dd5f 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -193,7 +193,7 @@ def cumprod(self, *args, **kwargs) -> df.DataFrame: def shift(self, periods=1) -> series.Series: window = core.WindowSpec( - grouping_keys=self._by_col_ids, + grouping_keys=tuple(self._by_col_ids), preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -201,7 +201,7 @@ def shift(self, periods=1) -> series.Series: def diff(self, periods=1) -> series.Series: window = core.WindowSpec( - grouping_keys=self._by_col_ids, + grouping_keys=tuple(self._by_col_ids), preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -210,7 +210,7 @@ def diff(self, periods=1) -> series.Series: 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, + grouping_keys=tuple(self._by_col_ids), preceding=window - 1, following=0, min_periods=min_periods or window, @@ -225,7 +225,7 @@ 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, + grouping_keys=tuple(self._by_col_ids), following=0, min_periods=min_periods, ) @@ -389,7 +389,7 @@ def _apply_window_op( ): """Apply window op to groupby. Defaults to grouped cumulative window.""" window_spec = window or core.WindowSpec( - grouping_keys=self._by_col_ids, following=0 + grouping_keys=tuple(self._by_col_ids), following=0 ) columns = self._aggregated_columns(numeric_only=numeric_only) block, result_ids = self._block.multi_apply_window_op( @@ -528,7 +528,7 @@ def cumcount(self, *args, **kwargs) -> series.Series: def shift(self, periods=1) -> series.Series: """Shift index by desired number of periods.""" window = core.WindowSpec( - grouping_keys=self._by_col_ids, + grouping_keys=tuple(self._by_col_ids), preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -536,7 +536,7 @@ def shift(self, periods=1) -> series.Series: def diff(self, periods=1) -> series.Series: window = core.WindowSpec( - grouping_keys=self._by_col_ids, + grouping_keys=tuple(self._by_col_ids), preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -545,7 +545,7 @@ def diff(self, periods=1) -> series.Series: 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, + grouping_keys=tuple(self._by_col_ids), preceding=window - 1, following=0, min_periods=min_periods or window, @@ -564,7 +564,7 @@ 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, + grouping_keys=tuple(self._by_col_ids), following=0, min_periods=min_periods, ) @@ -597,7 +597,7 @@ def _apply_window_op( ): """Apply window op to groupby. Defaults to grouped cumulative window.""" window_spec = window or core.WindowSpec( - grouping_keys=self._by_col_ids, following=0 + grouping_keys=tuple(self._by_col_ids), following=0 ) label = self._value_name if not discard_name else None diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index d18a0a38efe..f6ce084714a 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -311,7 +311,7 @@ def _loc_getitem_series_or_dataframe( values = [entry[i] for entry in key] index_cols_dict[index_name] = values keys_df = bigframes.dataframe.DataFrame( - index_cols_dict, session=series_or_dataframe._get_block().expr._session + index_cols_dict, session=series_or_dataframe._get_block().expr.session ) keys_df = keys_df.set_index(temporary_index_names, drop=True) keys_df = keys_df.rename_axis(original_index_names) @@ -324,7 +324,7 @@ def _loc_getitem_series_or_dataframe( index_name = "unnamed_col" keys_df = bigframes.dataframe.DataFrame( {index_name: key}, - session=series_or_dataframe._get_block().expr._session, + session=series_or_dataframe._get_block().expr.session, ) keys_df = keys_df.set_index(index_name, drop=True) if index_name_is_none: @@ -343,7 +343,7 @@ def _loc_getitem_series_or_dataframe( 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 + {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 diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index b9ffdff21ec..6c66c36062a 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -26,8 +26,7 @@ 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.joins.name_resolution as join_names +import bigframes.core.joins as joining import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.dtypes @@ -402,7 +401,7 @@ def to_pandas(self) -> pandas.Index: dtypes = dict(zip(index_columns, self.dtypes)) expr = self._expr.select_columns(index_columns) results, _ = expr.start_query() - df = expr._session._rows_to_dataframe(results, dtypes) + df = expr.session._rows_to_dataframe(results, dtypes) df = df.set_index(index_columns) index = df.index index.names = list(self._block._index_labels) @@ -461,11 +460,10 @@ def join_mono_indexed( ) -> 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( + get_column_left, get_column_right = joining.JOIN_NAME_REMAPPER( left_expr.column_ids, right_expr.column_ids ) - combined_expr = joins.join_by_column( - left._block.expr, + combined_expr = left._block.expr.join( left._block.index_columns, right._block.expr, right._block.index_columns, @@ -520,12 +518,11 @@ def join_multi_indexed( left_expr = left._block.expr right_expr = right._block.expr - get_column_left, get_column_right = join_names.JOIN_NAME_REMAPPER( + get_column_left, get_column_right = joining.JOIN_NAME_REMAPPER( left_expr.column_ids, right_expr.column_ids ) - combined_expr = joins.join_by_column( - left_expr, + combined_expr = left_expr.join( left_join_ids, right_expr, right_join_ids, diff --git a/bigframes/core/joins/__init__.py b/bigframes/core/joins/__init__.py index 3f9447aef0d..5d407ec22b2 100644 --- a/bigframes/core/joins/__init__.py +++ b/bigframes/core/joins/__init__.py @@ -15,11 +15,6 @@ """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 +from bigframes.core.joins.name_resolution import JOIN_NAME_REMAPPER, JoinNameRemapper -__all__ = ( - "join_by_row_identity", - "join_by_column", - "merge", -) +__all__ = ("merge", "JoinNameRemapper", "JOIN_NAME_REMAPPER") diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py new file mode 100644 index 00000000000..7b252b164f6 --- /dev/null +++ b/bigframes/core/nodes.py @@ -0,0 +1,245 @@ +# 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, field +import functools +import typing +from typing import Optional, Tuple + +import pandas + +import bigframes.core.guid +from bigframes.core.ordering import OrderingColumnReference +import bigframes.core.window_spec as window +import bigframes.dtypes +import bigframes.operations as ops +import bigframes.operations.aggregations as agg_ops + +if typing.TYPE_CHECKING: + import ibis.expr.types as ibis_types + + import bigframes.core.ordering as orderings + import bigframes.session + + +@dataclass(frozen=True) +class BigFrameNode: + """ + Immutable node for representing 2D typed array as a tree of operators. + + All subclasses must be hashable so as to be usable as caching key. + """ + + @property + def deterministic(self) -> bool: + """Whether this node will evaluates deterministically.""" + return True + + @property + def child_nodes(self) -> typing.Sequence[BigFrameNode]: + """Direct children of this node""" + return tuple([]) + + @functools.cached_property + def session(self): + sessions = [] + for child in self.child_nodes: + if child.session is not None: + sessions.append(child.session) + unique_sessions = len(set(sessions)) + if unique_sessions > 1: + raise ValueError("Cannot use combine sources from multiple sessions.") + elif unique_sessions == 1: + return sessions[0] + return None + + +@dataclass(frozen=True) +class UnaryNode(BigFrameNode): + child: BigFrameNode + + @property + def child_nodes(self) -> typing.Sequence[BigFrameNode]: + return (self.child,) + + +@dataclass(frozen=True) +class JoinNode(BigFrameNode): + left_child: BigFrameNode + right_child: BigFrameNode + left_column_ids: typing.Tuple[str, ...] + right_column_ids: typing.Tuple[str, ...] + how: typing.Literal[ + "inner", + "left", + "outer", + "right", + ] + allow_row_identity_join: bool = True + + @property + def child_nodes(self) -> typing.Sequence[BigFrameNode]: + return (self.left_child, self.right_child) + + +@dataclass(frozen=True) +class ConcatNode(BigFrameNode): + children: Tuple[BigFrameNode, ...] + + @property + def child_nodes(self) -> typing.Sequence[BigFrameNode]: + return self.children + + +# Input Nodex +@dataclass(frozen=True) +class ReadLocalNode(BigFrameNode): + feather_bytes: bytes + column_ids: typing.Tuple[str, ...] + + +# TODO: Refactor to take raw gbq object reference +@dataclass(frozen=True) +class ReadGbqNode(BigFrameNode): + table: ibis_types.Table = field() + table_session: bigframes.session.Session = field() + columns: Tuple[ibis_types.Value, ...] = field() + hidden_ordering_columns: Tuple[ibis_types.Value, ...] = field() + ordering: orderings.ExpressionOrdering = field() + + @property + def session(self): + return (self.table_session,) + + +# Unary nodes +@dataclass(frozen=True) +class DropColumnsNode(UnaryNode): + columns: Tuple[str, ...] + + +@dataclass(frozen=True) +class PromoteOffsetsNode(UnaryNode): + col_id: str + + +@dataclass(frozen=True) +class FilterNode(UnaryNode): + predicate_id: str + keep_null: bool = False + + +@dataclass(frozen=True) +class OrderByNode(UnaryNode): + by: Tuple[OrderingColumnReference, ...] + stable: bool = False + + +@dataclass(frozen=True) +class ReversedNode(UnaryNode): + pass + + +@dataclass(frozen=True) +class SelectNode(UnaryNode): + column_ids: typing.Tuple[str, ...] + + +@dataclass(frozen=True) +class ProjectUnaryOpNode(UnaryNode): + input_id: str + op: ops.UnaryOp + output_id: Optional[str] = None + + +@dataclass(frozen=True) +class ProjectBinaryOpNode(UnaryNode): + left_input_id: str + right_input_id: str + op: ops.BinaryOp + output_id: str + + +@dataclass(frozen=True) +class ProjectTernaryOpNode(UnaryNode): + input_id1: str + input_id2: str + input_id3: str + op: ops.TernaryOp + output_id: str + + +@dataclass(frozen=True) +class AggregateNode(UnaryNode): + aggregations: typing.Tuple[typing.Tuple[str, agg_ops.AggregateOp, str], ...] + by_column_ids: typing.Tuple[str, ...] = tuple([]) + dropna: bool = True + + +# TODO: Unify into aggregate +@dataclass(frozen=True) +class CorrNode(UnaryNode): + corr_aggregations: typing.Tuple[typing.Tuple[str, str, str], ...] + + +@dataclass(frozen=True) +class WindowOpNode(UnaryNode): + column_name: str + op: agg_ops.WindowOp + window_spec: window.WindowSpec + output_name: typing.Optional[str] = None + never_skip_nulls: bool = False + skip_reproject_unsafe: bool = False + + +@dataclass(frozen=True) +class ReprojectOpNode(UnaryNode): + pass + + +@dataclass(frozen=True) +class UnpivotNode(UnaryNode): + row_labels: typing.Tuple[typing.Hashable, ...] + unpivot_columns: typing.Tuple[ + typing.Tuple[str, typing.Tuple[typing.Optional[str], ...]], ... + ] + passthrough_columns: typing.Tuple[str, ...] = () + index_col_ids: typing.Tuple[str, ...] = ("index",) + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Tuple[bigframes.dtypes.Dtype, ...] + ] = (pandas.Float64Dtype(),) + how: typing.Literal["left", "right"] = "left" + + +@dataclass(frozen=True) +class AssignNode(UnaryNode): + source_id: str + destination_id: str + + +@dataclass(frozen=True) +class AssignConstantNode(UnaryNode): + destination_id: str + value: typing.Hashable + dtype: typing.Optional[bigframes.dtypes.Dtype] + + +@dataclass(frozen=True) +class RandomSampleNode(UnaryNode): + fraction: float + + @property + def deterministic(self) -> bool: + return False diff --git a/bigframes/core/ordering.py b/bigframes/core/ordering.py index d5f07ecf91a..2cecd2fe7b2 100644 --- a/bigframes/core/ordering.py +++ b/bigframes/core/ordering.py @@ -86,7 +86,7 @@ class IntegerEncoding: class ExpressionOrdering: """Immutable object that holds information about the ordering of rows in a ArrayValue object.""" - ordering_value_columns: Sequence[OrderingColumnReference] = () + ordering_value_columns: typing.Tuple[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. @@ -170,7 +170,7 @@ def with_column_remap(self, mapping: typing.Mapping[str, str]): mapping.get(col_id, col_id) for col_id in self.total_ordering_columns ) return ExpressionOrdering( - new_value_columns, + tuple(new_value_columns), integer_encoding=self.integer_encoding, string_encoding=self.string_encoding, total_ordering_columns=new_total_order, diff --git a/bigframes/core/window_spec.py b/bigframes/core/window_spec.py new file mode 100644 index 00000000000..3458bfb1b8c --- /dev/null +++ b/bigframes/core/window_spec.py @@ -0,0 +1,35 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +import typing + +import bigframes.core.ordering as orderings + + +@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.Tuple[str, ...] = tuple() + ordering: typing.Tuple[orderings.OrderingColumnReference, ...] = tuple() + preceding: typing.Optional[int] = None + following: typing.Optional[int] = None + min_periods: int = 0 diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 3fd83198768..9d22c02d871 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -170,9 +170,7 @@ def __init__( if isinstance(dt, pandas.ArrowDtype) ) ): - self._block = blocks.block_from_local( - pd_dataframe, session or bigframes.pandas.get_global_session() - ) + self._block = blocks.block_from_local(pd_dataframe) elif session: self._block = session.read_pandas(pd_dataframe)._get_block() else: @@ -299,7 +297,7 @@ def values(self) -> numpy.ndarray: @property def _session(self) -> bigframes.Session: - return self._get_block().expr._session + return self._get_block().expr.session def __len__(self): rows, _ = self.shape @@ -1107,7 +1105,7 @@ def _assign_single_item( ) local_df = bigframes.dataframe.DataFrame( - {k: v}, session=self._get_block().expr._session + {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 @@ -2203,7 +2201,7 @@ def to_csv( field_delimiter=sep, header=header, ) - _, query_job = self._block.expr._session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query(export_data_statement) self._set_internal_query_job(query_job) def to_json( @@ -2245,7 +2243,7 @@ def to_json( format="JSON", export_options={}, ) - _, query_job = self._block.expr._session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query(export_data_statement) self._set_internal_query_job(query_job) def to_gbq( @@ -2274,7 +2272,7 @@ def to_gbq( write_disposition=dispositions[if_exists], destination=bigquery.table.TableReference.from_string( destination_table, - default_project=self._block.expr._session.bqclient.project, + default_project=self._block.expr.session.bqclient.project, ), ) @@ -2321,7 +2319,7 @@ def to_parquet( format="PARQUET", export_options=export_options, ) - _, query_job = self._block.expr._session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query(export_data_statement) self._set_internal_query_job(query_job) def to_dict( @@ -2464,7 +2462,7 @@ def _run_io_query( """Executes a query job presenting this dataframe and returns the destination table.""" expr = self._block.expr - session = expr._session + 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 diff --git a/bigframes/ml/metrics.py b/bigframes/ml/metrics.py index 3bcb621f745..5731b946ca9 100644 --- a/bigframes/ml/metrics.py +++ b/bigframes/ml/metrics.py @@ -96,7 +96,7 @@ def roc_curve( y_true_series, y_score_series = utils.convert_to_series(y_true, y_score) - session = y_true_series._block.expr._session + 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 diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index b9abb2cc038..d33befe4da6 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -94,9 +94,7 @@ def __init__( if isinstance(dt, pd.ArrowDtype) ) ): - self._block = blocks.block_from_local( - pd_dataframe, session or bigframes.pandas.get_global_session() - ) + self._block = blocks.block_from_local(pd_dataframe) elif session: self._block = session.read_pandas(pd_dataframe)._get_block() else: diff --git a/bigframes/series.py b/bigframes/series.py index c1914527834..37d00d16f32 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -29,7 +29,6 @@ import bigframes.constants as constants 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 @@ -43,6 +42,7 @@ import bigframes.core.scalar as scalars import bigframes.core.utils as utils import bigframes.core.window +import bigframes.core.window_spec import bigframes.dataframe import bigframes.dtypes import bigframes.formatting_helpers as formatter @@ -367,43 +367,43 @@ def between(self, left, right, inclusive="both"): def cumsum(self) -> Series: return self._apply_window_op( - agg_ops.sum_op, bigframes.core.WindowSpec(following=0) + agg_ops.sum_op, bigframes.core.window_spec.WindowSpec(following=0) ) def ffill(self, *, limit: typing.Optional[int] = None) -> Series: - window = bigframes.core.WindowSpec(preceding=limit, following=0) + window = bigframes.core.window_spec.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) + window = bigframes.core.window_spec.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) + agg_ops.max_op, bigframes.core.window_spec.WindowSpec(following=0) ) def cummin(self) -> Series: return self._apply_window_op( - agg_ops.min_op, bigframes.core.WindowSpec(following=0) + agg_ops.min_op, bigframes.core.window_spec.WindowSpec(following=0) ) def cumprod(self) -> Series: return self._apply_window_op( - agg_ops.product_op, bigframes.core.WindowSpec(following=0) + agg_ops.product_op, bigframes.core.window_spec.WindowSpec(following=0) ) def shift(self, periods: int = 1) -> Series: - window = bigframes.core.WindowSpec( + window = bigframes.core.window_spec.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, periods: int = 1) -> Series: - window = bigframes.core.WindowSpec( + window = bigframes.core.window_spec.WindowSpec( preceding=periods if periods > 0 else None, following=-periods if periods < 0 else None, ) @@ -805,7 +805,7 @@ def mode(self) -> Series: block, max_value_count_col_id = block.apply_window_op( value_count_col_id, agg_ops.max_op, - window_spec=WindowSpec(), + window_spec=bigframes.core.window_spec.WindowSpec(), ) block, is_mode_col_id = block.apply_binary_op( value_count_col_id, @@ -1009,9 +1009,7 @@ def _apply_aggregation(self, op: agg_ops.AggregateOp) -> Any: return self._block.get_stat(self._value_column, op) def _apply_window_op( - self, - op: agg_ops.WindowOp, - window_spec: bigframes.core.WindowSpec, + self, op: agg_ops.WindowOp, window_spec: bigframes.core.window_spec.WindowSpec ): block = self._block block, result_id = block.apply_window_op( @@ -1070,7 +1068,7 @@ def sort_index(self, *, axis=0, ascending=True, na_position="last") -> Series: 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( + window_spec = bigframes.core.window_spec.WindowSpec( preceding=window - 1, following=0, min_periods=min_periods or window ) return bigframes.core.window.Window( @@ -1078,7 +1076,9 @@ def rolling(self, window: int, min_periods=None) -> bigframes.core.window.Window ) def expanding(self, min_periods: int = 1) -> bigframes.core.window.Window: - window_spec = WindowSpec(following=0, min_periods=min_periods) + window_spec = bigframes.core.window_spec.WindowSpec( + following=0, min_periods=min_periods + ) return bigframes.core.window.Window( self._block, window_spec, self._block.value_columns, is_series=True ) @@ -1251,7 +1251,7 @@ def reindex(self, index=None, *, validate: typing.Optional[bool] = None): "Cannot reindex with index with different nlevels" ) new_indexer = bigframes.dataframe.DataFrame( - index=index, session=self._get_block().expr._session + index=index, session=self._get_block().expr.session )[[]] # multiindex join is senstive to index names, so we will set all these result = new_indexer.rename_axis(range(new_indexer.index.nlevels)).join( @@ -1415,7 +1415,7 @@ def map( elif isinstance(arg, Mapping): map_df = bigframes.dataframe.DataFrame( {"keys": list(arg.keys()), self.name: list(arg.values())}, - session=self._get_block().expr._session, + session=self._get_block().expr.session, ) map_df = map_df.set_index("keys") elif callable(arg): diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index af1f70d54d5..473de62f53a 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -68,6 +68,7 @@ import bigframes.core.blocks as blocks import bigframes.core.guid as guid from bigframes.core.ordering import IntegerEncoding, OrderingColumnReference +import bigframes.core.ordering as orderings import bigframes.core.utils as utils import bigframes.dataframe as dataframe import bigframes.formatting_helpers as formatting_helpers @@ -206,6 +207,10 @@ def _session_dataset_id(self): def _project(self): return self.bqclient.project + def __hash__(self): + # Stable hash needed to use in expression tree + return hash(self._session_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""" @@ -592,11 +597,13 @@ def _read_gbq_table( # primary key(s) are set on a table. The query engine assumes such # columns are unique, even if not enforced. is_total_ordering = True - ordering = core.ExpressionOrdering( - ordering_value_columns=[ - core.OrderingColumnReference(column_id) - for column_id in total_ordering_cols - ], + ordering = orderings.ExpressionOrdering( + ordering_value_columns=tuple( + [ + core.OrderingColumnReference(column_id) + for column_id in total_ordering_cols + ] + ), total_ordering_columns=frozenset(total_ordering_cols), ) @@ -634,10 +641,13 @@ def _read_gbq_table( 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 - ], + ordering = orderings.ExpressionOrdering( + ordering_value_columns=tuple( + [ + core.OrderingColumnReference(column_id) + for column_id in index_cols + ] + ), total_ordering_columns=frozenset(index_cols), ) @@ -713,7 +723,7 @@ def _read_gbq_with_ordering( index_cols: Iterable[str] = (), index_labels: Iterable[Optional[str]] = (), hidden_cols: Iterable[str] = (), - ordering: core.ExpressionOrdering, + ordering: orderings.ExpressionOrdering, is_total_ordering: bool = False, api_name: str, ) -> dataframe.DataFrame: @@ -826,7 +836,7 @@ def _read_ibis( index_labels: Iterable[blocks.Label], column_keys: Iterable[str], column_labels: Iterable[blocks.Label], - ordering: core.ExpressionOrdering, + ordering: orderings.ExpressionOrdering, ) -> dataframe.DataFrame: """Turns a table expression (plus index column) into a DataFrame.""" @@ -843,7 +853,7 @@ def _read_ibis( hidden_ordering_columns.append(table_expression[ref.column_id]) block = blocks.Block( - core.ArrayValue( + core.ArrayValue.from_ibis( self, table_expression, columns, hidden_ordering_columns, ordering ), index_columns=[index_col.get_name() for index_col in index_cols], @@ -959,8 +969,8 @@ def _read_pandas( ) self._start_generic_job(load_job) - ordering = core.ExpressionOrdering( - ordering_value_columns=[OrderingColumnReference(ordering_col)], + ordering = orderings.ExpressionOrdering( + ordering_value_columns=tuple([OrderingColumnReference(ordering_col)]), total_ordering_columns=frozenset([ordering_col]), integer_encoding=IntegerEncoding(True, is_sequential=True), ) @@ -1303,7 +1313,7 @@ def _create_sequential_ordering( table: ibis_types.Table, index_cols: Iterable[str] = (), api_name: str = "", - ) -> Tuple[ibis_types.Table, core.ExpressionOrdering]: + ) -> Tuple[ibis_types.Table, orderings.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_") @@ -1320,8 +1330,8 @@ def _create_sequential_ordering( 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], + ordering = orderings.ExpressionOrdering( + ordering_value_columns=tuple([ordering_reference]), total_ordering_columns=frozenset([default_ordering_name]), integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), ) diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index f7fc4eaa8f6..084b723fba6 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -98,7 +98,7 @@ def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex): def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): bf.options.display.progress_bar = "terminal" - penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = ( + penguins_df_default_index._block._expr.session.bqclient.default_query_job_config.use_query_cache = ( False ) penguins_df_default_index.to_pandas() @@ -117,7 +117,7 @@ def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): 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 = ( + penguins_df_default_index._block._expr.session.bqclient.default_query_job_config.use_query_cache = ( False ) penguins_df_default_index.to_pandas() diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index c9510290b66..05d8b841855 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2858,7 +2858,7 @@ def test_map_series_input(scalars_dfs): 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_map_series, session=scalars_df._get_block().expr.session ) pd_result = scalars_pandas_df.int64_too.map(pd_map_series) @@ -2877,7 +2877,7 @@ def test_map_series_input_duplicates_error(scalars_dfs): 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_map_series, session=scalars_df._get_block().expr.session ) with pytest.raises(pd.errors.InvalidIndexError): diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 127a88a760b..bf72e444eba 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -318,7 +318,6 @@ 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.to_pandas() expected = scalars_pandas_df @@ -350,9 +349,8 @@ def test_read_pandas_rowid_exists_adds_suffix(session, scalars_pandas_df_default pandas_df = scalars_pandas_df_default_index.copy() pandas_df["rowid"] = np.arange(pandas_df.shape[0]) - 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" + df_roundtrip = session.read_pandas(pandas_df).to_pandas() + pd.testing.assert_frame_equal(df_roundtrip, pandas_df, check_dtype=False) def test_read_pandas_tokyo( @@ -385,7 +383,6 @@ def test_read_csv_gcs_default_engine(session, scalars_dfs, gcs_folder): # 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) @@ -441,7 +438,6 @@ def test_read_csv_local_default_engine(session, scalars_dfs, 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) @@ -976,7 +972,6 @@ def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder): 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, diff --git a/tests/unit/core/test_blocks.py b/tests/unit/core/test_blocks.py index a7e9b5a84b7..86715d090cb 100644 --- a/tests/unit/core/test_blocks.py +++ b/tests/unit/core/test_blocks.py @@ -18,8 +18,6 @@ import bigframes.core.blocks as blocks -from .. import resources - @pytest.mark.parametrize( ("data",), @@ -76,9 +74,8 @@ ) def test_block_from_local(data): expected = pandas.DataFrame(data) - session = resources.create_pandas_session({}) - block = blocks.block_from_local(data, session=session) + block = blocks.block_from_local(data) pandas.testing.assert_index_equal(block.column_labels, expected.columns) assert tuple(block.index_labels) == tuple(expected.index.names) diff --git a/tests/unit/resources.py b/tests/unit/resources.py index 0a68600a359..f660d774f06 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -22,6 +22,7 @@ import bigframes import bigframes.core as core +import bigframes.core.ordering import bigframes.session.clients """Utilities for creating test resources.""" @@ -61,14 +62,20 @@ def create_pandas_session(tables: Dict[str, pandas.DataFrame]) -> bigframes.Sess def create_arrayvalue( df: pandas.DataFrame, total_ordering_columns: List[str] -) -> bigframes.core.ArrayValue: +) -> 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], + ordering = bigframes.core.ordering.ExpressionOrdering( + tuple( + [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 + return core.ArrayValue.from_ibis( + session=session, + table=ibis_table, + columns=columns, + hidden_ordering_columns=(), + ordering=ordering, ) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 69b9e798075..d9672b2635b 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -16,6 +16,7 @@ import pandas import bigframes.core as core +import bigframes.core.ordering import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -37,15 +38,19 @@ def test_arrayvalue_constructor_from_ibis_table_adds_all_columns(): ) ibis_table = session.ibis_client.table("test_table") columns = (ibis_table["col1"], ibis_table["col2"], ibis_table["col3"]) - ordering = core.ExpressionOrdering( - [core.OrderingColumnReference("col1")], + ordering = bigframes.core.ordering.ExpressionOrdering( + tuple([core.OrderingColumnReference("col1")]), total_ordering_columns=frozenset(["col1"]), ) - actual = core.ArrayValue( - session=session, table=ibis_table, columns=columns, ordering=ordering + actual = core.ArrayValue.from_ibis( + session=session, + table=ibis_table, + columns=columns, + ordering=ordering, + hidden_ordering_columns=(), ) - assert actual._table is ibis_table - assert len(actual.columns) == 3 + assert actual.compile()._table is ibis_table + assert len(actual.column_ids) == 3 def test_arrayvalue_with_get_column_type(): @@ -78,7 +83,7 @@ def test_arrayvalue_with_get_column(): ), total_ordering_columns=["col1"], ) - col1 = value._get_ibis_column("col1") + col1 = value.compile()._get_ibis_column("col1") assert isinstance(col1, ibis_types.Value) assert col1.get_name() == "col1" assert col1.type().is_int64() @@ -95,7 +100,7 @@ def test_arrayvalues_to_ibis_expr_with_get_column(): ), total_ordering_columns=["col1"], ) - expr = value._get_ibis_column("col1") + expr = value.compile()._get_ibis_column("col1") assert expr.get_name() == "col1" assert expr.type().is_int64() @@ -112,7 +117,7 @@ def test_arrayvalues_to_ibis_expr_with_concat(): total_ordering_columns=["col1"], ) expr = value.concat([value]) - actual = expr._to_ibis_expr("unordered") + actual = expr.compile()._to_ibis_expr("unordered") assert len(actual.columns) == 3 # TODO(ashleyxu, b/299631930): test out the union expression assert actual.columns[0] == "column_0" @@ -131,8 +136,8 @@ def test_arrayvalues_to_ibis_expr_with_project_unary_op(): ), total_ordering_columns=["col1"], ) - expr = value.project_unary_op("col1", ops.AsTypeOp("string")) - assert value.columns[0].type().is_int64() + expr = value.project_unary_op("col1", ops.AsTypeOp("string")).compile() + assert value.compile().columns[0].type().is_int64() assert expr.columns[0].type().is_string() @@ -147,7 +152,7 @@ def test_arrayvalues_to_ibis_expr_with_project_binary_op(): ), total_ordering_columns=["col1"], ) - expr = value.project_binary_op("col2", "col3", ops.add_op, "col4") + expr = value.project_binary_op("col2", "col3", ops.add_op, "col4").compile() assert expr.columns[3].type().is_float64() actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 4 @@ -166,7 +171,9 @@ def test_arrayvalues_to_ibis_expr_with_project_ternary_op(): ), total_ordering_columns=["col1"], ) - expr = value.project_ternary_op("col2", "col3", "col4", ops.where_op, "col5") + expr = value.project_ternary_op( + "col2", "col3", "col4", ops.where_op, "col5" + ).compile() assert expr.columns[4].type().is_float64() actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 5 @@ -188,7 +195,7 @@ def test_arrayvalue_to_ibis_expr_with_aggregate(): aggregations=(("col1", agg_ops.sum_op, "col4"),), by_column_ids=["col1"], dropna=False, - ) + ).compile() actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 2 assert actual.columns[0] == "col1" @@ -207,7 +214,7 @@ def test_arrayvalue_to_ibis_expr_with_corr_aggregate(): ), total_ordering_columns=["col1"], ) - expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]) + expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]).compile() actual = expr._to_ibis_expr("unordered") assert len(expr.columns) == 1 assert actual.columns[0] == "col4" From 29032d06811569121f7be2a7de915740df7daf6e Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Thu, 26 Oct 2023 18:20:58 -0700 Subject: [PATCH 118/197] fix: fix bug with column names under repeated column assignment (#150) 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/dataframe.py | 23 +++++++++++------------ tests/system/small/test_dataframe.py | 22 ++++++++++++++++++++-- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 9d22c02d871..01117d3e0ab 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1110,19 +1110,18 @@ 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_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, - 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" + new_column_block = local_df._block + original_index_column_ids = self._block.index_columns + self_block = self._block.reset_index(drop=False) + result_index, (get_column_left, get_column_right) = self_block.index.join( + new_column_block.index, how="left", block_identity_join=True ) - return join_result.drop(columns=[this_offsets_col_id]) + result_block = result_index._block + result_block = result_block.set_index( + [get_column_left[col_id] for col_id in original_index_column_ids], + index_labels=self._block.index_labels, + ) + return DataFrame(result_block) else: return self._assign_scalar(k, v) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a746a1867cf..e459e3bee3b 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -505,14 +505,32 @@ def test_assign_new_column_w_setitem_list(scalars_dfs): pd.testing.assert_frame_equal(bf_result, pd_result) +def test_assign_new_column_w_setitem_list_repeated(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_df["new_col_2"] = [1, 3, 2, 5, 4, 7, 6, 9, 8] + pd_df["new_col_2"] = [1, 3, 2, 5, 4, 7, 6, 9, 8] + 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_result["new_col_2"] = pd_result["new_col_2"].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") + pd_df = pd_df.set_index(["string_col", "int64_col"]) + bf_df = bf_df.set_index(["string_col", "int64_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] From 2d7128d9b2107c3667a5ad7f153d446bfdc04df5 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 27 Oct 2023 22:46:13 +0000 Subject: [PATCH 119/197] test: refactor remote function tests (#147) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This changes moves the tests that deploy cloud function to large remote function tests, and the tests that do not make call to bigquery service to unit tests. 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/large/test_remote_function.py | 90 ++++++++++++ tests/system/small/test_remote_function.py | 156 ++++----------------- tests/unit/test_remote_function.py | 28 ++++ 3 files changed, 148 insertions(+), 126 deletions(-) create mode 100644 tests/unit/test_remote_function.py diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index 730a1dbde4e..c8f8f66ebaa 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -27,6 +27,7 @@ import pytest import test_utils.prefixer +import bigframes from bigframes.remote_function import ( get_cloud_function_name, get_remote_function_locations, @@ -1120,3 +1121,92 @@ def plusone(x): ) for dir_ in dirs_to_cleanup: shutil.rmtree(dir_) + + +@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.bq_connection = bq_cf_connection + session = bigframes.connect(context) + + try: + # 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).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, session.cloudfunctionsclient, square + ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_default_connection(session, scalars_dfs, dataset_id): + try: + + @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).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, session.cloudfunctionsclient, square + ) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index d024a57dedc..89907a53dfc 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -13,14 +13,11 @@ # 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 import remote_function as rf -import bigframes.pandas as bpd from tests.system.utils import assert_pandas_df_equal_ignore_ordering @@ -65,45 +62,14 @@ 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(bq_connection=bq_cf_connection)) - - -@pytest.fixture(scope="module") -def session_with_bq_connection_location_specified( - bq_cf_connection_location, -) -> bigframes.Session: - return bigframes.Session( - bigframes.BigQueryOptions(bq_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(bq_connection=bq_cf_connection_location_mistached) - ) - - -@pytest.fixture(scope="module") -def session_with_bq_connection_location_project_specified( - bq_cf_connection_location_project, +def session_with_bq_connection_and_permanent_dataset( + bq_cf_connection, dataset_id_permanent ) -> bigframes.Session: - return bigframes.Session( - bigframes.BigQueryOptions(bq_connection=bq_cf_connection_location_project) + session = bigframes.Session( + bigframes.BigQueryOptions(bq_connection=bq_cf_connection) ) - - -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 + session._session_dataset = bigquery.Dataset(dataset_id_permanent) + return session @pytest.mark.flaky(retries=2, delay=120) @@ -311,11 +277,13 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_direct_session_param(session_with_bq_connection, scalars_dfs): +def test_remote_function_direct_session_param( + session_with_bq_connection_and_permanent_dataset, scalars_dfs +): @rf.remote_function( [int], int, - session=session_with_bq_connection, + session=session_with_bq_connection_and_permanent_dataset, ) def square(x): return x * x @@ -345,7 +313,9 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_via_session_default(session_with_bq_connection, scalars_dfs): +def test_remote_function_via_session_default( + session_with_bq_connection_and_permanent_dataset, 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 @@ -353,7 +323,7 @@ def test_remote_function_via_session_default(session_with_bq_connection, scalars # 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) + @session_with_bq_connection_and_permanent_dataset.remote_function([int], int) def square(x): return x * x @@ -421,87 +391,15 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_via_session_context_connection_setter( - scalars_dfs, dataset_id, bq_cf_connection +def test_dataframe_applymap( + session_with_bq_connection_and_permanent_dataset, scalars_dfs ): - # Creating a session scoped only to this test as we would be setting a - # property in it - context = bigframes.BigQueryOptions() - context.bq_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).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_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): return x + 1 - remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + remote_add_one = session_with_bq_connection_and_permanent_dataset.remote_function( + [int], int + )(add_one) scalars_df, scalars_pandas_df = scalars_dfs int64_cols = ["int64_col", "int64_too"] @@ -524,11 +422,15 @@ def add_one(x): @pytest.mark.flaky(retries=2, delay=120) -def test_dataframe_applymap_na_ignore(session_with_bq_connection, scalars_dfs): +def test_dataframe_applymap_na_ignore( + session_with_bq_connection_and_permanent_dataset, scalars_dfs +): def add_one(x): return x + 1 - remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + remote_add_one = session_with_bq_connection_and_permanent_dataset.remote_function( + [int], int + )(add_one) scalars_df, scalars_pandas_df = scalars_dfs int64_cols = ["int64_col", "int64_too"] @@ -549,11 +451,13 @@ def add_one(x): @pytest.mark.flaky(retries=2, delay=120) -def test_series_map(session_with_bq_connection, scalars_dfs): +def test_series_map(session_with_bq_connection_and_permanent_dataset, scalars_dfs): def add_one(x): return x + 1 - remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) + remote_add_one = session_with_bq_connection_and_permanent_dataset.remote_function( + [int], int + )(add_one) scalars_df, scalars_pandas_df = scalars_dfs @@ -635,7 +539,7 @@ def square1(x): @pytest.mark.flaky(retries=2, delay=120) -def test_read_gbq_function_reads_udfs(bigquery_client, scalars_dfs, dataset_id): +def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) arg = bigquery.RoutineArgument( name="x", diff --git a/tests/unit/test_remote_function.py b/tests/unit/test_remote_function.py new file mode 100644 index 00000000000..540f4020d36 --- /dev/null +++ b/tests/unit/test_remote_function.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. + +from ibis.backends.bigquery import datatypes as bq_types +from ibis.expr import datatypes as ibis_types + +from bigframes import remote_function as rf + + +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 From 4e4409c5b235171f3770aec852193026519948fd Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 27 Oct 2023 19:22:28 -0700 Subject: [PATCH 120/197] feat: add dataframe melt (#116) 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/core/blocks.py | 41 +++++++++++++++++- bigframes/dataframe.py | 38 ++++++++++++++++ tests/system/small/test_dataframe.py | 43 +++++++++++++++++++ tests/system/small/test_multiindex.py | 28 ++++++++++++ .../bigframes_vendored/pandas/core/frame.py | 28 ++++++++++++ 5 files changed, 176 insertions(+), 2 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index cc13edeaf96..635e7db8656 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -1356,13 +1356,50 @@ def stack(self, how="left", levels: int = 1): index_columns = [*added_index_columns, *self.index_columns] index_labels = [*new_index_level_names, *self._index_labels] - block = Block( + return Block( unpivot_expr, index_columns=index_columns, column_labels=result_index, index_labels=index_labels, ) - return block + + def melt( + self, + id_vars=typing.Sequence[str], + value_vars=typing.Sequence[str], + var_names=typing.Sequence[typing.Hashable], + value_name: typing.Hashable = "value", + ): + # TODO: Implement col_level and ignore_index + unpivot_col_id = guid.generate_guid() + var_col_ids = tuple([guid.generate_guid() for _ in var_names]) + # single unpivot col + unpivot_col = (unpivot_col_id, tuple(value_vars)) + value_labels = [self.col_id_to_label[col_id] for col_id in value_vars] + id_labels = [self.col_id_to_label[col_id] for col_id in id_vars] + + dtype = self._expr.get_column_type(value_vars[0]) + + unpivot_expr = self._expr.unpivot( + row_labels=value_labels, + passthrough_columns=id_vars, + unpivot_columns=(unpivot_col,), + index_col_ids=var_col_ids, + dtype=dtype, + how="right", + ) + index_id = guid.generate_guid() + unpivot_expr = unpivot_expr.promote_offsets(index_id) + # Need to reorder to get id_vars before var_col and unpivot_col + unpivot_expr = unpivot_expr.select_columns( + [index_id, *id_vars, *var_col_ids, unpivot_col_id] + ) + + return Block( + unpivot_expr, + column_labels=[*id_labels, *var_names, value_name], + index_columns=[index_id], + ) def _create_stack_column( self, col_label: typing.Tuple, stack_labels: typing.Sequence[typing.Tuple] diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 01117d3e0ab..49d7ad991a2 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1673,6 +1673,44 @@ def idxmin(self) -> bigframes.series.Series: def idxmax(self) -> bigframes.series.Series: return bigframes.series.Series(block_ops.idxmax(self._block)) + def melt( + self, + id_vars: typing.Optional[typing.Iterable[typing.Hashable]] = None, + value_vars: typing.Optional[typing.Iterable[typing.Hashable]] = None, + var_name: typing.Union[ + typing.Hashable, typing.Sequence[typing.Hashable] + ] = None, + value_name: typing.Hashable = "value", + ): + if var_name is None: + # Determine default var_name. Attempt to use column labels if they are unique + if self.columns.nlevels > 1: + if len(set(self.columns.names)) == len(self.columns.names): + var_name = self.columns.names + else: + var_name = [f"variable_{i}" for i in range(len(self.columns.names))] + else: + var_name = self.columns.name or "variable" + + var_name = tuple(var_name) if utils.is_list_like(var_name) else (var_name,) + + if id_vars is not None: + id_col_ids = [self._resolve_label_exact(col) for col in id_vars] + else: + id_col_ids = [] + if value_vars is not None: + val_col_ids = [self._resolve_label_exact(col) for col in value_vars] + else: + val_col_ids = [ + col_id + for col_id in self._block.value_columns + if col_id not in id_col_ids + ] + + return DataFrame( + self._block.melt(id_col_ids, val_col_ids, var_name, value_name) + ) + def describe(self) -> DataFrame: df_numeric = self._drop_non_numeric(keep_bool=False) if len(df_numeric.columns) == 0: diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index e459e3bee3b..b503f9a31d3 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1919,6 +1919,49 @@ def test_df_stack(scalars_dfs): pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False) +def test_df_melt_default(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].melt().to_pandas() + pd_result = scalars_pandas_df[columns].melt() + + # Pandas produces int64 index, Bigframes produces Int64 (nullable) + pd.testing.assert_frame_equal( + bf_result, pd_result, check_index_type=False, check_dtype=False + ) + + +def test_df_melt_parameterized(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 + + bf_result = scalars_df.melt( + var_name="alice", + value_name="bob", + id_vars=["string_col"], + value_vars=["int64_col", "int64_too"], + ).to_pandas() + pd_result = scalars_pandas_df.melt( + var_name="alice", + value_name="bob", + id_vars=["string_col"], + value_vars=["int64_col", "int64_too"], + ) + + # Pandas produces int64 index, Bigframes produces Int64 (nullable) + pd.testing.assert_frame_equal( + bf_result, pd_result, check_index_type=False, check_dtype=False + ) + + def test_df_unstack(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs # To match bigquery dataframes diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index a87dacae044..d6bf46f77c8 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -752,6 +752,34 @@ def test_column_multi_index_stack(level): ) +def test_column_multi_index_melt(): + 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"]) + level2 = pandas.Index(["a", "b", "b"]) + level3 = pandas.Index(["b", "b", "a"]) + + 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.melt().to_pandas() + pd_result = pd_df.melt() + + # BigFrames uses different string and int types, but values are identical + pandas.testing.assert_frame_equal( + bf_result, pd_result, check_index_type=False, check_dtype=False + ) + + 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]") diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 13a81b46454..67836a8fd29 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2010,6 +2010,34 @@ def idxmax(self): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def melt(self, id_vars, value_vars, var_name, value_name): + """ + Unpivot a DataFrame from wide to long format, optionally leaving identifiers set. + + This function is useful to massage a DataFrame into a format where one + or more columns are identifier variables (`id_vars`), while all other + columns, considered measured variables (`value_vars`), are "unpivoted" to + the row axis, leaving just two non-identifier columns, 'variable' and + 'value'. + + Parameters + ---------- + id_vars (tuple, list, or ndarray, optional): + Column(s) to use as identifier variables. + value_vars (tuple, list, or ndarray, optional): + Column(s) to unpivot. If not specified, uses all columns that + are not set as `id_vars`. + var_name (scalar): + Name to use for the 'variable' column. If None it uses + ``frame.columns.name`` or 'variable'. + value_name (scalar, default 'value'): + Name to use for the 'value' column. + + Returns: + DataFrame: Unpivoted DataFrame. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def nunique(self): """ Count number of distinct elements in specified axis. From ac44ccd3936cdb28755d2bbe16377d489f08d5e5 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Mon, 30 Oct 2023 10:07:58 -0700 Subject: [PATCH 121/197] docs: add artithmetic df sample code (#153) * docs: add artithmetic df sample code * fix: address comments --- bigframes/session/__init__.py | 4 +- .../bigframes_vendored/pandas/core/frame.py | 494 +++++++++++++++++- 2 files changed, 492 insertions(+), 6 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 473de62f53a..932a41f2830 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -352,7 +352,7 @@ def read_gbq_query( >>> import bigframes.pandas as bpd >>> bpd.options.display.progress_bar = None - Simple query input: + Simple query input: >>> df = bpd.read_gbq_query(''' ... SELECT @@ -368,7 +368,7 @@ def read_gbq_query( [2 rows x 3 columns] - Preserve ordering in a query input. + Preserve ordering in a query input. >>> df = bpd.read_gbq_query(''' ... SELECT diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 67836a8fd29..013d1701146 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -697,6 +697,7 @@ def align( Join method is specified for each axis Index. + Args: other (DataFrame or Series): join ({{'outer', 'inner', 'left', 'right'}}, default 'outer'): @@ -978,9 +979,9 @@ def sort_values( Sort ascending vs. descending. Specify list for multiple sort orders. If this is a list of bools, must match the length of the by. - kind (str, default `quicksort`): - Choice of sorting algorithm. Accepts 'quicksort’, ‘mergesort’, - ‘heapsort’, ‘stable’. Ignored except when determining whether to + 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 @@ -1014,6 +1015,29 @@ def eq(self, other, axis: str | int = "columns") -> DataFrame: Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis (rows or columns) and level for comparison. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can use method name: + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].eq(360) + circle True + triangle False + rectangle True + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``==``: + >>> df["degrees"] == 360 + circle True + triangle False + rectangle True + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1036,6 +1060,30 @@ def ne(self, other, axis: str | int = "columns") -> DataFrame: Equivalent to `==`, `!=`, `<=`, `<`, `>=`, `>` with support to choose axis (rows or columns) and level for comparison. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can use method name: + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].ne(360) + circle False + triangle True + rectangle False + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``!=``: + + >>> df["degrees"] != 360 + circle False + triangle True + rectangle False + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1061,6 +1109,30 @@ def le(self, other, axis: str | int = "columns") -> DataFrame: floating point columns are considered different (i.e. `NaN` != `NaN`). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can use method name: + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].le(180) + circle False + triangle True + rectangle False + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``<=``: + + >>> df["degrees"] <= 180 + circle False + triangle True + rectangle False + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1087,6 +1159,30 @@ def lt(self, other, axis: str | int = "columns") -> DataFrame: floating point columns are considered different (i.e. `NaN` != `NaN`). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can use method name: + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].lt(180) + circle False + triangle False + rectangle False + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``<``: + + >>> df["degrees"] < 180 + circle False + triangle False + rectangle False + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1113,6 +1209,30 @@ def ge(self, other, axis: str | int = "columns") -> DataFrame: floating point columns are considered different (i.e. `NaN` != `NaN`). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can use method name: + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].ge(360) + circle True + triangle False + rectangle True + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``>=``: + + >>> df["degrees"] >= 360 + circle True + triangle False + rectangle True + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1139,6 +1259,28 @@ def gt(self, other, axis: str | int = "columns") -> DataFrame: floating point columns are considered different (i.e. `NaN` != `NaN`). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'angles': [0, 3, 4], + ... 'degrees': [360, 180, 360]}, + ... index=['circle', 'triangle', 'rectangle']) + >>> df["degrees"].gt(360) + circle False + triangle False + rectangle False + Name: degrees, dtype: boolean + + You can also use arithmetic operator ``>``: + + >>> df["degrees"] > 360 + circle False + triangle False + rectangle False + Name: degrees, dtype: boolean + Args: other (scalar, sequence, Series, or DataFrame): Any single or multiple element data structure, or list-like object. @@ -1162,6 +1304,32 @@ def add(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].add(df['B']) + 0 5 + 1 7 + 2 9 + dtype: Int64 + + You can also use arithmetic operator ``+``: + + >>> df['A'] + (df['B']) + 0 5 + 1 7 + 2 9 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1185,6 +1353,32 @@ def sub(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].sub(df['B']) + 0 -3 + 1 -3 + 2 -3 + dtype: Int64 + + You can also use arithmetic operator ``-``: + + >>> df['A'] - (df['B']) + 0 -3 + 1 -3 + 2 -3 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1208,6 +1402,29 @@ def rsub(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df['A'].rsub(df['B']) + 0 3 + 1 3 + 2 3 + dtype: Int64 + + It's equivalent to using arithmetic operator: ``-``: + + >>> df['B'] - (df['A']) + 0 3 + 1 3 + 2 3 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1231,6 +1448,32 @@ def mul(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].mul(df['B']) + 0 4 + 1 10 + 2 18 + dtype: Int64 + + You can also use arithmetic operator ``*``: + + >>> df['A'] * (df['B']) + 0 4 + 1 10 + 2 18 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1254,6 +1497,32 @@ def truediv(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].truediv(df['B']) + 0 0.25 + 1 0.4 + 2 0.5 + dtype: Float64 + + You can also use arithmetic operator ``/``: + + >>> df['A'] / (df['B']) + 0 0.25 + 1 0.4 + 2 0.5 + dtype: Float64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1277,6 +1546,29 @@ def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df['A'].rtruediv(df['B']) + 0 4.0 + 1 2.5 + 2 2.0 + dtype: Float64 + + It's equivalent to using arithmetic operator: ``/``: + + >>> df['B'] / (df['A']) + 0 4.0 + 1 2.5 + 2 2.0 + dtype: Float64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1300,6 +1592,32 @@ def floordiv(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].floordiv(df['B']) + 0 0 + 1 0 + 2 0 + dtype: Int64 + + You can also use arithmetic operator ``//``: + + >>> df['A'] // (df['B']) + 0 0 + 1 0 + 2 0 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1323,6 +1641,29 @@ def rfloordiv(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df['A'].rfloordiv(df['B']) + 0 4 + 1 2 + 2 2 + dtype: Int64 + + It's equivalent to using arithmetic operator: ``//``: + + >>> df['B'] // (df['A']) + 0 4 + 1 2 + 2 2 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1346,6 +1687,32 @@ def mod(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].mod(df['B']) + 0 1 + 1 2 + 2 3 + dtype: Int64 + + You can also use arithmetic operator ``%``: + + >>> df['A'] % (df['B']) + 0 1 + 1 2 + 2 3 + dtype: Int64 + Args: other: Any single or multiple element data structure, or list-like object. @@ -1369,6 +1736,29 @@ def rmod(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df['A'].rmod(df['B']) + 0 0 + 1 1 + 2 0 + dtype: Int64 + + It's equivalent to using arithmetic operator: ``%``: + + >>> df['B'] % (df['A']) + 0 0 + 1 1 + 2 0 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1382,7 +1772,7 @@ 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`). + """Get Exponential power of dataframe and other, element-wise (binary operator `**`). Equivalent to ``dataframe ** other``, but with support to substitute a fill_value for missing data in one of the inputs. With reverse version, `rpow`. @@ -1393,6 +1783,32 @@ def pow(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + + You can use method name: + + >>> df['A'].pow(df['B']) + 0 1 + 1 32 + 2 729 + dtype: Int64 + + You can also use arithmetic operator ``**``: + + >>> df['A'] ** (df['B']) + 0 1 + 1 32 + 2 729 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1417,6 +1833,29 @@ def rpow(self, other, axis: str | int = "columns") -> DataFrame: .. note:: Mismatched indices will be unioned together. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df['A'].rpow(df['B']) + 0 4 + 1 25 + 2 216 + dtype: Int64 + + It's equivalent to using arithmetic operator: ``**``: + + >>> df['B'] ** (df['A']) + 0 4 + 1 25 + 2 216 + dtype: Int64 + Args: other (float, int, or Series): Any single or multiple element data structure, or list-like object. @@ -1438,6 +1877,21 @@ def combine( to element-wise combine columns. The row and column indexes of the resulting DataFrame will be the union of the two. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df1 = bpd.DataFrame({'A': [0, 0], 'B': [4, 4]}) + >>> df2 = bpd.DataFrame({'A': [1, 1], 'B': [3, 3]}) + >>> take_smaller = lambda s1, s2: s1 if s1.sum() < s2.sum() else s2 + >>> df1.combine(df2, take_smaller) + A B + 0 0 3 + 1 0 3 + + [2 rows x 2 columns] + Args: other (DataFrame): The DataFrame to merge column-wise. @@ -1468,6 +1922,20 @@ def combine_first(self, other) -> DataFrame: second.loc[index, col] are not missing values, upon calling first.combine_first(second). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df1 = bpd.DataFrame({'A': [None, 0], 'B': [None, 4]}) + >>> df2 = bpd.DataFrame({'A': [1, 1], 'B': [3, 3]}) + >>> df1.combine_first(df2) + A B + 0 1.0 3.0 + 1 0.0 4.0 + + [2 rows x 2 columns] + Args: other (DataFrame): Provided DataFrame to use to fill null values. @@ -1485,6 +1953,24 @@ def update( Aligns on indices. There is no return value. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'A': [1, 2, 3], + ... 'B': [400, 500, 600]}) + >>> new_df = bpd.DataFrame({'B': [4, 5, 6], + ... 'C': [7, 8, 9]}) + >>> df.update(new_df) + >>> df + A B + 0 1 4 + 1 2 5 + 2 3 6 + + [3 rows x 2 columns] + Args: other (DataFrame, or object coercible into a DataFrame): Should have at least one matching index/column label From 79a638eda80c482b640b523426ffd95c42747edc Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Mon, 30 Oct 2023 18:56:14 +0000 Subject: [PATCH 122/197] feat: Implement operator `@` for `DataFrame.dot` (#139) 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 b/297502513 🦕 --- bigframes/dataframe.py | 2 ++ tests/system/small/test_dataframe.py | 33 +++++++++++++++++++++++++++ tests/system/small/test_multiindex.py | 16 +++++++++++++ 3 files changed, 51 insertions(+) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 49d7ad991a2..3369fb4868e 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2707,3 +2707,5 @@ def get_right_id(id): result = result[other.name].rename() return result + + __matmul__ = dot diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index b503f9a31d3..c96faa3526b 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3264,6 +3264,23 @@ def test_df_dot( ) +def test_df_dot_operator( + matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df +): + bf_result = (matrix_2by3_df @ matrix_3by4_df).to_pandas() + pd_result = matrix_2by3_pandas_df @ matrix_3by4_pandas_df + + # Patch pandas dtypes for testing parity + # Pandas result is object instead of Int64 (nullable) dtype. + for name in pd_result.columns: + pd_result[name] = pd_result[name].astype(pd.Int64Dtype()) + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + def test_df_dot_series( matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df ): @@ -3278,3 +3295,19 @@ def test_df_dot_series( bf_result, pd_result, ) + + +def test_df_dot_operator_series( + matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df +): + bf_result = (matrix_2by3_df @ matrix_3by4_df["x"]).to_pandas() + pd_result = matrix_2by3_pandas_df @ matrix_3by4_pandas_df["x"] + + # Patch pandas dtypes for testing parity + # Pandas result is object instead of Int64 (nullable) dtype. + pd_result = pd_result.astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index d6bf46f77c8..bc35f633fdb 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -998,6 +998,9 @@ def test_df_multi_index_dot_not_supported(): with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): bf1.dot(bf2) + with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): + bf1 @ bf2 + # right multi-index right_index = pandas.MultiIndex.from_tuples([("a", "aa"), ("a", "ab"), ("b", "bb")]) bf1 = bpd.DataFrame(left_matrix) @@ -1005,6 +1008,9 @@ def test_df_multi_index_dot_not_supported(): with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): bf1.dot(bf2) + with pytest.raises(NotImplementedError, match="Multi-index input is not supported"): + bf1 @ bf2 + def test_column_multi_index_dot_not_supported(): left_matrix = [[1, 2, 3], [2, 5, 7]] @@ -1022,6 +1028,11 @@ def test_column_multi_index_dot_not_supported(): ): bf1.dot(bf2) + with pytest.raises( + NotImplementedError, match="Multi-level column input is not supported" + ): + bf1 @ bf2 + # right multi-columns bf1 = bpd.DataFrame(left_matrix) bf2 = bpd.DataFrame(right_matrix, columns=multi_level_columns) @@ -1029,3 +1040,8 @@ def test_column_multi_index_dot_not_supported(): NotImplementedError, match="Multi-level column input is not supported" ): bf1.dot(bf2) + + with pytest.raises( + NotImplementedError, match="Multi-level column input is not supported" + ): + bf1 @ bf2 From cfebfaa91f945f0024ef743d38acb0b2ec8c4079 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 31 Oct 2023 10:54:17 -0700 Subject: [PATCH 123/197] test: add code snippets for loading data from BigQuery Job (#154) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: add code snippets for loading data from BigQuery Job * fix: address the comments * fix: fix the broken test * use BigQuery Client library to get the job_id * feat: Implement operator `@` for `DataFrame.dot` (#139) 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 b/297502513 🦕 * fix: fix the comments --------- Co-authored-by: Shobhit Singh --- .../load_data_from_biquery_job_test.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 samples/snippets/load_data_from_biquery_job_test.py diff --git a/samples/snippets/load_data_from_biquery_job_test.py b/samples/snippets/load_data_from_biquery_job_test.py new file mode 100644 index 00000000000..5271574a492 --- /dev/null +++ b/samples/snippets/load_data_from_biquery_job_test.py @@ -0,0 +1,51 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_bigquery_job(): + from google.cloud import bigquery + + # Construct a BigQuery client object. + client = bigquery.Client(project="bigframes-dev", location="us") + + query = """ + SELECT * + FROM `bigquery-public-data.ml_datasets.penguins` + LIMIT 20 + """ + query_job = client.query(query) + JOB_ID = query_job.job_id + your_project_id = "bigframes-dev" + + # [START bigquery_dataframes_load_data_from_bigquery_job] + from google.cloud import bigquery + + import bigframes.pandas as bpd + + # Project ID inserted based on the query results selected to explore + project = your_project_id + # Location inserted based on the query results selected to explore + location = "us" + client = bigquery.Client(project=project, location=location) + + # Job ID inserted based on the query results selcted to explore + job_id = JOB_ID + job = client.get_job(job_id) + destination = str(job.destination) + + # Load data from a BigQuery table using BigFrames DataFrames: + bq_df = bpd.read_gbq_table(destination) + + # [END bigquery_dataframes_load_data_from_bigquery_job] + assert bq_df is not None From 63c7919e28d2e0b864142320b47374d807f07c03 Mon Sep 17 00:00:00 2001 From: Bradford Orr <15842009+orrbradford@users.noreply.github.com> Date: Tue, 31 Oct 2023 12:46:47 -0700 Subject: [PATCH 124/197] feat: add bigframes.options.compute.maximum_bytes_billed option that sets maximum bytes billed on query jobs (#133) -implement context manager for global options -maximum_bytes_billed only applies to query jobs. This limitation will be set per query. Operations that trigger multiple jobs may result in total usage beyond this setting --- bigframes/__init__.py | 3 +- bigframes/_config/__init__.py | 11 +++++ bigframes/_config/compute_options.py | 35 +++++++++++++++ bigframes/_config/display_options.py | 23 ++++------ bigframes/pandas/__init__.py | 4 ++ bigframes/session/__init__.py | 19 +++++--- docs/reference/bigframes/options.rst | 2 + docs/templates/toc.yml | 2 + tests/system/conftest.py | 7 --- tests/system/small/test_progress_bar.py | 17 ++++--- tests/unit/test_compute_options.py | 30 +++++++++++++ .../pandas/_config/config.py | 45 +++++++++++++++++++ 12 files changed, 162 insertions(+), 36 deletions(-) create mode 100644 bigframes/_config/compute_options.py create mode 100644 tests/unit/test_compute_options.py create mode 100644 third_party/bigframes_vendored/pandas/_config/config.py diff --git a/bigframes/__init__.py b/bigframes/__init__.py index 8f41790072b..bd1476957bf 100644 --- a/bigframes/__init__.py +++ b/bigframes/__init__.py @@ -14,7 +14,7 @@ """BigQuery DataFrames provides a DataFrame API scaled by the BigQuery engine.""" -from bigframes._config import options +from bigframes._config import option_context, options from bigframes._config.bigquery_options import BigQueryOptions from bigframes.core.global_session import close_session, get_global_session from bigframes.session import connect, Session @@ -28,4 +28,5 @@ "connect", "Session", "__version__", + "option_context", ] diff --git a/bigframes/_config/__init__.py b/bigframes/_config/__init__.py index e26eaf88009..8dcebfce6a2 100644 --- a/bigframes/_config/__init__.py +++ b/bigframes/_config/__init__.py @@ -18,8 +18,10 @@ """ import bigframes._config.bigquery_options as bigquery_options +import bigframes._config.compute_options as compute_options import bigframes._config.display_options as display_options import bigframes._config.sampling_options as sampling_options +import third_party.bigframes_vendored.pandas._config.config as pandas_config class Options: @@ -29,6 +31,7 @@ def __init__(self): self._bigquery_options = bigquery_options.BigQueryOptions() self._display_options = display_options.DisplayOptions() self._sampling_options = sampling_options.SamplingOptions() + self._compute_options = compute_options.ComputeOptions() @property def bigquery(self) -> bigquery_options.BigQueryOptions: @@ -49,6 +52,11 @@ def sampling(self) -> sampling_options.SamplingOptions: parameters in specific functions.""" return self._sampling_options + @property + def compute(self) -> compute_options.ComputeOptions: + """Options controlling object computation.""" + return self._compute_options + options = Options() """Global options for default session.""" @@ -58,3 +66,6 @@ def sampling(self) -> sampling_options.SamplingOptions: "Options", "options", ) + + +option_context = pandas_config.option_context diff --git a/bigframes/_config/compute_options.py b/bigframes/_config/compute_options.py new file mode 100644 index 00000000000..20c31d39066 --- /dev/null +++ b/bigframes/_config/compute_options.py @@ -0,0 +1,35 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Options for displaying objects.""" + +import dataclasses +from typing import Optional + + +@dataclasses.dataclass +class ComputeOptions: + """ + Encapsulates configuration for compute options. + + Attributes: + maximum_bytes_billed (int, Options): + Limits the bytes billed for query jobs. Queries that will have + bytes billed beyond this limit will fail (without incurring a + charge). If unspecified, this will be set to your project default. + See `maximum_bytes_billed `_. + + """ + + maximum_bytes_billed: Optional[int] = None diff --git a/bigframes/_config/display_options.py b/bigframes/_config/display_options.py index 8bd2743f175..ad3ea3f68cc 100644 --- a/bigframes/_config/display_options.py +++ b/bigframes/_config/display_options.py @@ -40,17 +40,12 @@ def pandas_repr(display_options: DisplayOptions): 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 + with pd.option_context( + "display.max_columns", + display_options.max_columns, + "display.max_rows", + display_options.max_rows, + "display.show_dimensions", + True, + ) as pandas_context: + yield (pandas_context) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 8d9726312fd..0fab1109dc7 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -462,6 +462,9 @@ def read_gbq_function(function_name: str): options = config.options """Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames.""" +option_context = config.option_context +"""Global :class:`~bigframes._config.option_context` to configure BigQuery DataFrames.""" + # Session management APIs get_global_session = global_session.get_global_session close_session = global_session.close_session @@ -494,6 +497,7 @@ def read_gbq_function(function_name: str): # Other public pandas attributes "NamedAgg", "options", + "option_context", # Session management APIs "get_global_session", "close_session", diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 932a41f2830..4858c7726a7 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1504,12 +1504,10 @@ def _start_query( max_results: Optional[int] = None, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """ - Starts query job and waits for results + 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) + job_config = self._prepare_job_config(job_config) + query_job = self.bqclient.query(sql, job_config=job_config) opts = bigframes.options.display if opts.progress_bar is not None and not query_job.configuration.dry_run: @@ -1538,6 +1536,17 @@ def _start_generic_job(self, job: formatting_helpers.GenericJob): else: job.result() + def _prepare_job_config( + self, job_config: Optional[bigquery.QueryJobConfig] = None + ) -> bigquery.QueryJobConfig: + if job_config is None: + job_config = self.bqclient.default_query_job_config + if bigframes.options.compute.maximum_bytes_billed is not None: + job_config.maximum_bytes_billed = ( + bigframes.options.compute.maximum_bytes_billed + ) + return job_config + def connect(context: Optional[bigquery_options.BigQueryOptions] = None) -> Session: return Session(context) diff --git a/docs/reference/bigframes/options.rst b/docs/reference/bigframes/options.rst index d831a519fe6..991399eb886 100644 --- a/docs/reference/bigframes/options.rst +++ b/docs/reference/bigframes/options.rst @@ -12,3 +12,5 @@ Options and settings .. autoclass:: bigframes._config.display_options.DisplayOptions .. autoclass:: bigframes._config.sampling_options.SamplingOptions + +.. autoclass:: bigframes._config.compute_options.ComputeOptions diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 4fe2ec1a6ad..9879721d286 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -13,6 +13,8 @@ uid: bigframes._config.display_options.DisplayOptions - name: SamplingOptions uid: bigframes._config.sampling_options.SamplingOptions + - name: ComputeOptions + uid: bigframes._config.compute_options.ComputeOptions name: Options and settings - items: - name: Session diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 8885b03d34c..f9f69c6c8ee 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -898,13 +898,6 @@ def usa_names_grouped_table( return session.bqclient.get_table(table_id) -@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 diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index 084b723fba6..30ea63b483a 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -135,12 +135,11 @@ def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame): 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 +def test_query_job_dry_run(penguins_df_default_index: bf.dataframe.DataFrame, capsys): + with bf.option_context("display.repr_mode", "deferred"): + 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/unit/test_compute_options.py b/tests/unit/test_compute_options.py new file mode 100644 index 00000000000..499a0a5fefa --- /dev/null +++ b/tests/unit/test_compute_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. +import bigframes as bf + +from . import resources + + +def test_maximum_bytes_option(): + session = resources.create_bigquery_session() + num_query_calls = 0 + with bf.option_context("compute.maximum_bytes_billed", 10000): + # clear initial method calls + session.bqclient.method_calls = [] + session._start_query("query") + for call in session.bqclient.method_calls: + _, _, kwargs = call + num_query_calls += 1 + assert kwargs["job_config"].maximum_bytes_billed == 10000 + assert num_query_calls > 0 diff --git a/third_party/bigframes_vendored/pandas/_config/config.py b/third_party/bigframes_vendored/pandas/_config/config.py new file mode 100644 index 00000000000..8abaca76c78 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/_config/config.py @@ -0,0 +1,45 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/_config/config.py +import contextlib +import operator + +import bigframes + + +class option_context(contextlib.ContextDecorator): + """ + Context manager to temporarily set options in the `with` statement context. + + You need to invoke as ``option_context(pat, val, [(pat, val), ...])``. + + Examples + -------- + >>> import bigframes + >>> with bigframes.option_context('display.max_rows', 10, 'display.max_columns', 5): + ... pass + """ + + def __init__(self, *args) -> None: + if len(args) % 2 != 0 or len(args) < 2: + raise ValueError( + "Need to invoke as option_context(pat, val, [(pat, val), ...])." + ) + + self.ops = list(zip(args[::2], args[1::2])) + + def __enter__(self) -> None: + self.undo = [ + (pat, operator.attrgetter(pat)(bigframes.options)) for pat, val in self.ops + ] + + for pat, val in self.ops: + self._set_option(pat, val) + + def __exit__(self, *args) -> None: + if self.undo: + for pat, val in self.undo: + self._set_option(pat, val) + + def _set_option(self, pat, val): + root, attr = pat.rsplit(".", 1) + parent = operator.attrgetter(root)(bigframes.options) + setattr(parent, attr, val) From 0801d96830dab467232277dea9fd2dacee41055c Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 31 Oct 2023 18:00:20 -0500 Subject: [PATCH 125/197] docs: fix indentation on `read_gbq_function` code sample (#163) 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/session/__init__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 4858c7726a7..5a61ed534f3 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1467,13 +1467,13 @@ def read_gbq_function( **Examples:** - >>> import bigframes.pandas as bpd - >>> bpd.options.display.progress_bar = None + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None - >>> function_name = "bqutil.fn.cw_lower_case_ascii_only" - >>> func = bpd.read_gbq_function(function_name=function_name) - >>> func.bigframes_remote_function - 'bqutil.fn.cw_lower_case_ascii_only' + >>> function_name = "bqutil.fn.cw_lower_case_ascii_only" + >>> func = bpd.read_gbq_function(function_name=function_name) + >>> func.bigframes_remote_function + 'bqutil.fn.cw_lower_case_ascii_only' Args: function_name (str): From d8baad5b71ec67a35a0fb6132ee16e4c7418c456 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Wed, 1 Nov 2023 13:41:27 -0700 Subject: [PATCH 126/197] feat: add pd.get_dummies (#149) * feat: add pd.get_dummies * remove unneeded prefix case * param/documentation fixes * be stricter about types in test * be stricter about types in series test * remove unneeded comment * adjust for type difference in pandas 1 * add example code (tested) * fix None columns and add test cases * variable names and _get_unique_values per-column * account for pandas 1 behavior difference * remove already_seen set * avoid unnecessary join/projection * fix column ordering edge case * adjust for picky examples checker * example tweak * make part of the example comments * use ellipsis in doctest comment * add to doctest string * extract parameter standardization * extract submethods --------- Co-authored-by: Henry J Solberg --- bigframes/pandas/__init__.py | 177 ++++++++++++++++++ tests/system/small/test_pandas.py | 112 +++++++++++ .../pandas/core/reshape/concat.py | 2 +- .../pandas/core/reshape/encoding.py | 119 ++++++++++++ .../pandas/core/reshape/merge.py | 1 - .../pandas/core/reshape/tile.py | 2 +- 6 files changed, 410 insertions(+), 3 deletions(-) create mode 100644 third_party/bigframes_vendored/pandas/core/reshape/encoding.py diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 0fab1109dc7..1c52b103fbe 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -45,14 +45,18 @@ ) import bigframes._config as config +import bigframes.constants as constants +import bigframes.core.blocks import bigframes.core.global_session as global_session import bigframes.core.indexes import bigframes.core.reshape import bigframes.dataframe +import bigframes.operations as ops import bigframes.series import bigframes.session import bigframes.session.clients import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat +import third_party.bigframes_vendored.pandas.core.reshape.encoding as vendored_pandas_encoding 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 @@ -134,6 +138,179 @@ def cut( cut.__doc__ = vendored_pandas_tile.cut.__doc__ +def get_dummies( + data: Union[DataFrame, Series], + prefix: Union[List, dict, str, None] = None, + prefix_sep: Union[List, dict, str, None] = "_", + dummy_na: bool = False, + columns: Optional[List] = None, + drop_first: bool = False, + dtype: Any = None, +) -> DataFrame: + # simplify input parameters into per-input-label lists + # also raise errors for invalid parameters + column_labels, prefixes, prefix_seps = _standardize_get_dummies_params( + data, prefix, prefix_sep, columns, dtype + ) + + # combine prefixes into per-column-id list + full_columns_prefixes, columns_ids = _determine_get_dummies_columns_from_labels( + data, column_labels, prefix is not None, prefixes, prefix_seps + ) + + # run queries to compute unique values + block = data._block + max_unique_value = ( + bigframes.core.blocks._BQ_MAX_COLUMNS + - len(block.value_columns) + - len(block.index_columns) + - 1 + ) // len(column_labels) + columns_values = [ + block._get_unique_values([col_id], max_unique_value) for col_id in columns_ids + ] + + # for each dummified column, add the content of the output columns via block operations + intermediate_col_ids = [] + for i in range(len(columns_values)): + level = columns_values[i].get_level_values(0).sort_values().dropna() + if drop_first: + level = level[1:] + column_label = full_columns_prefixes[i] + column_id = columns_ids[i] + block, new_intermediate_col_ids = _perform_get_dummies_block_operations( + block, level, column_label, column_id, dummy_na + ) + intermediate_col_ids.extend(new_intermediate_col_ids) + + # drop dummified columns (and the intermediate columns we added) + block = block.drop_columns(columns_ids + intermediate_col_ids) + return DataFrame(block) + + +get_dummies.__doc__ = vendored_pandas_encoding.get_dummies.__doc__ + + +def _standardize_get_dummies_params( + data: Union[DataFrame, Series], + prefix: Union[List, dict, str, None], + prefix_sep: Union[List, dict, str, None], + columns: Optional[List], + dtype: Any, +) -> Tuple[List, List[str], List[str]]: + block = data._block + + if isinstance(data, Series): + columns = [block.column_labels[0]] + if columns is not None and not pandas.api.types.is_list_like(columns): + raise TypeError("Input must be a list-like for parameter `columns`") + if dtype is not None and dtype not in [ + pandas.BooleanDtype, + bool, + "Boolean", + "boolean", + "bool", + ]: + raise NotImplementedError( + f"Only Boolean dtype is currently supported. {constants.FEEDBACK_LINK}" + ) + + if columns is None: + default_dummy_types = [pandas.StringDtype, "string[pyarrow]"] + columns = [] + columns_set = set() + for col_id in block.value_columns: + label = block.col_id_to_label[col_id] + if ( + label not in columns_set + and block.expr.get_column_type(col_id) in default_dummy_types + ): + columns.append(label) + columns_set.add(label) + + column_labels: List = typing.cast(List, columns) + + def parse_prefix_kwarg(kwarg, kwarg_name) -> Optional[List[str]]: + if kwarg is None: + return None + if isinstance(kwarg, str): + return [kwarg] * len(column_labels) + if isinstance(kwarg, dict): + return [kwarg[column] for column in column_labels] + kwarg = typing.cast(List, kwarg) + if pandas.api.types.is_list_like(kwarg) and len(kwarg) != len(column_labels): + raise ValueError( + f"Length of '{kwarg_name}' ({len(kwarg)}) did not match " + f"the length of the columns being encoded ({len(column_labels)})." + ) + if pandas.api.types.is_list_like(kwarg): + return list(map(str, kwarg)) + raise TypeError(f"{kwarg_name} kwarg must be a string, list, or dictionary") + + prefix_seps = parse_prefix_kwarg(prefix_sep or "_", "prefix_sep") + prefix_seps = typing.cast(List, prefix_seps) + prefixes = parse_prefix_kwarg(prefix, "prefix") + if prefixes is None: + prefixes = column_labels + prefixes = typing.cast(List, prefixes) + + return column_labels, prefixes, prefix_seps + + +def _determine_get_dummies_columns_from_labels( + data: Union[DataFrame, Series], + column_labels: List, + prefix_given: bool, + prefixes: List[str], + prefix_seps: List[str], +) -> Tuple[List[str], List[str]]: + block = data._block + + columns_ids = [] + columns_prefixes = [] + for i in range(len(column_labels)): + label = column_labels[i] + empty_prefix = label is None or (isinstance(data, Series) and not prefix_given) + full_prefix = "" if empty_prefix else prefixes[i] + prefix_seps[i] + + for col_id in block.label_to_col_id[label]: + columns_ids.append(col_id) + columns_prefixes.append(full_prefix) + + return columns_prefixes, columns_ids + + +def _perform_get_dummies_block_operations( + block: bigframes.core.blocks.Block, + level: pandas.Index, + column_label: str, + column_id: str, + dummy_na: bool, +) -> Tuple[bigframes.core.blocks.Block, List[str]]: + intermediate_col_ids = [] + for value in level: + new_column_label = f"{column_label}{value}" + if column_label == "": + new_column_label = value + new_block, new_id = block.apply_unary_op( + column_id, ops.BinopPartialLeft(ops.eq_op, value) + ) + intermediate_col_ids.append(new_id) + block, _ = new_block.apply_unary_op( + new_id, + ops.BinopPartialRight(ops.fillna_op, False), + result_label=new_column_label, + ) + if dummy_na: + # dummy column name for na depends on the dtype + na_string = str(pandas.Index([None], dtype=level.dtype)[0]) + new_column_label = f"{column_label}{na_string}" + block, _ = block.apply_unary_op( + column_id, ops.isnull_op, result_label=new_column_label + ) + return block, intermediate_col_ids + + def qcut( x: bigframes.series.Series, q: int, diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index f8fa78587f9..0292ebd2069 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -45,6 +45,118 @@ def test_concat_series(scalars_dfs): pd.testing.assert_series_equal(bf_result, pd_result) +@pytest.mark.parametrize( + ("kwargs"), + [ + { + "prefix": ["prefix1", "prefix2"], + "prefix_sep": "_", + "dummy_na": None, + "columns": ["bool_col", "int64_col"], + "drop_first": False, + }, + { + "prefix": "prefix", + "prefix_sep": ["_", ","], + "dummy_na": False, + "columns": ["int64_too", "string_col"], + "drop_first": False, + }, + { + "prefix": None, + "prefix_sep": ".", + "dummy_na": True, + "columns": ["time_col", "float64_col"], + "drop_first": True, + }, + ], +) +def test_get_dummies_dataframe(scalars_dfs, kwargs): + scalars_df, scalars_pandas_df = scalars_dfs + + bf_result = bpd.get_dummies(scalars_df, **kwargs, dtype=bool) + pd_result = pd.get_dummies(scalars_pandas_df, **kwargs, dtype=bool) + # dtype argument above is needed for pandas v1 only + + # adjust for expected dtype differences + for (column_name, type_name) in zip(pd_result.columns, pd_result.dtypes): + if type_name == "bool": + pd_result[column_name] = pd_result[column_name].astype("boolean") + + pd.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) + + +def test_get_dummies_dataframe_duplicate_labels(scalars_dfs): + if pd.__version__.startswith("1."): + pytest.skip("pandas has different behavior in 1.x") + + scalars_df, scalars_pandas_df = scalars_dfs + + scalars_renamed_df = scalars_df.rename( + columns={"int64_too": "int64_col", "float64_col": None, "string_col": None} + ) + scalars_renamed_pandas_df = scalars_pandas_df.rename( + columns={"int64_too": "int64_col", "float64_col": None, "string_col": None} + ) + + bf_result = bpd.get_dummies( + scalars_renamed_df, columns=["int64_col", None], dtype=bool + ) + pd_result = pd.get_dummies( + scalars_renamed_pandas_df, columns=["int64_col", None], dtype=bool + ) + # dtype argument above is needed for pandas v1 only + + # adjust for expected dtype differences + for (column_name, type_name) in zip(pd_result.columns, pd_result.dtypes): + if type_name == "bool": + pd_result[column_name] = pd_result[column_name].astype("boolean") + + pd.testing.assert_frame_equal(bf_result.to_pandas(), pd_result) + + +def test_get_dummies_series(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_series = scalars_df.date_col + pd_series = scalars_pandas_df.date_col + + bf_result = bpd.get_dummies(bf_series, dtype=bool) + pd_result = pd.get_dummies(pd_series, dtype=bool) + # dtype argument above is needed for pandas v1 only + + # adjust for expected dtype differences + for (column_name, type_name) in zip(pd_result.columns, pd_result.dtypes): + if type_name == "bool": + pd_result[column_name] = pd_result[column_name].astype("boolean") + pd_result.columns = pd_result.columns.astype(object) + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + +def test_get_dummies_series_nameless(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_series = scalars_df.date_col.rename(None) + pd_series = scalars_pandas_df.date_col.rename(None) + + bf_result = bpd.get_dummies(bf_series, dtype=bool) + pd_result = pd.get_dummies(pd_series, dtype=bool) + # dtype argument above is needed for pandas v1 only + + # adjust for expected dtype differences + for (column_name, type_name) in zip(pd_result.columns, pd_result.dtypes): + if type_name == "bool": + pd_result[column_name] = pd_result[column_name].astype("boolean") + pd_result.columns = pd_result.columns.astype(object) + + pd.testing.assert_frame_equal( + bf_result.to_pandas(), + pd_result, + ) + + @pytest.mark.parametrize( ("how"), [ diff --git a/third_party/bigframes_vendored/pandas/core/reshape/concat.py b/third_party/bigframes_vendored/pandas/core/reshape/concat.py index 6e6d2d8b5c9..b0472c524ab 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/concat.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/concat.py @@ -1,6 +1,6 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/concat.py """ -Concat routines. +Concat routines """ from __future__ import annotations diff --git a/third_party/bigframes_vendored/pandas/core/reshape/encoding.py b/third_party/bigframes_vendored/pandas/core/reshape/encoding.py new file mode 100644 index 00000000000..da92b58f505 --- /dev/null +++ b/third_party/bigframes_vendored/pandas/core/reshape/encoding.py @@ -0,0 +1,119 @@ +# Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/encoding.py +""" +Encoding routines +""" +from __future__ import annotations + +from bigframes import constants + + +def get_dummies( + data, + prefix=None, + prefix_sep="_", + dummy_na=False, + columns=None, + drop_first=False, + dtype=None, +): + """ + Convert categorical variable into dummy/indicator variables. + + Each variable is converted in as many 0/1 variables as there are + different values. Columns in the output are each named after a value; + if the input is a DataFrame, the name of the original variable is + prepended to the value. + + **Examples:** + >>> import bigframes.pandas as pd + >>> pd.options.display.progress_bar = None + >>> s = pd.Series(list('abca')) + >>> pd.get_dummies(s) + a b c + 0 True False False + 1 False True False + 2 False False True + 3 True False False + + [4 rows x 3 columns] + + >>> s1 = pd.Series(['a', 'b', None]) + >>> pd.get_dummies(s1) + a b + 0 True False + 1 False True + 2 False False + + [3 rows x 2 columns] + + >>> pd.get_dummies(s1, dummy_na=True) + a b + 0 True False False + 1 False True False + 2 False False True + + [3 rows x 3 columns] + + >>> df = pd.DataFrame({'A': ['a', 'b', 'a'], 'B': ['b', 'a', 'c'], 'C': [1, 2, 3]}) + >>> pd.get_dummies(df, prefix=['col1', 'col2']) + C col1_a col1_b col2_a col2_b col2_c + 0 1 True False False True False + 1 2 False True True False False + 2 3 True False False False True + + [3 rows x 6 columns] + + >>> pd.get_dummies(pd.Series(list('abcaa'))) + a b c + 0 True False False + 1 False True False + 2 False False True + 3 True False False + 4 True False False + + [5 rows x 3 columns] + + >>> pd.get_dummies(pd.Series(list('abcaa')), drop_first=True) + b c + 0 False False + 1 True False + 2 False True + 3 False False + 4 False False + + [5 rows x 2 columns] + + Args: + data (Series or DataFrame): + Data of which to get dummy indicators. + + prefix (str, list of str, or dict of str, default None): + String to append DataFrame column names. Pass a list with length + equal to the number of columns when calling get_dummies on a + DataFrame. Alternatively, prefix can be a dictionary mapping column + names to prefixes. + + prefix_sep (str, list of str, or dict of str, default '_'): + Separator/delimiter to use, appended to prefix. Or pass a list or + dictionary as with prefix. + + dummy_na (bool, default False): + Add a column to indicate NaNs, if False NaNs are ignored. + + columns (list-like, default None): + Column names in the DataFrame to be encoded. If columns is None + then only the columns with string dtype will be converted. + + drop_first (bool, default False): + Whether to get k-1 dummies out of k categorical levels by removing the + first level. + + dtype (dtype, default bool): + Data type for new columns. Only a single dtype is allowed. + + Returns: + DataFrame: Dummy-coded data. If data contains other columns than the + dummy-coded one(s), these will be prepended, unaltered, to the + result. + """ + 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 index cc81de405b9..b03f366fca7 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/merge.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/merge.py @@ -16,7 +16,6 @@ def merge( sort=False, suffixes=("_x", "_y"), ): - """ Merge DataFrame objects with a database-style join. diff --git a/third_party/bigframes_vendored/pandas/core/reshape/tile.py b/third_party/bigframes_vendored/pandas/core/reshape/tile.py index 24ea655a5fe..d4471ed68e4 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/tile.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/tile.py @@ -1,6 +1,6 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/reshape/tile.py """ -Quantilization functions and related stuff +Quantilization functions and related routines """ from __future__ import annotations From bf1ec89f8da2c7b2d042b7516a16a9e7cda6db06 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 1 Nov 2023 14:49:24 -0700 Subject: [PATCH 127/197] chore(main): release 0.12.0 (#151) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 29 +++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 93ebadb56fb..845d3634bc9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,35 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.12.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.11.0...v0.12.0) (2023-11-01) + + +### Features + +* Add `DataFrame.melt` ([#113](https://github.com/googleapis/python-bigquery-dataframes/issues/113)) ([4e4409c](https://github.com/googleapis/python-bigquery-dataframes/commit/4e4409c5b235171f3770aec852193026519948fd)) +* Add `DataFrame.to_pandas_batches()` to download large `DataFrame` objects ([#136](https://github.com/googleapis/python-bigquery-dataframes/issues/136)) ([3afd4a3](https://github.com/googleapis/python-bigquery-dataframes/commit/3afd4a35f4c38dad86dab17ff62444cd418cab88)) +* Add bigframes.options.compute.maximum_bytes_billed option that sets maximum bytes billed on query jobs ([#133](https://github.com/googleapis/python-bigquery-dataframes/issues/133)) ([63c7919](https://github.com/googleapis/python-bigquery-dataframes/commit/63c7919e28d2e0b864142320b47374d807f07c03)) +* Add pandas.qcut ([#104](https://github.com/googleapis/python-bigquery-dataframes/issues/104)) ([8e44518](https://github.com/googleapis/python-bigquery-dataframes/commit/8e4451841ba09099b0ed5433f9102511741dfbed)) +* Add pd.get_dummies ([#149](https://github.com/googleapis/python-bigquery-dataframes/issues/149)) ([d8baad5](https://github.com/googleapis/python-bigquery-dataframes/commit/d8baad5b71ec67a35a0fb6132ee16e4c7418c456)) +* Add unstack to series, add level param ([#115](https://github.com/googleapis/python-bigquery-dataframes/issues/115)) ([5edcd19](https://github.com/googleapis/python-bigquery-dataframes/commit/5edcd19e6200db9b9ebe3d4945816b3ebf1f7bcd)) +* Implement operator `@` for `DataFrame.dot` ([#139](https://github.com/googleapis/python-bigquery-dataframes/issues/139)) ([79a638e](https://github.com/googleapis/python-bigquery-dataframes/commit/79a638eda80c482b640b523426ffd95c42747edc)) +* Populate ibis version in user agent ([#140](https://github.com/googleapis/python-bigquery-dataframes/issues/140)) ([c639a36](https://github.com/googleapis/python-bigquery-dataframes/commit/c639a3657465e2b68a3b93c363bd3ae1e969d2cc)) + + +### Bug Fixes + +* Don't override the global logging config ([#138](https://github.com/googleapis/python-bigquery-dataframes/issues/138)) ([2ddbf74](https://github.com/googleapis/python-bigquery-dataframes/commit/2ddbf743efc2fd8ffb61ae8d3333fc4b98ce4b55)) +* Fix bug with column names under repeated column assignment ([#150](https://github.com/googleapis/python-bigquery-dataframes/issues/150)) ([29032d0](https://github.com/googleapis/python-bigquery-dataframes/commit/29032d06811569121f7be2a7de915740df7daf6e)) +* Resolve plotly rendering issue by using ipython html for job pro… ([#134](https://github.com/googleapis/python-bigquery-dataframes/issues/134)) ([39df43e](https://github.com/googleapis/python-bigquery-dataframes/commit/39df43e243ac0374d1a1eb2a75779324825afbe9)) +* Use indexee's session for loc listlike cases ([#152](https://github.com/googleapis/python-bigquery-dataframes/issues/152)) ([27c5725](https://github.com/googleapis/python-bigquery-dataframes/commit/27c57255c7fe11e1ef9b9826d988d80fc17442a6)) + + +### Documentation + +* Add artithmetic df sample code ([#153](https://github.com/googleapis/python-bigquery-dataframes/issues/153)) ([ac44ccd](https://github.com/googleapis/python-bigquery-dataframes/commit/ac44ccd3936cdb28755d2bbe16377d489f08d5e5)) +* Fix indentation on `read_gbq_function` code sample ([#163](https://github.com/googleapis/python-bigquery-dataframes/issues/163)) ([0801d96](https://github.com/googleapis/python-bigquery-dataframes/commit/0801d96830dab467232277dea9fd2dacee41055c)) +* Link to ML.EVALUATE BQML page for score() methods ([#137](https://github.com/googleapis/python-bigquery-dataframes/issues/137)) ([45c617f](https://github.com/googleapis/python-bigquery-dataframes/commit/45c617fee7becc42f1c129246ffdc32f3a963f12)) + ## [0.11.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.10.0...v0.11.0) (2023-10-26) diff --git a/bigframes/version.py b/bigframes/version.py index 18edfa56150..b324ed72340 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.11.0" +__version__ = "0.12.0" From 3d7a0d6f7172b8b38ac1825990e612ce03645743 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Wed, 1 Nov 2023 15:45:08 -0700 Subject: [PATCH 128/197] test: add code snippets for using bigframes.ml (#159) * test: add code snippets for using bigframes.ml --- samples/snippets/clustering_model_test.py | 35 ++++++++++++++ samples/snippets/gen_ai_model_test.py | 39 ++++++++++++++++ samples/snippets/regression_model_test.py | 57 +++++++++++++++++++++++ 3 files changed, 131 insertions(+) create mode 100644 samples/snippets/clustering_model_test.py create mode 100644 samples/snippets/gen_ai_model_test.py create mode 100644 samples/snippets/regression_model_test.py diff --git a/samples/snippets/clustering_model_test.py b/samples/snippets/clustering_model_test.py new file mode 100644 index 00000000000..a407fc78058 --- /dev/null +++ b/samples/snippets/clustering_model_test.py @@ -0,0 +1,35 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_clustering_model(): + # [START bigquery_dataframes_clustering_model] + from bigframes.ml.cluster import KMeans + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Create the KMeans model + cluster_model = KMeans(n_clusters=10) + cluster_model.fit(bq_df["culmen_length_mm"], bq_df["sex"]) + + # Predict using the model + result = cluster_model.predict(bq_df) + # Score the model + score = cluster_model.score(bq_df) + # [END bigquery_dataframes_clustering_model] + assert result is not None + assert score is not None diff --git a/samples/snippets/gen_ai_model_test.py b/samples/snippets/gen_ai_model_test.py new file mode 100644 index 00000000000..7cbc90d4c03 --- /dev/null +++ b/samples/snippets/gen_ai_model_test.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. + + +def test_llm_model(): + PROJECT_ID = "bigframes-dev" + REGION = "us" + CONN_NAME = "bigframes-ml" + # [START bigquery_dataframes_gen_ai_model] + from bigframes.ml.llm import PaLM2TextGenerator + import bigframes.pandas as bpd + + # Create the LLM model + session = bpd.get_global_session() + connection = f"{PROJECT_ID}.{REGION}.{CONN_NAME}" + model = PaLM2TextGenerator(session=session, connection_name=connection) + + df_api = bpd.read_csv("gs://cloud-samples-data/vertex-ai/bigframe/df.csv") + + # Prepare the prompts and send them to the LLM model for prediction + df_prompt_prefix = "Generate Pandas sample code for DataFrame." + df_prompt = df_prompt_prefix + df_api["API"] + + # Predict using the model + df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024) + # [END bigquery_dataframes_gen_ai_model] + assert df_pred["ml_generate_text_llm_result"] is not None + assert df_pred["ml_generate_text_llm_result"].iloc[0] is not None diff --git a/samples/snippets/regression_model_test.py b/samples/snippets/regression_model_test.py new file mode 100644 index 00000000000..7d1bde689cb --- /dev/null +++ b/samples/snippets/regression_model_test.py @@ -0,0 +1,57 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_regression_model(): + # [START bigquery_dataframes_regression_model] + from bigframes.ml.linear_model import LinearRegression + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Filter down to the data to the Adelie Penguin species + adelie_data = bq_df[bq_df.species == "Adelie Penguin (Pygoscelis adeliae)"] + + # Drop the species column + adelie_data = adelie_data.drop(columns=["species"]) + + # Drop rows with nulls to get training data + training_data = adelie_data.dropna() + + # Specify your feature (or input) columns and the label (or output) column: + feature_columns = training_data[ + ["island", "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", "sex"] + ] + label_columns = training_data[["body_mass_g"]] + + test_data = adelie_data[adelie_data.body_mass_g.isnull()] + + # Create the linear model + model = LinearRegression() + model.fit(feature_columns, label_columns) + + # Score the model + score = model.score(feature_columns, label_columns) + + # Predict using the model + result = model.predict(test_data) + # [END bigquery_dataframes_regression_model] + assert test_data is not None + assert feature_columns is not None + assert label_columns is not None + assert model is not None + assert score is not None + assert result is not None From b9cb55c5b9354f9ff60de0aad66fe60049876055 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 1 Nov 2023 17:44:12 -0700 Subject: [PATCH 129/197] feat: add interpolate() to series and dataframe (#157) --- bigframes/core/block_transforms.py | 91 +++++++++++++++++++ bigframes/dataframe.py | 4 + bigframes/series.py | 4 + tests/system/small/test_dataframe.py | 16 ++++ tests/system/small/test_series.py | 26 ++++++ .../bigframes_vendored/pandas/core/frame.py | 37 ++++++++ .../bigframes_vendored/pandas/core/series.py | 32 +++++++ 7 files changed, 210 insertions(+) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 3706bf1681b..917edac0ded 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -105,6 +105,97 @@ def indicate_duplicates( ) +def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block: + if method != "linear": + raise NotImplementedError( + f"Only 'linear' interpolate method supported. {constants.FEEDBACK_LINK}" + ) + backwards_window = windows.WindowSpec(following=0) + forwards_window = windows.WindowSpec(preceding=0) + + output_column_ids = [] + + original_columns = block.value_columns + original_labels = block.column_labels + block, offsets = block.promote_offsets() + for column in original_columns: + # null in same places column is null + should_interpolate = block._column_type(column) in [ + pd.Float64Dtype(), + pd.Int64Dtype(), + ] + if should_interpolate: + block, notnull = block.apply_unary_op(column, ops.notnull_op) + block, masked_offsets = block.apply_binary_op( + offsets, notnull, ops.partial_arg3(ops.where_op, None) + ) + + block, previous_value = block.apply_window_op( + column, agg_ops.LastNonNullOp(), backwards_window + ) + block, next_value = block.apply_window_op( + column, agg_ops.FirstNonNullOp(), forwards_window + ) + block, previous_value_offset = block.apply_window_op( + masked_offsets, + agg_ops.LastNonNullOp(), + backwards_window, + skip_reproject_unsafe=True, + ) + block, next_value_offset = block.apply_window_op( + masked_offsets, + agg_ops.FirstNonNullOp(), + forwards_window, + skip_reproject_unsafe=True, + ) + + block, prediction_id = _interpolate( + block, + previous_value_offset, + previous_value, + next_value_offset, + next_value, + offsets, + ) + + block, interpolated_column = block.apply_binary_op( + column, prediction_id, ops.fillna_op + ) + # Pandas performs ffill-like behavior to extrapolate forwards + block, interpolated_and_ffilled = block.apply_binary_op( + interpolated_column, previous_value, ops.fillna_op + ) + + output_column_ids.append(interpolated_and_ffilled) + else: + output_column_ids.append(column) + + # Force reproject since used `skip_project_unsafe` perviously + block = block.select_columns(output_column_ids)._force_reproject() + return block.with_column_labels(original_labels) + + +def _interpolate( + block: blocks.Block, + x0_id: str, + y0_id: str, + x1_id: str, + y1_id: str, + xpredict_id: str, +) -> typing.Tuple[blocks.Block, str]: + """Applies linear interpolation equation to predict y values for xpredict.""" + block, x1x0diff = block.apply_binary_op(x1_id, x0_id, ops.sub_op) + block, y1y0diff = block.apply_binary_op(y1_id, y0_id, ops.sub_op) + block, xpredictx0diff = block.apply_binary_op(xpredict_id, x0_id, ops.sub_op) + + block, y1_weight = block.apply_binary_op(y1y0diff, x1x0diff, ops.div_op) + block, y1_part = block.apply_binary_op(xpredictx0diff, y1_weight, ops.mul_op) + + block, prediction_id = block.apply_binary_op(y0_id, y1_part, ops.add_op) + block = block.drop_columns([x1x0diff, y1y0diff, xpredictx0diff, y1_weight, y1_part]) + return block, prediction_id + + def drop_duplicates( block: blocks.Block, columns: typing.Sequence[str], keep: str = "first" ) -> blocks.Block: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 3369fb4868e..ffcaf0d613d 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1434,6 +1434,10 @@ def _reindex_columns(self, columns): def reindex_like(self, other: DataFrame, *, validate: typing.Optional[bool] = None): return self.reindex(index=other.index, columns=other.columns, validate=validate) + def interpolate(self, method: str = "linear") -> DataFrame: + result = block_ops.interpolate(self._block, method) + return DataFrame(result) + def fillna(self, value=None) -> DataFrame: return self._apply_binop(value, ops.fillna_op, how="left") diff --git a/bigframes/series.py b/bigframes/series.py index 37d00d16f32..824757cf52b 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -468,6 +468,10 @@ def replace( ) return Series(block.select_column(result_col)) + def interpolate(self, method: str = "linear") -> Series: + result = block_ops.interpolate(self._block, method) + return Series(result) + def dropna( self, *, diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index c96faa3526b..2b710d692a4 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -711,6 +711,22 @@ def test_df_dropna(scalars_dfs, axis, how, ignore_index): pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_df_interpolate(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + columns = ["int64_col", "int64_too", "float64_col"] + bf_result = scalars_df[columns].interpolate().to_pandas() + # Pandas can only interpolate on "float64" columns + # https://github.com/pandas-dev/pandas/issues/40252 + pd_result = scalars_pandas_df[columns].astype("float64").interpolate() + + pandas.testing.assert_frame_equal( + bf_result, + pd_result, + check_index_type=False, + check_dtype=False, + ) + + def test_df_fillna(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs df = scalars_df[["int64_col", "float64_col"]].fillna(3) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 05d8b841855..183ba01c0e0 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -272,6 +272,32 @@ def test_series_replace_list_scalar(scalars_dfs): ) +@pytest.mark.parametrize( + ("values",), + ( + ([None, 1, 2, None, None, 16, None],), + ([None, None, 3.6, None],), + ([403.2, None, 352.1, None, None, 111.9],), + ), +) +def test_series_interpolate(values): + pd_series = pd.Series(values) + bf_series = series.Series(pd_series) + + # Pandas can only interpolate on "float64" columns + # https://github.com/pandas-dev/pandas/issues/40252 + pd_result = pd_series.astype("float64").interpolate() + bf_result = bf_series.interpolate().to_pandas() + + # pd uses non-null types, while bf uses nullable types + pd.testing.assert_series_equal( + pd_result, + bf_result, + check_index_type=False, + check_dtype=False, + ) + + @pytest.mark.parametrize( ("ignore_index",), ( diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 013d1701146..12bd0531795 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2756,6 +2756,43 @@ def value_counts( """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def interpolate(self, method: str = "linear"): + """ + Fill NaN values using an interpolation method. + + Args: + method (str, default 'linear'): + Interpolation technique to use. Only 'linear' supported. + 'linear': Ignore the index and treat the values as equally spaced. + This is the only method supported on MultiIndexes. + + Returns: + DataFrame: + Returns the same object type as the caller, interpolated at + some or all ``NaN`` values + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3, None, None, 6], + ... 'B': [None, 6, None, 2, None, 3], + ... }) + >>> df.interpolate() + A B + 0 1.0 + 1 2.0 6.0 + 2 3.0 4.0 + 3 4.0 2.0 + 4 5.0 2.5 + 5 6.0 3.0 + + [6 rows x 2 columns] + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def fillna(self, value): """ Fill NA/NaN values using the specified method. diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index f0e13e16f5a..b569e5699c7 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -916,6 +916,38 @@ def droplevel(self, level, axis): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def interpolate(self, method: str = "linear"): + """ + Fill NaN values using an interpolation method. + + Args: + method (str, default 'linear'): + Interpolation technique to use. Only 'linear' supported. + 'linear': Ignore the index and treat the values as equally spaced. + This is the only method supported on MultiIndexes. + + Returns: + Series: + Returns the same object type as the caller, interpolated at + some or all ``NaN`` values + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> series = bpd.Series([1, 2, 3, None, None, 6]) + >>> series.interpolate() + 0 1.0 + 1 2.0 + 2 3.0 + 3 4.0 + 4 5.0 + 5 6.0 + dtype: Float64 + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def fillna( self, value=None, From e1817c9201ba4ea7fd2f8b6f4a667b010a6fec1b Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 2 Nov 2023 10:55:01 -0500 Subject: [PATCH 130/197] feat: `to_gbq` without a destination table writes to a temporary table (#158) * feat: `to_gbq` without a destination table writes to a temporary table * add unit test covering happy path for to_gbq * update to_gbq docs --- bigframes/constants.py | 4 ++ bigframes/dataframe.py | 47 +++++++++++---- bigframes/session/__init__.py | 11 ++++ bigframes/session/_io/bigquery.py | 20 +++++++ tests/unit/resources.py | 34 ++++++++++- tests/unit/session/test_io_bigquery.py | 26 +++++++- tests/unit/test_dataframe.py | 59 +++++++++++++++++++ tests/unit/test_pandas.py | 4 +- .../bigframes_vendored/pandas/core/frame.py | 40 +++++++++++-- 9 files changed, 225 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_dataframe.py diff --git a/bigframes/constants.py b/bigframes/constants.py index 90837c79eba..82b48dc9675 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime + """Constants used across BigQuery DataFrames. This module should not depend on any others in the package. @@ -23,3 +25,5 @@ ) 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}" + +DEFAULT_EXPIRATION = datetime.timedelta(days=1) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index ffcaf0d613d..4932008f093 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2289,25 +2289,51 @@ def to_json( def to_gbq( self, - destination_table: str, + destination_table: Optional[str] = None, *, - if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", + if_exists: Optional[Literal["fail", "replace", "append"]] = None, index: bool = True, ordering_id: Optional[str] = None, - ) -> None: - if "." not in destination_table: - raise ValueError( - "Invalid Table Name. Should be of the form 'datasetId.tableId' or " - "'projectId.datasetId.tableId'" - ) - + ) -> str: dispositions = { "fail": bigquery.WriteDisposition.WRITE_EMPTY, "replace": bigquery.WriteDisposition.WRITE_TRUNCATE, "append": bigquery.WriteDisposition.WRITE_APPEND, } + + if destination_table is None: + # TODO(swast): If there have been no modifications to the DataFrame + # since the last time it was written (cached), then return that. + # For `read_gbq` nodes, return the underlying table clone. + destination_table = bigframes.session._io.bigquery.create_temp_table( + self._session.bqclient, + self._session._anonymous_dataset, + # TODO(swast): allow custom expiration times, probably via session configuration. + constants.DEFAULT_EXPIRATION, + ) + + if if_exists is not None and if_exists != "replace": + raise ValueError( + f"Got invalid value {repr(if_exists)} for if_exists. " + "When no destination table is specified, a new table is always created. " + "None or 'replace' are the only valid options in this case." + ) + if_exists = "replace" + + if "." not in destination_table: + raise ValueError( + f"Got invalid value for destination_table {repr(destination_table)}. " + "Should be of the form 'datasetId.tableId' or 'projectId.datasetId.tableId'." + ) + + if if_exists is None: + if_exists = "fail" + if if_exists not in dispositions: - raise ValueError("'{0}' is not valid for if_exists".format(if_exists)) + raise ValueError( + f"Got invalid value {repr(if_exists)} for if_exists. " + f"Valid options include None or one of {dispositions.keys()}." + ) job_config = bigquery.QueryJobConfig( write_disposition=dispositions[if_exists], @@ -2318,6 +2344,7 @@ def to_gbq( ) self._run_io_query(index=index, ordering_id=ordering_id, job_config=job_config) + return destination_table def to_numpy( self, dtype=None, copy=False, na_value=None, **kwargs diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 5a61ed534f3..a1eae697158 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -223,6 +223,17 @@ def _create_and_bind_bq_session(self): query_job.result() # blocks until finished self._session_id = query_job.session_info.session_id + # The anonymous dataset is used by BigQuery to write query results and + # session tables. BigQuery DataFrames also writes temp tables directly + # to the dataset, no BigQuery Session required. Note: there is a + # different anonymous dataset per location. See: + # https://cloud.google.com/bigquery/docs/cached-results#how_cached_results_are_stored + query_destination = query_job.destination + self._anonymous_dataset = bigquery.DatasetReference( + query_destination.project, + query_destination.dataset_id, + ) + self.bqclient.default_query_job_config = bigquery.QueryJobConfig( connection_properties=[ bigquery.ConnectionProperty("session_id", self._session_id) diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index d47efbdddcf..d200a9a861d 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -18,10 +18,12 @@ import textwrap import types from typing import Dict, Iterable, Union +import uuid import google.cloud.bigquery as bigquery IO_ORDERING_ID = "bqdf_row_nums" +TEMP_TABLE_PREFIX = "bqdf{date}_{random_id}" def create_export_csv_statement( @@ -90,6 +92,24 @@ def create_snapshot_sql( ) +def create_temp_table( + bqclient: bigquery.Client, + dataset: bigquery.DatasetReference, + expiration: datetime.timedelta, +) -> str: + """Create an empty table with an expiration in the desired dataset.""" + now = datetime.datetime.now(datetime.timezone.utc) + random_id = uuid.uuid4().hex + table_id = TEMP_TABLE_PREFIX.format( + date=now.strftime("%Y%m%d"), random_id=random_id + ) + table_ref = dataset.table(table_id) + destination = bigquery.Table(table_ref) + destination.expires = now + expiration + bqclient.create_table(destination) + return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" + + # BigQuery REST API returns types in Legacy SQL format # https://cloud.google.com/bigquery/docs/data-types but we use Standard SQL # names diff --git a/tests/unit/resources.py b/tests/unit/resources.py index f660d774f06..8fc8acd1759 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -19,17 +19,21 @@ import google.cloud.bigquery import ibis import pandas +import pytest import bigframes import bigframes.core as core import bigframes.core.ordering +import bigframes.dataframe import bigframes.session.clients """Utilities for creating test resources.""" def create_bigquery_session( - bqclient: Optional[google.cloud.bigquery.Client] = None, session_id: str = "abcxyz" + bqclient: Optional[mock.Mock] = None, + session_id: str = "abcxyz", + anonymous_dataset: Optional[google.cloud.bigquery.DatasetReference] = None, ) -> bigframes.Session: credentials = mock.create_autospec( google.auth.credentials.Credentials, instance=True @@ -39,6 +43,21 @@ def create_bigquery_session( bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) bqclient.project = "test-project" + if anonymous_dataset is None: + anonymous_dataset = google.cloud.bigquery.DatasetReference( + "test-project", + "test_dataset", + ) + + query_job = mock.create_autospec(google.cloud.bigquery.QueryJob) + type(query_job).destination = mock.PropertyMock( + return_value=anonymous_dataset.table("test_table"), + ) + type(query_job).session_info = google.cloud.bigquery.SessionInfo( + {"sessionInfo": {"sessionId": session_id}}, + ) + bqclient.query.return_value = query_job + clients_provider = mock.create_autospec(bigframes.session.clients.ClientsProvider) type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient) clients_provider._credentials = credentials @@ -51,6 +70,19 @@ def create_bigquery_session( return session +def create_dataframe( + monkeypatch: pytest.MonkeyPatch, session: Optional[bigframes.Session] = None +) -> bigframes.dataframe.DataFrame: + if session is None: + session = create_bigquery_session() + + # Since this may create a ReadLocalNode, the session we explicitly pass in + # might not actually be used. Mock out the global session, too. + monkeypatch.setattr(bigframes.core.global_session, "_global_session", session) + bigframes.options.bigquery._session_started = True + return bigframes.dataframe.DataFrame({}, session=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. diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index d2255d5edf6..cb3003b1cc1 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -14,6 +14,7 @@ import datetime from typing import Iterable +import unittest.mock as mock import google.cloud.bigquery as bigquery import pytest @@ -37,7 +38,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): assert "`my-test-project`.`_e8166e0cdb`.`anonbb92cd`" in sql -def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): +def test_create_snapshot_sql_doesnt_timetravel_session_tables(): table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") sql = bigframes.session._io.bigquery.create_snapshot_sql( @@ -51,6 +52,29 @@ def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): assert "my-test-project" not in sql +def test_create_temp_table_default_expiration(): + """Make sure the created table has an expiration.""" + bqclient = mock.create_autospec(bigquery.Client) + dataset = bigquery.DatasetReference("test-project", "test_dataset") + now = datetime.datetime.now(datetime.timezone.utc) + expiration = datetime.timedelta(days=3) + expected_expires = now + expiration + + bigframes.session._io.bigquery.create_temp_table(bqclient, dataset, expiration) + + bqclient.create_table.assert_called_once() + call_args = bqclient.create_table.call_args + table = call_args.args[0] + assert table.project == "test-project" + assert table.dataset_id == "test_dataset" + assert table.table_id.startswith("bqdf") + assert ( + (expected_expires - datetime.timedelta(minutes=1)) + < table.expires + < (expected_expires + datetime.timedelta(minutes=1)) + ) + + @pytest.mark.parametrize( ("schema", "expected"), ( diff --git a/tests/unit/test_dataframe.py b/tests/unit/test_dataframe.py new file mode 100644 index 00000000000..17a82908893 --- /dev/null +++ b/tests/unit/test_dataframe.py @@ -0,0 +1,59 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import google.cloud.bigquery +import pytest + +from . import resources + + +def test_dataframe_to_gbq_invalid_destination(monkeypatch: pytest.MonkeyPatch): + dataframe = resources.create_dataframe(monkeypatch) + + with pytest.raises(ValueError, match="no_dataset_or_project"): + dataframe.to_gbq("no_dataset_or_project") + + +def test_dataframe_to_gbq_invalid_if_exists(monkeypatch: pytest.MonkeyPatch): + dataframe = resources.create_dataframe(monkeypatch) + + with pytest.raises(ValueError, match="notreallyanoption"): + # Even though the type is annotated with the literals we accept, users + # might not be using a type checker, especially not in an interactive + # notebook. + dataframe.to_gbq(if_exists="notreallyanoption") # type: ignore + + +def test_dataframe_to_gbq_invalid_if_exists_no_destination( + monkeypatch: pytest.MonkeyPatch, +): + dataframe = resources.create_dataframe(monkeypatch) + + with pytest.raises(ValueError, match="append"): + dataframe.to_gbq(if_exists="append") + + +def test_dataframe_to_gbq_writes_to_anonymous_dataset( + monkeypatch: pytest.MonkeyPatch, +): + anonymous_dataset_id = "my-anonymous-project.my_anonymous_dataset" + anonymous_dataset = google.cloud.bigquery.DatasetReference.from_string( + anonymous_dataset_id + ) + session = resources.create_bigquery_session(anonymous_dataset=anonymous_dataset) + dataframe = resources.create_dataframe(monkeypatch, session=session) + + destination = dataframe.to_gbq() + + assert destination.startswith(anonymous_dataset_id) diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index 5d4f69c7c01..70c5441c685 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_close_session_after_bq_session_ended(monkeypatch): +def test_close_session_after_bq_session_ended(monkeypatch: pytest.MonkeyPatch): bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) bqclient.project = "test-project" session = resources.create_bigquery_session( @@ -141,7 +141,7 @@ def test_close_session_after_bq_session_ended(monkeypatch): google.api_core.exceptions.BadRequest, match="Session JUST_A_TEST has expired and is no longer available.", ): - bpd.read_gbq("SELECT 1") + bpd.read_gbq("SELECT 'ABC'") # Even though the query to stop the session raises an exception, we should # still be able to close it without raising an error to the user. diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 12bd0531795..e267fac0f73 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -125,12 +125,12 @@ def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarra def to_gbq( self, - destination_table: str, + destination_table: Optional[str], *, - if_exists: Optional[Literal["fail", "replace", "append"]] = "fail", + if_exists: Optional[Literal["fail", "replace", "append"]] = None, index: bool = True, ordering_id: Optional[str] = None, - ) -> None: + ) -> str: """Write a DataFrame to a BigQuery table. **Examples:** @@ -138,17 +138,40 @@ def to_gbq( >>> import bigframes.pandas as bpd >>> bpd.options.display.progress_bar = None + Write a DataFrame to a BigQuery table. + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) >>> # destination_table = PROJECT_ID + "." + DATASET_ID + "." + TABLE_NAME >>> df.to_gbq("bigframes-dev.birds.test-numbers", if_exists="replace") + 'bigframes-dev.birds.test-numbers' + + Write a DataFrame to a temporary BigQuery table in the anonymous dataset. + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> destination = df.to_gbq(ordering_id="ordering_id") + >>> # The table created can be read outside of the current session. + >>> bpd.close_session() # For demonstration, only. + >>> bpd.read_gbq(destination, index_col="ordering_id") + col1 col2 + ordering_id + 0 1 3 + 1 2 4 + + [2 rows x 2 columns] Args: - destination_table (str): + destination_table (Optional[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: + If no ``destination_table`` is set, a new temporary table is + created in the BigQuery anonymous dataset. + + if_exists (Optional[str]): + Behavior when the destination table exists. When + ``destination_table`` is set, this defaults to ``'fail'``. When + ``destination_table`` is not set, this field is not applicable. + A new table is always created. Value can be one of: ``'fail'`` If table exists raise pandas_gbq.gbq.TableCreationError. @@ -163,6 +186,11 @@ def to_gbq( ordering_id (Optional[str], default None): If set, write the ordering of the DataFrame as a column in the result table with this name. + + Returns: + str: + The fully-qualified ID for the written table, in the form + ``project.dataset.tablename``. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 5f0ea37fffff792fc3fbed65e6ace846d8ef6a06 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 2 Nov 2023 10:49:11 -0700 Subject: [PATCH 131/197] feat: support 32k text-generation and multilingual embedding models (#161) * feat: support 32k text-generation and embedding multilingual models --- bigframes/ml/llm.py | 47 ++++++++++++++++------ tests/system/small/ml/conftest.py | 18 +++++++++ tests/system/small/ml/test_llm.py | 65 +++++++++++++++++++++++++++++++ 3 files changed, 118 insertions(+), 12 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index d78f467537c..2e5a9a1e5e5 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import cast, Optional, Union +from typing import cast, Literal, Optional, Union import bigframes from bigframes import clients, constants @@ -25,9 +25,11 @@ import bigframes.pandas as bpd _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" +_REMOTE_TEXT_GENERATOR_32K_MODEL_CODE = "text-bison-32k" _TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE = "CLOUD_AI_TEXT_EMBEDDING_MODEL_V1" +_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE = "textembedding-gecko-multilingual" _EMBED_TEXT_RESULT_COLUMN = "text_embedding" @@ -35,19 +37,25 @@ class PaLM2TextGenerator(base.Predictor): """PaLM2 text generator LLM model. Args: + model_name (str, Default to "text-bison"): + The model for natural language tasks. “text-bison” returns model fine-tuned to follow natural language instructions + and is suitable for a variety of language tasks. "text-bison-32k" supports up to 32k tokens per request. + Default to "text-bison". 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 ... + 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__( self, + model_name: Literal["text-bison", "text-bison-32k"] = "text-bison", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient @@ -80,11 +88,14 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - - options = { - "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, - } - + if self.model_name == "text-bison": + options = { + "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, + } + else: + options = { + "endpoint": _REMOTE_TEXT_GENERATOR_32K_MODEL_CODE, + } return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -118,7 +129,7 @@ def predict( 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). + 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. Possible values [1, 40]. @@ -175,6 +186,10 @@ class PaLM2TextEmbeddingGenerator(base.Predictor): """PaLM2 text embedding generator LLM model. Args: + model_name (str, Default to "textembedding-gecko"): + The model for text embedding. “textembedding-gecko” returns model embeddings for text inputs. + "textembedding-gecko-multilingual" returns model embeddings for text inputs which support over 100 languages + Default to "textembedding-gecko". session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): @@ -184,9 +199,13 @@ class PaLM2TextEmbeddingGenerator(base.Predictor): def __init__( self, + model_name: Literal[ + "textembedding-gecko", "textembedding-gecko-multilingual" + ] = "textembedding-gecko", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient @@ -219,10 +238,14 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - - options = { - "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, - } + if self.model_name == "textembedding-gecko": + options = { + "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, + } + else: + options = { + "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE, + } return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 1dd1c813b81..c11445b79a5 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -213,6 +213,13 @@ def palm2_text_generator_model(session, bq_connection) -> llm.PaLM2TextGenerator return llm.PaLM2TextGenerator(session=session, connection_name=bq_connection) +@pytest.fixture(scope="session") +def palm2_text_generator_32k_model(session, bq_connection) -> llm.PaLM2TextGenerator: + return llm.PaLM2TextGenerator( + model_name="text-bison-32k", session=session, connection_name=bq_connection + ) + + @pytest.fixture(scope="function") def ephemera_palm2_text_generator_model( session, bq_connection @@ -229,6 +236,17 @@ def palm2_embedding_generator_model( ) +@pytest.fixture(scope="session") +def palm2_embedding_generator_multilingual_model( + session, bq_connection +) -> llm.PaLM2TextEmbeddingGenerator: + return llm.PaLM2TextEmbeddingGenerator( + model_name="textembedding-gecko-multilingual", + session=session, + connection_name=bq_connection, + ) + + @pytest.fixture(scope="session") def time_series_bqml_arima_plus_model( session, time_series_arima_plus_model_name diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index b7257dde1be..79d3c40317d 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -26,6 +26,12 @@ 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_32k_model(palm2_text_generator_32k_model): + # Model creation doesn't return error + assert palm2_text_generator_32k_model is not None + assert palm2_text_generator_32k_model._bqml_model is not None + + @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 @@ -48,6 +54,30 @@ def test_create_text_generator_model_default_session(bq_connection, llm_text_pan assert all(series.str.len() > 20) +@pytest.mark.flaky(retries=2, delay=120) +def test_create_text_generator_32k_model_default_session( + bq_connection, llm_text_pandas_df +): + import bigframes.pandas as bpd + + bpd.close_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextGenerator(model_name="text-bison-32k") + 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 @@ -127,6 +157,14 @@ def test_create_embedding_generator_model(palm2_embedding_generator_model): assert palm2_embedding_generator_model._bqml_model is not None +def test_create_embedding_generator_multilingual_model( + palm2_embedding_generator_multilingual_model, +): + # Model creation doesn't return error + assert palm2_embedding_generator_multilingual_model is not None + assert palm2_embedding_generator_multilingual_model._bqml_model is not None + + def test_create_text_embedding_generator_model_defaults(bq_connection): import bigframes.pandas as bpd @@ -139,6 +177,20 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None +def test_create_text_embedding_generator_multilingual_model_defaults(bq_connection): + import bigframes.pandas as bpd + + bpd.close_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextEmbeddingGenerator( + model_name="textembedding-gecko-multilingual" + ) + assert model is not None + assert model._bqml_model is not None + + @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df @@ -152,6 +204,19 @@ def test_embedding_generator_predict_success( assert value.size == 768 +@pytest.mark.flaky(retries=2, delay=120) +def test_embedding_generator_multilingual_predict_success( + palm2_embedding_generator_multilingual_model, llm_text_df +): + df = palm2_embedding_generator_multilingual_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 + + @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_series_success( palm2_embedding_generator_model, llm_text_df From 6d1953b7b46b402a2ea52233141ed6fd338b0098 Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Thu, 2 Nov 2023 20:32:49 -0700 Subject: [PATCH 132/197] chore: update docfx minimum Python version (#167) * chore: update docfx minimum Python version Source-Link: https://github.com/googleapis/synthtool/commit/bc07fd415c39853b382bcf8315f8eeacdf334055 Post-Processor: gcr.io/cloud-devrel-public-resources/owlbot-python:latest@sha256:30470597773378105e239b59fce8eb27cc97375580d592699206d17d117143d0 * chore: remove restriction on noxfile.py --------- Co-authored-by: Owl Bot Co-authored-by: Dan Lee <71398022+dandhlee@users.noreply.github.com> --- .github/.OwlBot.lock.yaml | 4 ++-- .github/workflows/docs.yml | 2 +- noxfile.py | 4 +--- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index 7f291dbd5f9..ec696b558c3 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:4f9b3b106ad0beafc2c8a415e3f62c1a0cc23cabea115dbe841b848f581cfe99 -# created: 2023-10-18T20:26:37.410353675Z + digest: sha256:30470597773378105e239b59fce8eb27cc97375580d592699206d17d117143d0 +# created: 2023-11-03T00:57:07.335914631Z diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index e97d89e484c..221806cedf5 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -28,7 +28,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v4 with: - python-version: "3.9" + python-version: "3.10" - name: Install nox run: | python -m pip install --upgrade setuptools pip wheel diff --git a/noxfile.py b/noxfile.py index d0bbda80fd4..34b055de445 100644 --- a/noxfile.py +++ b/noxfile.py @@ -451,9 +451,7 @@ def docs(session): ) -# docfx doesn't yet support Python 3.10. -# https://github.com/googleapis/sphinx-docfx-yaml/issues/305 -@nox.session(python="3.9") +@nox.session(python=DEFAULT_PYTHON_VERSION) def docfx(session): """Build the docfx yaml files for this library.""" From 031f253890f8a212309097554c3462980654201d Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 3 Nov 2023 13:20:36 -0500 Subject: [PATCH 133/197] fix: use table clone instead of system time for `read_gbq_table` (#109) * fix: use table clone instead of system time for `read_gbq_table` * accept expiration datetime instead of timedelta for easier testing * don't use table clone on _session tables * remove unnecessary assert * add docstrings --- bigframes/constants.py | 2 +- bigframes/dataframe.py | 4 +- bigframes/session/__init__.py | 59 +++++++---------- bigframes/session/_io/bigquery.py | 89 ++++++++++++++++++++------ tests/system/small/test_session.py | 3 - tests/unit/session/test_io_bigquery.py | 64 +++++++++++------- 6 files changed, 137 insertions(+), 84 deletions(-) diff --git a/bigframes/constants.py b/bigframes/constants.py index 82b48dc9675..a1ffd2b755c 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -26,4 +26,4 @@ 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}" -DEFAULT_EXPIRATION = datetime.timedelta(days=1) +DEFAULT_EXPIRATION = datetime.timedelta(days=7) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 4932008f093..45dbcdc78da 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -16,6 +16,7 @@ from __future__ import annotations +import datetime import re import textwrap import typing @@ -2309,7 +2310,8 @@ def to_gbq( self._session.bqclient, self._session._anonymous_dataset, # TODO(swast): allow custom expiration times, probably via session configuration. - constants.DEFAULT_EXPIRATION, + datetime.datetime.now(datetime.timezone.utc) + + constants.DEFAULT_EXPIRATION, ) if if_exists is not None and if_exists != "replace": diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index a1eae697158..12ee91a13a8 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -16,6 +16,7 @@ from __future__ import annotations +import datetime import logging import os import re @@ -430,7 +431,9 @@ def _read_gbq_query( index_cols = list(index_col) destination, query_job = self._query_to_destination( - query, index_cols, api_name="read_gbq_query" + query, + index_cols, + api_name=api_name, ) # If there was no destination table, that means the query must have @@ -508,6 +511,12 @@ def _read_gbq_table_to_ibis_with_total_ordering( If we can get a total ordering from the table, such as via primary key column(s), then return those too so that ordering generation can be avoided. + + For tables that aren't already read-only, this creates Create a table + clone so that any changes to the underlying table don't affect the + DataFrame and break our assumptions, especially with regards to unique + index and ordering. See: + https://cloud.google.com/bigquery/docs/table-clones-create """ if table_ref.dataset_id.upper() == "_SESSION": # _SESSION tables aren't supported by the tables.get REST API. @@ -518,15 +527,24 @@ def _read_gbq_table_to_ibis_with_total_ordering( None, ) + now = datetime.datetime.now(datetime.timezone.utc) + destination = bigframes_io.create_table_clone( + table_ref, + self._anonymous_dataset, + # TODO(swast): Allow the default expiration to be configured. + now + constants.DEFAULT_EXPIRATION, + self, + api_name, + ) table_expression = self.ibis_client.table( - table_ref.table_id, - database=f"{table_ref.project}.{table_ref.dataset_id}", + destination.table_id, + database=f"{destination.project}.{destination.dataset_id}", ) # If there are primary keys defined, the query engine assumes these # columns are unique, even if the constraint is not enforced. We make # the same assumption and use these columns as the total ordering keys. - table = self.bqclient.get_table(table_ref) + table = self.bqclient.get_table(destination) # TODO(b/305264153): Use public properties to fetch primary keys once # added to google-cloud-bigquery. @@ -535,23 +553,7 @@ def _read_gbq_table_to_ibis_with_total_ordering( .get("primaryKey", {}) .get("columns") ) - - if not primary_keys: - return table_expression, None - else: - # Read from a snapshot since we won't have to copy the table data to create a total ordering. - job_config = bigquery.QueryJobConfig() - job_config.labels["bigframes-api"] = api_name - current_timestamp = list( - self.bqclient.query( - "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", - job_config=job_config, - ).result() - )[0][0] - table_expression = self.ibis_client.sql( - bigframes_io.create_snapshot_sql(table_ref, current_timestamp) - ) - return table_expression, primary_keys + return table_expression, primary_keys def _read_gbq_table( self, @@ -662,20 +664,7 @@ def _read_gbq_table( total_ordering_columns=frozenset(index_cols), ) - # 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: + 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) diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index d200a9a861d..fd3b1c59a7b 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -14,14 +14,21 @@ """Private module: Helpers for I/O operations.""" +from __future__ import annotations + import datetime import textwrap import types +import typing from typing import Dict, Iterable, Union import uuid import google.cloud.bigquery as bigquery +if typing.TYPE_CHECKING: + import bigframes.session + + IO_ORDERING_ID = "bqdf_row_nums" TEMP_TABLE_PREFIX = "bqdf{date}_{random_id}" @@ -69,43 +76,83 @@ def create_export_data_statement( ) -def create_snapshot_sql( - table_ref: bigquery.TableReference, current_timestamp: datetime.datetime -) -> str: - """Query a table via 'time travel' for consistent reads.""" +def random_table(dataset: bigquery.DatasetReference) -> bigquery.TableReference: + """Generate a random table ID with BigQuery DataFrames prefix. + + Args: + dataset (google.cloud.bigquery.DatasetReference): + The dataset to make the table reference in. Usually the anonymous + dataset for the session. + + Returns: + google.cloud.bigquery.TableReference: + Fully qualified table ID of a table that doesn't exist. + """ + now = datetime.datetime.now(datetime.timezone.utc) + random_id = uuid.uuid4().hex + table_id = TEMP_TABLE_PREFIX.format( + date=now.strftime("%Y%m%d"), random_id=random_id + ) + return dataset.table(table_id) + - # 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}`" +def table_ref_to_sql(table: bigquery.TableReference) -> str: + """Format a table reference as escaped SQL.""" + return f"`{table.project}`.`{table.dataset_id}`.`{table.table_id}`" + +def create_table_clone( + source: bigquery.TableReference, + dataset: bigquery.DatasetReference, + expiration: datetime.datetime, + session: bigframes.session.Session, + api_name: str, +) -> bigquery.TableReference: + """Create a table clone for consistent reads.""" # 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}`" + if source.dataset_id.startswith("_"): + return source - return textwrap.dedent( + fully_qualified_source_id = table_ref_to_sql(source) + destination = random_table(dataset) + fully_qualified_destination_id = table_ref_to_sql(destination) + + # Include a label so that Dataplex Lineage can identify temporary + # tables that BigQuery DataFrames creates. Googlers: See internal issue + # 296779699. + ddl = 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())}) + CREATE OR REPLACE TABLE + {fully_qualified_destination_id} + CLONE {fully_qualified_source_id} + OPTIONS( + expiration_timestamp=TIMESTAMP "{expiration.isoformat()}", + labels=[ + ("source", "bigquery-dataframes-temp"), + ("bigframes-api", {repr(api_name)}) + ] + ) """ ) + job_config = bigquery.QueryJobConfig() + job_config.labels = { + "source": "bigquery-dataframes-temp", + "bigframes-api": api_name, + } + session._start_query(ddl, job_config=job_config) + return destination def create_temp_table( bqclient: bigquery.Client, dataset: bigquery.DatasetReference, - expiration: datetime.timedelta, + expiration: datetime.datetime, ) -> str: """Create an empty table with an expiration in the desired dataset.""" - now = datetime.datetime.now(datetime.timezone.utc) - random_id = uuid.uuid4().hex - table_id = TEMP_TABLE_PREFIX.format( - date=now.strftime("%Y%m%d"), random_id=random_id - ) - table_ref = dataset.table(table_id) + table_ref = random_table(dataset) destination = bigquery.Table(table_ref) - destination.expires = now + expiration + destination.expires = expiration bqclient.create_table(destination) return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index bf72e444eba..28486a1269b 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -252,9 +252,6 @@ def test_read_gbq_w_primary_keys_table( sorted_result = result.sort_values(primary_keys) pd.testing.assert_frame_equal(result, sorted_result) - # Verify that we're working from a snapshot rather than a copy of the table. - assert "FOR SYSTEM_TIME AS OF TIMESTAMP" in df.sql - @pytest.mark.parametrize( ("query_or_table", "max_results"), diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index cb3003b1cc1..7a8691232ba 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -19,46 +19,63 @@ import google.cloud.bigquery as bigquery import pytest +import bigframes.session import bigframes.session._io.bigquery -def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): - table_ref = bigquery.TableReference.from_string( +def test_create_table_clone_doesnt_clone_anonymous_datasets(): + session = mock.create_autospec(bigframes.session.Session) + source = bigquery.TableReference.from_string( "my-test-project._e8166e0cdb.anonbb92cd" ) - sql = bigframes.session._io.bigquery.create_snapshot_sql( - table_ref, datetime.datetime.now(datetime.timezone.utc) + destination = bigframes.session._io.bigquery.create_table_clone( + source, + bigquery.DatasetReference("other-project", "other_dataset"), + datetime.datetime(2023, 11, 2, 15, 43, 21, tzinfo=datetime.timezone.utc), + session, + "test_api", ) - # Anonymous query results tables don't support time travel. - assert "SYSTEM_TIME" not in sql + # Anonymous query results tables don't support CLONE + assert destination is source + session._start_query.assert_not_called() - # Need fully-qualified table name. - assert "`my-test-project`.`_e8166e0cdb`.`anonbb92cd`" in sql - -def test_create_snapshot_sql_doesnt_timetravel_session_tables(): - table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") - - sql = bigframes.session._io.bigquery.create_snapshot_sql( - table_ref, datetime.datetime.now(datetime.timezone.utc) +def test_create_table_clone_sets_expiration(): + session = mock.create_autospec(bigframes.session.Session) + source = bigquery.TableReference.from_string( + "my-test-project.test_dataset.some_table" ) - # We aren't modifying _SESSION tables, so don't use time travel. - assert "SYSTEM_TIME" not in sql + expiration = datetime.datetime( + 2023, 11, 2, 15, 43, 21, tzinfo=datetime.timezone.utc + ) + bigframes.session._io.bigquery.create_table_clone( + source, + bigquery.DatasetReference("other-project", "other_dataset"), + expiration, + session, + "test_api", + ) - # Don't need the project ID for _SESSION tables. - assert "my-test-project" not in sql + session._start_query.assert_called_once() + call_args = session._start_query.call_args + query = call_args.args[0] + assert "CREATE OR REPLACE TABLE" in query + assert "CLONE" in query + assert f'expiration_timestamp=TIMESTAMP "{expiration.isoformat()}"' in query + assert '("source", "bigquery-dataframes-temp")' in query + assert call_args.kwargs["job_config"].labels["bigframes-api"] == "test_api" def test_create_temp_table_default_expiration(): """Make sure the created table has an expiration.""" bqclient = mock.create_autospec(bigquery.Client) dataset = bigquery.DatasetReference("test-project", "test_dataset") - now = datetime.datetime.now(datetime.timezone.utc) - expiration = datetime.timedelta(days=3) - expected_expires = now + expiration + expiration = datetime.datetime( + 2023, 11, 2, 13, 44, 55, 678901, datetime.timezone.utc + ) bigframes.session._io.bigquery.create_temp_table(bqclient, dataset, expiration) @@ -68,10 +85,11 @@ def test_create_temp_table_default_expiration(): assert table.project == "test-project" assert table.dataset_id == "test_dataset" assert table.table_id.startswith("bqdf") + # TODO(swast): Why isn't the expiration exactly what we set it to? assert ( - (expected_expires - datetime.timedelta(minutes=1)) + (expiration - datetime.timedelta(minutes=1)) < table.expires - < (expected_expires + datetime.timedelta(minutes=1)) + < (expiration + datetime.timedelta(minutes=1)) ) From c065071028c2f4ac80ee7f84dbeb1df385c2a512 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 3 Nov 2023 13:42:14 -0700 Subject: [PATCH 134/197] feat: add __iter__, iterrows, itertuples, keys methods (#164) 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/dataframe.py | 18 +++++ bigframes/series.py | 6 ++ tests/system/small/test_dataframe.py | 49 ++++++++++++ .../bigframes_vendored/pandas/core/frame.py | 79 +++++++++++++++++++ .../bigframes_vendored/pandas/core/generic.py | 31 +++++++- 5 files changed, 182 insertions(+), 1 deletion(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 45dbcdc78da..40f12671aef 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -304,6 +304,9 @@ def __len__(self): rows, _ = self.shape return rows + def __iter__(self): + return iter(self.columns) + def astype( self, dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype], @@ -1477,12 +1480,27 @@ def isin(self, values) -> DataFrame: f"isin(), you passed a [{type(values).__name__}]" ) + def keys(self) -> pandas.Index: + return self.columns + 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 iterrows(self) -> Iterable[tuple[typing.Any, pandas.Series]]: + for df in self.to_pandas_batches(): + for item in df.iterrows(): + yield item + + def itertuples( + self, index: bool = True, name: typing.Optional[str] = "Pandas" + ) -> Iterable[tuple[typing.Any, ...]]: + for df in self.to_pandas_batches(): + for item in df.itertuples(index=index, name=name): + yield item + def dropna( self, *, diff --git a/bigframes/series.py b/bigframes/series.py index 824757cf52b..032bdf6c429 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -16,6 +16,7 @@ from __future__ import annotations +import itertools import numbers import textwrap import typing @@ -148,6 +149,11 @@ def _set_internal_query_job(self, query_job: bigquery.QueryJob): def __len__(self): return self.shape[0] + def __iter__(self) -> typing.Iterator: + return itertools.chain.from_iterable( + map(lambda x: x.index, self._block.to_pandas_batches()) + ) + def copy(self) -> Series: return Series(self._block) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 2b710d692a4..bd5930e5089 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -803,6 +803,55 @@ def test_apply_series_scalar_callable( pandas.testing.assert_series_equal(bf_result, pd_result) +def test_df_keys( + scalars_df_index, + scalars_pandas_df_index, +): + pandas.testing.assert_index_equal( + scalars_df_index.keys(), scalars_pandas_df_index.keys() + ) + + +def test_df_iter( + scalars_df_index, + scalars_pandas_df_index, +): + for bf_i, df_i in zip(scalars_df_index, scalars_pandas_df_index): + assert bf_i == df_i + + +def test_iterrows( + scalars_df_index, + scalars_pandas_df_index, +): + for (bf_index, bf_series), (pd_index, pd_series) in zip( + scalars_df_index.iterrows(), scalars_pandas_df_index.iterrows() + ): + assert bf_index == pd_index + pandas.testing.assert_series_equal(bf_series, pd_series) + + +@pytest.mark.parametrize( + ( + "index", + "name", + ), + [ + ( + True, + "my_df", + ), + (False, None), + ], +) +def test_itertuples(scalars_df_index, index, name): + # Numeric has slightly different representation as a result of conversions. + bf_tuples = scalars_df_index.itertuples(index, name) + pd_tuples = scalars_df_index.to_pandas().itertuples(index, name) + for bf_tuple, pd_tuple in zip(bf_tuples, pd_tuples): + assert bf_tuple == pd_tuple + + 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 e267fac0f73..6f4f6be35d6 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -975,6 +975,85 @@ def isin(self, values): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def keys(self): + """ + Get the 'info axis'. + + This is index for Series, columns for DataFrame. + + Returns: + Index: Info axis. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> df.keys() + Index(['A', 'B'], dtype='object') + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def iterrows(self): + """ + Iterate over DataFrame rows as (index, Series) pairs. + + Yields: + a tuple (index, data) where data contains row values as a Series + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> index, row = next(df.iterrows()) + >>> index + 0 + >>> row + A 1 + B 4 + Name: 0, dtype: object + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def itertuples(self, index: bool = True, name: str | None = "Pandas"): + """ + Iterate over DataFrame rows as namedtuples. + + Args: + index (bool, default True): + If True, return the index as the first element of the tuple. + name (str or None, default "Pandas"): + The name of the returned namedtuples or None to return regular + tuples. + + Returns: + iterator: + An object to iterate over namedtuples for each row in the + DataFrame with the first field possibly being the index and + following fields being the column values. + + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> next(df.itertuples(name="Pair")) + Pair(Index=0, A=1, B=4) + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def items(self): """ Iterate over (column name, Series) pairs. diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 27d2e845376..127efe6a3d6 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -1,7 +1,7 @@ # 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 typing import Iterator, Literal, Optional from bigframes import constants from third_party.bigframes_vendored.pandas.core import indexing @@ -35,6 +35,35 @@ def size(self) -> int: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def __iter__(self) -> Iterator: + """ + Iterate over info axis. + + Returns + iterator: Info axis as iterator. + + **Examples:** + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3], + ... 'B': [4, 5, 6], + ... }) + >>> for x in df: + ... print(x) + A + B + + >>> series = bpd.Series(["a", "b", "c"], index=[10, 20, 30]) + >>> for x in series: + ... print(x) + 10 + 20 + 30 + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ------------------------------------------------------------------------- # Unary Methods From dfcc2d3f6918785a3048681c61e0f5f6c99d9d95 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 3 Nov 2023 21:44:23 -0500 Subject: [PATCH 135/197] Revert "fix: use table clone instead of system time for `read_gbq_table` (#109)" (#171) This reverts commit 031f253890f8a212309097554c3462980654201d. --- bigframes/constants.py | 2 +- bigframes/dataframe.py | 4 +- bigframes/session/__init__.py | 59 ++++++++++------- bigframes/session/_io/bigquery.py | 89 ++++++-------------------- tests/system/small/test_session.py | 3 + tests/unit/session/test_io_bigquery.py | 64 +++++++----------- 6 files changed, 84 insertions(+), 137 deletions(-) diff --git a/bigframes/constants.py b/bigframes/constants.py index a1ffd2b755c..82b48dc9675 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -26,4 +26,4 @@ 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}" -DEFAULT_EXPIRATION = datetime.timedelta(days=7) +DEFAULT_EXPIRATION = datetime.timedelta(days=1) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 40f12671aef..04a5456e265 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -16,7 +16,6 @@ from __future__ import annotations -import datetime import re import textwrap import typing @@ -2328,8 +2327,7 @@ def to_gbq( self._session.bqclient, self._session._anonymous_dataset, # TODO(swast): allow custom expiration times, probably via session configuration. - datetime.datetime.now(datetime.timezone.utc) - + constants.DEFAULT_EXPIRATION, + constants.DEFAULT_EXPIRATION, ) if if_exists is not None and if_exists != "replace": diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 12ee91a13a8..a1eae697158 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -16,7 +16,6 @@ from __future__ import annotations -import datetime import logging import os import re @@ -431,9 +430,7 @@ def _read_gbq_query( index_cols = list(index_col) destination, query_job = self._query_to_destination( - query, - index_cols, - api_name=api_name, + query, index_cols, api_name="read_gbq_query" ) # If there was no destination table, that means the query must have @@ -511,12 +508,6 @@ def _read_gbq_table_to_ibis_with_total_ordering( If we can get a total ordering from the table, such as via primary key column(s), then return those too so that ordering generation can be avoided. - - For tables that aren't already read-only, this creates Create a table - clone so that any changes to the underlying table don't affect the - DataFrame and break our assumptions, especially with regards to unique - index and ordering. See: - https://cloud.google.com/bigquery/docs/table-clones-create """ if table_ref.dataset_id.upper() == "_SESSION": # _SESSION tables aren't supported by the tables.get REST API. @@ -527,24 +518,15 @@ def _read_gbq_table_to_ibis_with_total_ordering( None, ) - now = datetime.datetime.now(datetime.timezone.utc) - destination = bigframes_io.create_table_clone( - table_ref, - self._anonymous_dataset, - # TODO(swast): Allow the default expiration to be configured. - now + constants.DEFAULT_EXPIRATION, - self, - api_name, - ) table_expression = self.ibis_client.table( - destination.table_id, - database=f"{destination.project}.{destination.dataset_id}", + table_ref.table_id, + database=f"{table_ref.project}.{table_ref.dataset_id}", ) # If there are primary keys defined, the query engine assumes these # columns are unique, even if the constraint is not enforced. We make # the same assumption and use these columns as the total ordering keys. - table = self.bqclient.get_table(destination) + table = self.bqclient.get_table(table_ref) # TODO(b/305264153): Use public properties to fetch primary keys once # added to google-cloud-bigquery. @@ -553,7 +535,23 @@ def _read_gbq_table_to_ibis_with_total_ordering( .get("primaryKey", {}) .get("columns") ) - return table_expression, primary_keys + + if not primary_keys: + return table_expression, None + else: + # Read from a snapshot since we won't have to copy the table data to create a total ordering. + job_config = bigquery.QueryJobConfig() + job_config.labels["bigframes-api"] = api_name + current_timestamp = list( + self.bqclient.query( + "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", + job_config=job_config, + ).result() + )[0][0] + table_expression = self.ibis_client.sql( + bigframes_io.create_snapshot_sql(table_ref, current_timestamp) + ) + return table_expression, primary_keys def _read_gbq_table( self, @@ -664,7 +662,20 @@ def _read_gbq_table( 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) diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index fd3b1c59a7b..d200a9a861d 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -14,21 +14,14 @@ """Private module: Helpers for I/O operations.""" -from __future__ import annotations - import datetime import textwrap import types -import typing from typing import Dict, Iterable, Union import uuid import google.cloud.bigquery as bigquery -if typing.TYPE_CHECKING: - import bigframes.session - - IO_ORDERING_ID = "bqdf_row_nums" TEMP_TABLE_PREFIX = "bqdf{date}_{random_id}" @@ -76,83 +69,43 @@ def create_export_data_statement( ) -def random_table(dataset: bigquery.DatasetReference) -> bigquery.TableReference: - """Generate a random table ID with BigQuery DataFrames prefix. - - Args: - dataset (google.cloud.bigquery.DatasetReference): - The dataset to make the table reference in. Usually the anonymous - dataset for the session. - - Returns: - google.cloud.bigquery.TableReference: - Fully qualified table ID of a table that doesn't exist. - """ - now = datetime.datetime.now(datetime.timezone.utc) - random_id = uuid.uuid4().hex - table_id = TEMP_TABLE_PREFIX.format( - date=now.strftime("%Y%m%d"), random_id=random_id - ) - return dataset.table(table_id) - - -def table_ref_to_sql(table: bigquery.TableReference) -> str: - """Format a table reference as escaped SQL.""" - return f"`{table.project}`.`{table.dataset_id}`.`{table.table_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}`" -def create_table_clone( - source: bigquery.TableReference, - dataset: bigquery.DatasetReference, - expiration: datetime.datetime, - session: bigframes.session.Session, - api_name: str, -) -> bigquery.TableReference: - """Create a table clone for consistent reads.""" # If we have an anonymous query results table, it can't be modified and # there isn't any BigQuery time travel. - if source.dataset_id.startswith("_"): - return source - - fully_qualified_source_id = table_ref_to_sql(source) - destination = random_table(dataset) - fully_qualified_destination_id = table_ref_to_sql(destination) + if table_ref.dataset_id.startswith("_"): + return f"SELECT * FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}`" - # Include a label so that Dataplex Lineage can identify temporary - # tables that BigQuery DataFrames creates. Googlers: See internal issue - # 296779699. - ddl = textwrap.dedent( + return textwrap.dedent( f""" - CREATE OR REPLACE TABLE - {fully_qualified_destination_id} - CLONE {fully_qualified_source_id} - OPTIONS( - expiration_timestamp=TIMESTAMP "{expiration.isoformat()}", - labels=[ - ("source", "bigquery-dataframes-temp"), - ("bigframes-api", {repr(api_name)}) - ] - ) + SELECT * + FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}` + FOR SYSTEM_TIME AS OF TIMESTAMP({repr(current_timestamp.isoformat())}) """ ) - job_config = bigquery.QueryJobConfig() - job_config.labels = { - "source": "bigquery-dataframes-temp", - "bigframes-api": api_name, - } - session._start_query(ddl, job_config=job_config) - return destination def create_temp_table( bqclient: bigquery.Client, dataset: bigquery.DatasetReference, - expiration: datetime.datetime, + expiration: datetime.timedelta, ) -> str: """Create an empty table with an expiration in the desired dataset.""" - table_ref = random_table(dataset) + now = datetime.datetime.now(datetime.timezone.utc) + random_id = uuid.uuid4().hex + table_id = TEMP_TABLE_PREFIX.format( + date=now.strftime("%Y%m%d"), random_id=random_id + ) + table_ref = dataset.table(table_id) destination = bigquery.Table(table_ref) - destination.expires = expiration + destination.expires = now + expiration bqclient.create_table(destination) return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 28486a1269b..bf72e444eba 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -252,6 +252,9 @@ def test_read_gbq_w_primary_keys_table( sorted_result = result.sort_values(primary_keys) pd.testing.assert_frame_equal(result, sorted_result) + # Verify that we're working from a snapshot rather than a copy of the table. + assert "FOR SYSTEM_TIME AS OF TIMESTAMP" in df.sql + @pytest.mark.parametrize( ("query_or_table", "max_results"), diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index 7a8691232ba..cb3003b1cc1 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -19,63 +19,46 @@ import google.cloud.bigquery as bigquery import pytest -import bigframes.session import bigframes.session._io.bigquery -def test_create_table_clone_doesnt_clone_anonymous_datasets(): - session = mock.create_autospec(bigframes.session.Session) - source = bigquery.TableReference.from_string( +def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): + table_ref = bigquery.TableReference.from_string( "my-test-project._e8166e0cdb.anonbb92cd" ) - destination = bigframes.session._io.bigquery.create_table_clone( - source, - bigquery.DatasetReference("other-project", "other_dataset"), - datetime.datetime(2023, 11, 2, 15, 43, 21, tzinfo=datetime.timezone.utc), - session, - "test_api", + sql = bigframes.session._io.bigquery.create_snapshot_sql( + table_ref, datetime.datetime.now(datetime.timezone.utc) ) - # Anonymous query results tables don't support CLONE - assert destination is source - session._start_query.assert_not_called() + # 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_table_clone_sets_expiration(): - session = mock.create_autospec(bigframes.session.Session) - source = bigquery.TableReference.from_string( - "my-test-project.test_dataset.some_table" - ) - expiration = datetime.datetime( - 2023, 11, 2, 15, 43, 21, tzinfo=datetime.timezone.utc - ) - bigframes.session._io.bigquery.create_table_clone( - source, - bigquery.DatasetReference("other-project", "other_dataset"), - expiration, - session, - "test_api", +def test_create_snapshot_sql_doesnt_timetravel_session_tables(): + table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") + + sql = bigframes.session._io.bigquery.create_snapshot_sql( + table_ref, datetime.datetime.now(datetime.timezone.utc) ) - session._start_query.assert_called_once() - call_args = session._start_query.call_args - query = call_args.args[0] - assert "CREATE OR REPLACE TABLE" in query - assert "CLONE" in query - assert f'expiration_timestamp=TIMESTAMP "{expiration.isoformat()}"' in query - assert '("source", "bigquery-dataframes-temp")' in query - assert call_args.kwargs["job_config"].labels["bigframes-api"] == "test_api" + # 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 def test_create_temp_table_default_expiration(): """Make sure the created table has an expiration.""" bqclient = mock.create_autospec(bigquery.Client) dataset = bigquery.DatasetReference("test-project", "test_dataset") - expiration = datetime.datetime( - 2023, 11, 2, 13, 44, 55, 678901, datetime.timezone.utc - ) + now = datetime.datetime.now(datetime.timezone.utc) + expiration = datetime.timedelta(days=3) + expected_expires = now + expiration bigframes.session._io.bigquery.create_temp_table(bqclient, dataset, expiration) @@ -85,11 +68,10 @@ def test_create_temp_table_default_expiration(): assert table.project == "test-project" assert table.dataset_id == "test_dataset" assert table.table_id.startswith("bqdf") - # TODO(swast): Why isn't the expiration exactly what we set it to? assert ( - (expiration - datetime.timedelta(minutes=1)) + (expected_expires - datetime.timedelta(minutes=1)) < table.expires - < (expiration + datetime.timedelta(minutes=1)) + < (expected_expires + datetime.timedelta(minutes=1)) ) From 4ff26cdf862e9f9b91a3a1d2abfa7fbdf0af9c5b Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 6 Nov 2023 12:35:06 -0600 Subject: [PATCH 136/197] fix: update default temp table expiration to 7 days (#174) --- bigframes/constants.py | 2 +- bigframes/dataframe.py | 4 ++- bigframes/session/__init__.py | 4 ++- bigframes/session/_io/bigquery.py | 36 ++++++++++++++++++++------ tests/unit/session/test_io_bigquery.py | 10 +++---- 5 files changed, 40 insertions(+), 16 deletions(-) diff --git a/bigframes/constants.py b/bigframes/constants.py index 82b48dc9675..a1ffd2b755c 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -26,4 +26,4 @@ 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}" -DEFAULT_EXPIRATION = datetime.timedelta(days=1) +DEFAULT_EXPIRATION = datetime.timedelta(days=7) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 04a5456e265..40f12671aef 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -16,6 +16,7 @@ from __future__ import annotations +import datetime import re import textwrap import typing @@ -2327,7 +2328,8 @@ def to_gbq( self._session.bqclient, self._session._anonymous_dataset, # TODO(swast): allow custom expiration times, probably via session configuration. - constants.DEFAULT_EXPIRATION, + datetime.datetime.now(datetime.timezone.utc) + + constants.DEFAULT_EXPIRATION, ) if if_exists is not None and if_exists != "replace": diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index a1eae697158..2537e81e19c 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -430,7 +430,9 @@ def _read_gbq_query( index_cols = list(index_col) destination, query_job = self._query_to_destination( - query, index_cols, api_name="read_gbq_query" + query, + index_cols, + api_name=api_name, ) # If there was no destination table, that means the query must have diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index d200a9a861d..06d240fec65 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -14,6 +14,8 @@ """Private module: Helpers for I/O operations.""" +from __future__ import annotations + import datetime import textwrap import types @@ -69,6 +71,29 @@ def create_export_data_statement( ) +def random_table(dataset: bigquery.DatasetReference) -> bigquery.TableReference: + """Generate a random table ID with BigQuery DataFrames prefix. + Args: + dataset (google.cloud.bigquery.DatasetReference): + The dataset to make the table reference in. Usually the anonymous + dataset for the session. + Returns: + google.cloud.bigquery.TableReference: + Fully qualified table ID of a table that doesn't exist. + """ + now = datetime.datetime.now(datetime.timezone.utc) + random_id = uuid.uuid4().hex + table_id = TEMP_TABLE_PREFIX.format( + date=now.strftime("%Y%m%d"), random_id=random_id + ) + return dataset.table(table_id) + + +def table_ref_to_sql(table: bigquery.TableReference) -> str: + """Format a table reference as escaped SQL.""" + return f"`{table.project}`.`{table.dataset_id}`.`{table.table_id}`" + + def create_snapshot_sql( table_ref: bigquery.TableReference, current_timestamp: datetime.datetime ) -> str: @@ -95,17 +120,12 @@ def create_snapshot_sql( def create_temp_table( bqclient: bigquery.Client, dataset: bigquery.DatasetReference, - expiration: datetime.timedelta, + expiration: datetime.datetime, ) -> str: """Create an empty table with an expiration in the desired dataset.""" - now = datetime.datetime.now(datetime.timezone.utc) - random_id = uuid.uuid4().hex - table_id = TEMP_TABLE_PREFIX.format( - date=now.strftime("%Y%m%d"), random_id=random_id - ) - table_ref = dataset.table(table_id) + table_ref = random_table(dataset) destination = bigquery.Table(table_ref) - destination.expires = now + expiration + destination.expires = expiration bqclient.create_table(destination) return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index cb3003b1cc1..03470208e42 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -56,9 +56,9 @@ def test_create_temp_table_default_expiration(): """Make sure the created table has an expiration.""" bqclient = mock.create_autospec(bigquery.Client) dataset = bigquery.DatasetReference("test-project", "test_dataset") - now = datetime.datetime.now(datetime.timezone.utc) - expiration = datetime.timedelta(days=3) - expected_expires = now + expiration + expiration = datetime.datetime( + 2023, 11, 2, 13, 44, 55, 678901, datetime.timezone.utc + ) bigframes.session._io.bigquery.create_temp_table(bqclient, dataset, expiration) @@ -69,9 +69,9 @@ def test_create_temp_table_default_expiration(): assert table.dataset_id == "test_dataset" assert table.table_id.startswith("bqdf") assert ( - (expected_expires - datetime.timedelta(minutes=1)) + (expiration - datetime.timedelta(minutes=1)) < table.expires - < (expected_expires + datetime.timedelta(minutes=1)) + < (expiration + datetime.timedelta(minutes=1)) ) From 8b6b1c610750e2d75d465a36f829af0b8c835b19 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Tue, 7 Nov 2023 16:56:16 +0000 Subject: [PATCH 137/197] chore(main): release 0.13.0 (#165) :robot: I have created a release *beep* *boop* --- ## [0.13.0](https://togithub.com/googleapis/python-bigquery-dataframes/compare/v0.12.0...v0.13.0) (2023-11-07) ### Features * `to_gbq` without a destination table writes to a temporary table ([#158](https://togithub.com/googleapis/python-bigquery-dataframes/issues/158)) ([e1817c9](https://togithub.com/googleapis/python-bigquery-dataframes/commit/e1817c9201ba4ea7fd2f8b6f4a667b010a6fec1b)) * Add `DataFrame.__iter__`, `DataFrame.iterrows`, `DataFrame.itertuples`, and `DataFrame.keys` methods ([#164](https://togithub.com/googleapis/python-bigquery-dataframes/issues/164)) ([c065071](https://togithub.com/googleapis/python-bigquery-dataframes/commit/c065071028c2f4ac80ee7f84dbeb1df385c2a512)) * Add `Series.__iter__` method ([#164](https://togithub.com/googleapis/python-bigquery-dataframes/issues/164)) ([c065071](https://togithub.com/googleapis/python-bigquery-dataframes/commit/c065071028c2f4ac80ee7f84dbeb1df385c2a512)) * Add interpolate() to series and dataframe ([#157](https://togithub.com/googleapis/python-bigquery-dataframes/issues/157)) ([b9cb55c](https://togithub.com/googleapis/python-bigquery-dataframes/commit/b9cb55c5b9354f9ff60de0aad66fe60049876055)) * Support 32k text-generation and multilingual embedding models ([#161](https://togithub.com/googleapis/python-bigquery-dataframes/issues/161)) ([5f0ea37](https://togithub.com/googleapis/python-bigquery-dataframes/commit/5f0ea37fffff792fc3fbed65e6ace846d8ef6a06)) ### Bug Fixes * Update default temp table expiration to 7 days ([#174](https://togithub.com/googleapis/python-bigquery-dataframes/issues/174)) ([4ff26cd](https://togithub.com/googleapis/python-bigquery-dataframes/commit/4ff26cdf862e9f9b91a3a1d2abfa7fbdf0af9c5b)) --- 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 845d3634bc9..fc327b2e966 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,22 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.13.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.12.0...v0.13.0) (2023-11-07) + + +### Features + +* `to_gbq` without a destination table writes to a temporary table ([#158](https://github.com/googleapis/python-bigquery-dataframes/issues/158)) ([e1817c9](https://github.com/googleapis/python-bigquery-dataframes/commit/e1817c9201ba4ea7fd2f8b6f4a667b010a6fec1b)) +* Add `DataFrame.__iter__`, `DataFrame.iterrows`, `DataFrame.itertuples`, and `DataFrame.keys` methods ([#164](https://github.com/googleapis/python-bigquery-dataframes/issues/164)) ([c065071](https://github.com/googleapis/python-bigquery-dataframes/commit/c065071028c2f4ac80ee7f84dbeb1df385c2a512)) +* Add `Series.__iter__` method ([#164](https://github.com/googleapis/python-bigquery-dataframes/issues/164)) ([c065071](https://github.com/googleapis/python-bigquery-dataframes/commit/c065071028c2f4ac80ee7f84dbeb1df385c2a512)) +* Add interpolate() to series and dataframe ([#157](https://github.com/googleapis/python-bigquery-dataframes/issues/157)) ([b9cb55c](https://github.com/googleapis/python-bigquery-dataframes/commit/b9cb55c5b9354f9ff60de0aad66fe60049876055)) +* Support 32k text-generation and multilingual embedding models ([#161](https://github.com/googleapis/python-bigquery-dataframes/issues/161)) ([5f0ea37](https://github.com/googleapis/python-bigquery-dataframes/commit/5f0ea37fffff792fc3fbed65e6ace846d8ef6a06)) + + +### Bug Fixes + +* Update default temp table expiration to 7 days ([#174](https://github.com/googleapis/python-bigquery-dataframes/issues/174)) ([4ff26cd](https://github.com/googleapis/python-bigquery-dataframes/commit/4ff26cdf862e9f9b91a3a1d2abfa7fbdf0af9c5b)) + ## [0.12.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.11.0...v0.12.0) (2023-11-01) diff --git a/bigframes/version.py b/bigframes/version.py index b324ed72340..0a5df274799 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.12.0" +__version__ = "0.13.0" From a8a409ab0bd1f99dfb442df0703bf8786e0fe58e Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 7 Nov 2023 12:50:46 -0800 Subject: [PATCH 138/197] fix: deprecate the `remote_service_type` in llm model (#180) * fix: deprecate the remote_service_type in llm model * fix: address comments --- bigframes/ml/llm.py | 32 +++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 2e5a9a1e5e5..3cfc28e61f8 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -24,12 +24,14 @@ from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd -_REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" -_REMOTE_TEXT_GENERATOR_32K_MODEL_CODE = "text-bison-32k" +_REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT = "text-bison" +_REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT = "text-bison-32k" _TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" -_REMOTE_EMBEDDING_GENERATOR_MODEL_CODE = "CLOUD_AI_TEXT_EMBEDDING_MODEL_V1" -_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE = "textembedding-gecko-multilingual" +_REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT = "textembedding-gecko" +_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT = ( + "textembedding-gecko-multilingual" +) _EMBED_TEXT_RESULT_COLUMN = "text_embedding" @@ -88,14 +90,18 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - if self.model_name == "text-bison": + if self.model_name == _REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT: options = { - "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, + "endpoint": _REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT, } - else: + elif self.model_name == _REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT: options = { - "endpoint": _REMOTE_TEXT_GENERATOR_32K_MODEL_CODE, + "endpoint": _REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT, } + else: + raise ValueError( + f"Model name {self.model_name} is not supported. We only support {_REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT} and {_REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT}." + ) return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -240,12 +246,16 @@ def _create_bqml_model(self): ) if self.model_name == "textembedding-gecko": options = { - "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, + "endpoint": _REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT, } - else: + elif self.model_name == _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT: options = { - "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE, + "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT, } + else: + raise ValueError( + f"Model name {self.model_name} is not supported. We only support {_REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT} and {_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT}." + ) return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options From 0c8bd33806bb99206b8b12dbdf7d7485c6ffb759 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 7 Nov 2023 15:16:14 -0800 Subject: [PATCH 139/197] feat: support `series.at[row_label] = scalar` (#173) Just quickly pipe the existing loc solution to at --- bigframes/core/indexers.py | 12 ++++++++++++ tests/system/small/test_series.py | 11 +++++++++++ 2 files changed, 23 insertions(+) diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index f6ce084714a..69048b6845a 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -117,6 +117,18 @@ def __getitem__( ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]: return self._series.loc[key] + def __setitem__( + self, + key: LocSingleKey, + value: bigframes.core.scalar.Scalar, + ): + if not pd.api.types.is_scalar(value): + raise NotImplementedError( + "series.at.__setitem__ only supports scalar right-hand values. " + f"{constants.FEEDBACK_LINK}" + ) + self._series.loc[key] = value + class LocDataFrameIndexer: def __init__(self, dataframe: bigframes.dataframe.DataFrame): diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 183ba01c0e0..c8bd0f7afdd 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -1012,6 +1012,17 @@ def test_loc_setitem_cell(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_series_equal(bf_original.to_pandas(), pd_original) +def test_at_setitem_row_label_scalar(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_series = scalars_df["int64_col"] + pd_series = scalars_pandas_df["int64_col"].copy() + bf_series.at[1] = 1000 + pd_series.at[1] = 1000 + bf_result = bf_series.to_pandas() + pd_result = pd_series.astype("Int64") + pd.testing.assert_series_equal(bf_result, pd_result) + + def test_ne_obj_series(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "string_col" From c1d1f42a21cc089877f79ebb46a39ddef6958e04 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 7 Nov 2023 16:00:28 -0800 Subject: [PATCH 140/197] fix: match pandas behavior when assigning listlike to empty dfs (#172) Previously, we weren't able to assign to empty dfs at all, which is certainly not right. This change also matches the behavior for when there is a multi-index but with empty index columns. The pandas behavior is to throw an error with a somewhat odd error message. (I made our error message explicit.) Note: this bug also occurs for assigning series to empty dfs. However, in that case, fixing it would come at the cost of adding a query to check if there are no rows. That's a big cost, so I haven't fixed the series case. --- bigframes/dataframe.py | 49 +++++++++++++++++----------- tests/system/small/test_dataframe.py | 28 ++++++++++++++++ 2 files changed, 58 insertions(+), 19 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 40f12671aef..c5b48822fb1 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1101,23 +1101,38 @@ def _assign_single_item( 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})" - ) + return self._assign_single_item_listlike(k, v) + else: + return self._assign_scalar(k, v) - local_df = bigframes.dataframe.DataFrame( - {k: v}, session=self._get_block().expr.session + def _assign_single_item_listlike(self, k: str, v: Sequence) -> DataFrame: + given_rows = len(v) + actual_rows = len(self) + assigning_to_empty_df = len(self.columns) == 0 and actual_rows == 0 + if not assigning_to_empty_df and given_rows != actual_rows: + raise ValueError( + f"Length of values ({given_rows}) does not match length of index ({actual_rows})" ) - # 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 - new_column_block = local_df._block - original_index_column_ids = self._block.index_columns - self_block = self._block.reset_index(drop=False) - result_index, (get_column_left, get_column_right) = self_block.index.join( + local_df = bigframes.dataframe.DataFrame( + {k: v}, session=self._get_block().expr.session + ) + # local_df is likely (but not guaranteed) to be cached locally + # since the original list came from memory and so is probably < MAX_INLINE_DF_SIZE + + new_column_block = local_df._block + original_index_column_ids = self._block.index_columns + self_block = self._block.reset_index(drop=False) + if assigning_to_empty_df: + if len(self._block.index_columns) > 1: + # match error raised by pandas here + raise ValueError( + "Assigning listlike to a first column under multiindex is not supported." + ) + result_block = new_column_block.with_index_labels(self._block.index_labels) + result_block = result_block.with_column_labels([k]) + else: + result_index, (get_column_left, get_column_right,) = self_block.index.join( new_column_block.index, how="left", block_identity_join=True ) result_block = result_index._block @@ -1125,13 +1140,9 @@ def _assign_single_item( [get_column_left[col_id] for col_id in original_index_column_ids], index_labels=self._block.index_labels, ) - return DataFrame(result_block) - else: - return self._assign_scalar(k, v) + return DataFrame(result_block) 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._block.cols_matching_label(label) block, constant_col_id = self._block.create_constant(value, label) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index bd5930e5089..a2583e08155 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -567,6 +567,34 @@ def test_assign_existing_column(scalars_dfs): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +def test_assign_listlike_to_empty_df(session): + empty_df = dataframe.DataFrame(session=session) + empty_pandas_df = pd.DataFrame() + + bf_result = empty_df.assign(new_col=[1, 2, 3]) + pd_result = empty_pandas_df.assign(new_col=[1, 2, 3]) + + pd_result["new_col"] = pd_result["new_col"].astype("Int64") + pd_result.index = pd_result.index.astype("Int64") + assert_pandas_df_equal_ignore_ordering(bf_result.to_pandas(), pd_result) + + +def test_assign_to_empty_df_multiindex_error(session): + empty_df = dataframe.DataFrame(session=session) + empty_pandas_df = pd.DataFrame() + empty_df["empty_col_1"] = [] + empty_df["empty_col_2"] = [] + empty_pandas_df["empty_col_1"] = [] + empty_pandas_df["empty_col_2"] = [] + empty_df = empty_df.set_index(["empty_col_1", "empty_col_2"]) + empty_pandas_df = empty_pandas_df.set_index(["empty_col_1", "empty_col_2"]) + + with pytest.raises(ValueError): + empty_df.assign(new_col=[1, 2, 3, 4, 5, 6, 7, 8, 9]) + with pytest.raises(ValueError): + empty_pandas_df.assign(new_col=[1, 2, 3, 4, 5, 6, 7, 8, 9]) + + def test_assign_series(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs column_name = "int64_col" From 800d44eb5eb77da5d87b2e005f5a2ed53842e7b5 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 7 Nov 2023 16:44:18 -0800 Subject: [PATCH 141/197] fix: Use anonymous dataset instead of session dataset for temp tables (#181) 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/core/__init__.py | 20 +++++---- bigframes/session/__init__.py | 72 +++++++++++-------------------- bigframes/session/_io/bigquery.py | 8 +++- 3 files changed, 44 insertions(+), 56 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 4653f0ab6aa..866be9c4003 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -32,6 +32,7 @@ import bigframes.dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops +import bigframes.session._io.bigquery if typing.TYPE_CHECKING: from bigframes.session import Session @@ -153,25 +154,28 @@ def start_query( 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.""" - compiled = self.compile() - ibis_expr = compiled._to_ibis_expr("unordered", expose_hidden_cols=True) - destination = self.session._ibis_to_session_table( - ibis_expr, cluster_cols=cluster_cols, api_name="cache" + compiled_value = self.compile() + ibis_expr = compiled_value._to_ibis_expr( + ordering_mode="unordered", expose_hidden_cols=True ) + tmp_table = self.session._ibis_to_session_table( + ibis_expr, cluster_cols=cluster_cols, api_name="cached" + ) + table_expression = self.session.ibis_client.table( - f"{destination.project}.{destination.dataset_id}.{destination.table_id}" + f"{tmp_table.project}.{tmp_table.dataset_id}.{tmp_table.table_id}" ) - new_columns = [table_expression[column] for column in compiled.column_ids] + new_columns = [table_expression[column] for column in compiled_value.column_ids] new_hidden_columns = [ table_expression[column] - for column in compiled._hidden_ordering_column_names + for column in compiled_value._hidden_ordering_column_names ] return ArrayValue.from_ibis( self.session, table_expression, columns=new_columns, hidden_ordering_columns=new_hidden_columns, - ordering=compiled._ordering, + ordering=compiled_value._ordering, ) # Operations diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 2537e81e19c..b49e2469a91 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -16,6 +16,7 @@ from __future__ import annotations +import datetime import logging import os import re @@ -325,9 +326,15 @@ def _query_to_destination( # 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) + cluster_cols = [ + item.name + for item in schema + if (item.name in index_cols) and _can_cluster_bq(item) + ][:_MAX_CLUSTER_COLUMNS] + temp_table = self._create_empty_temp_table(schema, cluster_cols) job_config = bigquery.QueryJobConfig() + job_config.labels["bigframes-api"] = api_name job_config.destination = temp_table try: @@ -422,7 +429,7 @@ def _read_gbq_query( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, - api_name: str, + api_name: str = "read_gbq_query", ) -> dataframe.DataFrame: if isinstance(index_col, str): index_cols = [index_col] @@ -430,9 +437,7 @@ def _read_gbq_query( index_cols = list(index_col) destination, query_job = self._query_to_destination( - query, - index_cols, - api_name=api_name, + query, index_cols, api_name=api_name ) # If there was no destination table, that means the query must have @@ -1273,53 +1278,26 @@ def _create_session_table(self) -> bigquery.TableReference: ) return dataset.table(table_name) - def _create_session_table_empty( + def _create_empty_temp_table( 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) + dataset = self._anonymous_dataset + expiration = ( + datetime.datetime.now(datetime.timezone.utc) + constants.DEFAULT_EXPIRATION + ) - # Use fully-qualified name instead of `_SESSION` name so that the - # created table can be used as the destination table. - return query_job.destination + table = bigframes_io.create_temp_table( + self.bqclient, + dataset, + expiration, + schema=schema, + cluster_columns=cluster_cols, + ) + return bigquery.TableReference.from_string(table) def _create_sequential_ordering( self, @@ -1356,13 +1334,13 @@ def _ibis_to_session_table( cluster_cols: Iterable[str], api_name: str, ) -> bigquery.TableReference: - desination, _ = self._query_to_destination( + destination, _ = self._query_to_destination( self.ibis_client.compile(table), index_cols=list(cluster_cols), api_name=api_name, ) # There should always be a destination table for this query type. - return typing.cast(bigquery.TableReference, desination) + return typing.cast(bigquery.TableReference, destination) def remote_function( self, diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index 06d240fec65..badc91e3565 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -19,7 +19,7 @@ import datetime import textwrap import types -from typing import Dict, Iterable, Union +from typing import Dict, Iterable, Optional, Union import uuid import google.cloud.bigquery as bigquery @@ -121,11 +121,17 @@ def create_temp_table( bqclient: bigquery.Client, dataset: bigquery.DatasetReference, expiration: datetime.datetime, + *, + schema: Optional[Iterable[bigquery.SchemaField]] = None, + cluster_columns: Optional[list[str]] = None, ) -> str: """Create an empty table with an expiration in the desired dataset.""" table_ref = random_table(dataset) destination = bigquery.Table(table_ref) destination.expires = expiration + destination.schema = schema + if cluster_columns: + destination.clustering_fields = cluster_columns bqclient.create_table(destination) return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" From f95000d3f88662be4d88c8b0152f1b838e99ec55 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 7 Nov 2023 17:36:15 -0800 Subject: [PATCH 142/197] fix: for reset_index on unnamed multiindex, always use level_[n] label (#182) 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/core/blocks.py | 2 +- tests/system/small/test_dataframe.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 635e7db8656..e831b42752f 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -282,7 +282,7 @@ def reset_index(self, drop: bool = True) -> Block: column_labels_modified = self.column_labels for level, label in enumerate(index_labels): if label is None: - if "index" not in self.column_labels: + if "index" not in self.column_labels and len(index_labels) <= 1: label = "index" else: label = f"level_{level}" diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a2583e08155..3d7ba867ea4 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1235,6 +1235,28 @@ def test_reset_index_with_unnamed_index( pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_reset_index_with_unnamed_multiindex( + scalars_df_index, + scalars_pandas_df_index, +): + bf_df = dataframe.DataFrame( + ([1, 2, 3], [2, 5, 7]), + index=pd.MultiIndex.from_tuples([("a", "aa"), ("a", "aa")]), + ) + pd_df = pd.DataFrame( + ([1, 2, 3], [2, 5, 7]), + index=pd.MultiIndex.from_tuples([("a", "aa"), ("a", "aa")]), + ) + + bf_df = bf_df.reset_index() + pd_df = pd_df.reset_index() + + assert pd_df.columns[0] == "level_0" + assert bf_df.columns[0] == "level_0" + assert pd_df.columns[1] == "level_1" + assert bf_df.columns[1] == "level_1" + + def test_reset_index_with_unnamed_index_and_index_column( scalars_df_index, scalars_pandas_df_index, From aee0e8e2518c59bd1e0b07940c3309871fde8899 Mon Sep 17 00:00:00 2001 From: Alexander Potts Date: Thu, 9 Nov 2023 03:47:29 +1100 Subject: [PATCH 143/197] feat: support `date_series.astype("string[pyarrow]")` to cast DATE to STRING (#186) --- bigframes/dtypes.py | 2 +- tests/system/small/test_series.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 079f0cc27a2..cd35e380c02 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -375,7 +375,7 @@ def cast_ibis_value( ), ibis_dtypes.float64: (ibis_dtypes.string, ibis_dtypes.int64), ibis_dtypes.string: (ibis_dtypes.int64, ibis_dtypes.float64), - ibis_dtypes.date: (), + ibis_dtypes.date: (ibis_dtypes.string,), ibis_dtypes.Decimal(precision=38, scale=9): (ibis_dtypes.float64,), ibis_dtypes.Decimal(precision=76, scale=38): (ibis_dtypes.float64,), ibis_dtypes.time: (), diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index c8bd0f7afdd..6bb5c4755ca 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2498,6 +2498,7 @@ def test_mask_custom_value(scalars_dfs): # with timezone conversions, so we'll allow it. ("timestamp_col", pd.ArrowDtype(pa.timestamp("us"))), ("datetime_col", pd.ArrowDtype(pa.timestamp("us", tz="UTC"))), + ("date_col", "string[pyarrow]"), # TODO(bmil): fix Ibis bug: BigQuery backend rounds to nearest int # ("float64_col", "Int64"), # TODO(bmil): decide whether to fix Ibis bug: BigQuery backend From 765446a929abe1ac076c3037afa7892f64105356 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 8 Nov 2023 14:28:13 -0800 Subject: [PATCH 144/197] feat: add 'cross' join support (#176) * feat: add 'cross' join support --- bigframes/core/__init__.py | 1 + bigframes/core/blocks.py | 1 + bigframes/core/compile/single_column.py | 3 +- bigframes/core/joins/merge.py | 1 + bigframes/core/nodes.py | 1 + bigframes/dataframe.py | 26 +++++++++++ bigframes/pandas/__init__.py | 1 + tests/system/small/test_dataframe.py | 44 ++++++++++++++----- tests/system/small/test_pandas.py | 21 +++++++++ .../bigframes_vendored/pandas/core/frame.py | 5 +++ .../pandas/core/reshape/merge.py | 2 + 11 files changed, 93 insertions(+), 13 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 866be9c4003..9426df54f07 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -417,6 +417,7 @@ def join( "left", "outer", "right", + "cross", ], allow_row_identity_join: bool = True, ): diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index e831b42752f..bf3b9321c3b 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -1531,6 +1531,7 @@ def merge( "left", "outer", "right", + "cross", ], left_join_ids: typing.Sequence[str], right_join_ids: typing.Sequence[str], diff --git a/bigframes/core/compile/single_column.py b/bigframes/core/compile/single_column.py index b992aa1d1d9..93ba3f16f1c 100644 --- a/bigframes/core/compile/single_column.py +++ b/bigframes/core/compile/single_column.py @@ -40,6 +40,7 @@ def join_by_column( "left", "outer", "right", + "cross", ], allow_row_identity_join: bool = True, ) -> compiled.CompiledArrayValue: @@ -107,7 +108,7 @@ def join_by_column( left_table, right_table, predicates=join_conditions, - how=how, + how=how, # type: ignore ) # Preserve ordering accross joins. diff --git a/bigframes/core/joins/merge.py b/bigframes/core/joins/merge.py index fac16b36078..c65e1bdd54d 100644 --- a/bigframes/core/joins/merge.py +++ b/bigframes/core/joins/merge.py @@ -32,6 +32,7 @@ def merge( "left", "outer", "right", + "cross", ] = "inner", on: Optional[str] = None, *, diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py index 7b252b164f6..8f1e2e5e736 100644 --- a/bigframes/core/nodes.py +++ b/bigframes/core/nodes.py @@ -86,6 +86,7 @@ class JoinNode(BigFrameNode): "left", "outer", "right", + "cross", ] allow_row_identity_join: bool = True diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c5b48822fb1..701c26f42f6 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1933,6 +1933,7 @@ def merge( "left", "outer", "right", + "cross", ] = "inner", # TODO(garrettwu): Currently can take inner, outer, left and right. To support # cross joins @@ -1943,6 +1944,19 @@ def merge( sort: bool = False, suffixes: tuple[str, str] = ("_x", "_y"), ) -> DataFrame: + if how == "cross": + if on is not None: + raise ValueError("'on' is not supported for cross join.") + result_block = self._block.merge( + right._block, + left_join_ids=[], + right_join_ids=[], + suffixes=suffixes, + how=how, + sort=True, + ) + return DataFrame(result_block) + if on is None: if left_on is None or right_on is None: raise ValueError("Must specify `on` or `left_on` + `right_on`.") @@ -1996,6 +2010,18 @@ def join( raise NotImplementedError( f"Deduping column names is not implemented. {constants.FEEDBACK_LINK}" ) + if how == "cross": + if on is not None: + raise ValueError("'on' is not supported for cross join.") + result_block = left._block.merge( + right._block, + left_join_ids=[], + right_join_ids=[], + suffixes=("", ""), + how="cross", + sort=True, + ) + return DataFrame(result_block) # Join left columns with right index if on is not None: diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 1c52b103fbe..d35f8383665 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -332,6 +332,7 @@ def merge( "left", "outer", "right", + "cross", ] = "inner", on: Optional[str] = None, *, diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 3d7ba867ea4..605d4abc1d0 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -914,6 +914,26 @@ def test_df_isin_dict(scalars_dfs): pandas.testing.assert_frame_equal(bf_result, pd_result.astype("boolean")) +def test_df_cross_merge(scalars_dfs): + 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) + + bf_result = left.merge(right, "cross").to_pandas() + + pd_result = scalars_pandas_df[left_columns].merge( + scalars_pandas_df[right_columns].assign( + rowindex_2=scalars_pandas_df["rowindex_2"] + 2 + ), + "cross", + ) + pd.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) + + @pytest.mark.parametrize( ("merge_how",), [ @@ -1745,12 +1765,7 @@ def test_series_binop_add_different_table( all_joins = pytest.mark.parametrize( ("how",), - ( - ("outer",), - ("left",), - ("right",), - ("inner",), - ), + (("outer",), ("left",), ("right",), ("inner",), ("cross",)), ) @@ -1795,13 +1810,18 @@ def test_join_param_on(scalars_dfs, how): 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) + if how == "cross": + with pytest.raises(ValueError): + bf_df_a.join(bf_df_b, on="rowindex_2", how=how) + else: + 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( diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index 0292ebd2069..b88901f3bcf 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -289,6 +289,27 @@ def test_merge_left_on_right_on(scalars_dfs, merge_how): assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) +def test_pd_merge_cross(scalars_dfs): + 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, "cross", sort=True) + bf_result = df.to_pandas() + + pd_result = pd.merge( + scalars_pandas_df[left_columns], + scalars_pandas_df[right_columns], + "cross", + sort=True, + ) + + pd.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) + + @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 6f4f6be35d6..3e4cec284c9 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2194,6 +2194,8 @@ def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame: 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. + ``cross``: creates the cartesian product from both frames, preserves + the order of the left keys. Returns: bigframes.dataframe.DataFrame: A dataframe containing columns from both the caller and `other`. @@ -2208,6 +2210,7 @@ def merge( "left", "outer", "right", + "cross", ] = "inner", on: Optional[str] = None, *, @@ -2243,6 +2246,8 @@ def merge( 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. + ``cross``: creates the cartesian product from both frames, preserves the order + of the left keys. on (label or list of labels): Columns to join on. It must be found in both DataFrames. Either on or left_on + right_on diff --git a/third_party/bigframes_vendored/pandas/core/reshape/merge.py b/third_party/bigframes_vendored/pandas/core/reshape/merge.py index b03f366fca7..704e50f5162 100644 --- a/third_party/bigframes_vendored/pandas/core/reshape/merge.py +++ b/third_party/bigframes_vendored/pandas/core/reshape/merge.py @@ -49,6 +49,8 @@ def merge( 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. + ``cross``: creates the cartesian product from both frames, preserves the order + of the left keys. on (label or list of labels): Columns to join on. It must be found in both DataFrames. Either on or left_on + right_on From 6a2840349a23035bdfdabacd1e231b41bbb5ed7a Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 8 Nov 2023 15:24:55 -0800 Subject: [PATCH 145/197] feat: add 'index', 'pad', 'nearest' interpolate methods (#162) * feat: add 'index', 'pad', 'nearest' interpolate methods --- bigframes/core/block_transforms.py | 207 ++++++++++++++---- bigframes/dataframe.py | 2 + bigframes/series.py | 2 + tests/system/small/test_series.py | 21 +- .../bigframes_vendored/pandas/core/frame.py | 51 +++-- .../bigframes_vendored/pandas/core/series.py | 50 +++-- 6 files changed, 243 insertions(+), 90 deletions(-) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index 917edac0ded..e095f21f6b7 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -22,6 +22,7 @@ import bigframes.core.blocks as blocks import bigframes.core.ordering as ordering import bigframes.core.window_spec as windows +import bigframes.dtypes as dtypes import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops @@ -106,18 +107,33 @@ def indicate_duplicates( def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block: - if method != "linear": + supported_methods = [ + "linear", + "values", + "index", + "nearest", + "zero", + "slinear", + ] + if method not in supported_methods: raise NotImplementedError( - f"Only 'linear' interpolate method supported. {constants.FEEDBACK_LINK}" + f"Method {method} not supported, following interpolate methods supported: {', '.join(supported_methods)}. {constants.FEEDBACK_LINK}" ) - backwards_window = windows.WindowSpec(following=0) - forwards_window = windows.WindowSpec(preceding=0) - output_column_ids = [] original_columns = block.value_columns original_labels = block.column_labels - block, offsets = block.promote_offsets() + + if method == "linear": # Assumes evenly spaced, ignore index + block, xvalues = block.promote_offsets() + else: + index_columns = block.index_columns + if len(index_columns) != 1: + raise ValueError("only method 'linear' supports multi-index") + xvalues = block.index_columns[0] + if block.index_dtypes[0] not in dtypes.NUMERIC_BIGFRAMES_TYPES: + raise ValueError("Can only interpolate on numeric index.") + for column in original_columns: # null in same places column is null should_interpolate = block._column_type(column) in [ @@ -125,48 +141,25 @@ def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block: pd.Int64Dtype(), ] if should_interpolate: - block, notnull = block.apply_unary_op(column, ops.notnull_op) - block, masked_offsets = block.apply_binary_op( - offsets, notnull, ops.partial_arg3(ops.where_op, None) - ) - - block, previous_value = block.apply_window_op( - column, agg_ops.LastNonNullOp(), backwards_window - ) - block, next_value = block.apply_window_op( - column, agg_ops.FirstNonNullOp(), forwards_window - ) - block, previous_value_offset = block.apply_window_op( - masked_offsets, - agg_ops.LastNonNullOp(), - backwards_window, - skip_reproject_unsafe=True, - ) - block, next_value_offset = block.apply_window_op( - masked_offsets, - agg_ops.FirstNonNullOp(), - forwards_window, - skip_reproject_unsafe=True, - ) - - block, prediction_id = _interpolate( + interpolate_method_map = { + "linear": "linear", + "values": "linear", + "index": "linear", + "slinear": "linear", + "zero": "ffill", + "nearest": "nearest", + } + extrapolating_methods = ["linear", "values", "index"] + interpolate_method = interpolate_method_map[method] + do_extrapolate = method in extrapolating_methods + block, interpolated = _interpolate_column( block, - previous_value_offset, - previous_value, - next_value_offset, - next_value, - offsets, + column, + xvalues, + interpolate_method=interpolate_method, + do_extrapolate=do_extrapolate, ) - - block, interpolated_column = block.apply_binary_op( - column, prediction_id, ops.fillna_op - ) - # Pandas performs ffill-like behavior to extrapolate forwards - block, interpolated_and_ffilled = block.apply_binary_op( - interpolated_column, previous_value, ops.fillna_op - ) - - output_column_ids.append(interpolated_and_ffilled) + output_column_ids.append(interpolated) else: output_column_ids.append(column) @@ -175,7 +168,80 @@ def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block: return block.with_column_labels(original_labels) -def _interpolate( +def _interpolate_column( + block: blocks.Block, + column: str, + x_values: str, + interpolate_method: str, + do_extrapolate: bool = True, +) -> typing.Tuple[blocks.Block, str]: + if interpolate_method not in ["linear", "nearest", "ffill"]: + raise ValueError("interpolate method not supported") + window_ordering = (ordering.OrderingColumnReference(x_values),) + backwards_window = windows.WindowSpec(following=0, ordering=window_ordering) + forwards_window = windows.WindowSpec(preceding=0, ordering=window_ordering) + + # Note, this method may + block, notnull = block.apply_unary_op(column, ops.notnull_op) + block, masked_offsets = block.apply_binary_op( + x_values, notnull, ops.partial_arg3(ops.where_op, None) + ) + + block, previous_value = block.apply_window_op( + column, agg_ops.LastNonNullOp(), backwards_window + ) + block, next_value = block.apply_window_op( + column, agg_ops.FirstNonNullOp(), forwards_window + ) + block, previous_value_offset = block.apply_window_op( + masked_offsets, + agg_ops.LastNonNullOp(), + backwards_window, + skip_reproject_unsafe=True, + ) + block, next_value_offset = block.apply_window_op( + masked_offsets, + agg_ops.FirstNonNullOp(), + forwards_window, + skip_reproject_unsafe=True, + ) + + if interpolate_method == "linear": + block, prediction_id = _interpolate_points_linear( + block, + previous_value_offset, + previous_value, + next_value_offset, + next_value, + x_values, + ) + elif interpolate_method == "nearest": + block, prediction_id = _interpolate_points_nearest( + block, + previous_value_offset, + previous_value, + next_value_offset, + next_value, + x_values, + ) + else: # interpolate_method == 'ffill': + block, prediction_id = _interpolate_points_ffill( + block, + previous_value_offset, + previous_value, + next_value_offset, + next_value, + x_values, + ) + if do_extrapolate: + block, prediction_id = block.apply_binary_op( + prediction_id, previous_value, ops.fillna_op + ) + + return block.apply_binary_op(column, prediction_id, ops.fillna_op) + + +def _interpolate_points_linear( block: blocks.Block, x0_id: str, y0_id: str, @@ -196,6 +262,53 @@ def _interpolate( return block, prediction_id +def _interpolate_points_nearest( + block: blocks.Block, + x0_id: str, + y0_id: str, + x1_id: str, + y1_id: str, + xpredict_id: str, +) -> typing.Tuple[blocks.Block, str]: + """Interpolate by taking the y value of the nearest x value""" + block, left_diff = block.apply_binary_op(xpredict_id, x0_id, ops.sub_op) + block, right_diff = block.apply_binary_op(x1_id, xpredict_id, ops.sub_op) + # If diffs equal, choose left + block, choose_left = block.apply_binary_op(left_diff, right_diff, ops.le_op) + block, choose_left = block.apply_unary_op( + choose_left, ops.partial_right(ops.fillna_op, False) + ) + + block, nearest = block.apply_ternary_op(y0_id, choose_left, y1_id, ops.where_op) + + block, y0_exists = block.apply_unary_op(y0_id, ops.notnull_op) + block, y1_exists = block.apply_unary_op(y1_id, ops.notnull_op) + block, is_interpolation = block.apply_binary_op(y0_exists, y1_exists, ops.and_op) + + block, prediction_id = block.apply_binary_op( + nearest, is_interpolation, ops.partial_arg3(ops.where_op, None) + ) + + return block, prediction_id + + +def _interpolate_points_ffill( + block: blocks.Block, + x0_id: str, + y0_id: str, + x1_id: str, + y1_id: str, + xpredict_id: str, +) -> typing.Tuple[blocks.Block, str]: + """Interpolates by using the preceding values""" + # check for existance of y1, otherwise we are extrapolating instead of interpolating + block, y1_exists = block.apply_unary_op(y1_id, ops.notnull_op) + block, prediction_id = block.apply_binary_op( + y0_id, y1_exists, ops.partial_arg3(ops.where_op, None) + ) + return block, prediction_id + + def drop_duplicates( block: blocks.Block, columns: typing.Sequence[str], keep: str = "first" ) -> blocks.Block: diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 701c26f42f6..1e143144fe3 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1450,6 +1450,8 @@ def reindex_like(self, other: DataFrame, *, validate: typing.Optional[bool] = No return self.reindex(index=other.index, columns=other.columns, validate=validate) def interpolate(self, method: str = "linear") -> DataFrame: + if method == "pad": + return self.ffill() result = block_ops.interpolate(self._block, method) return DataFrame(result) diff --git a/bigframes/series.py b/bigframes/series.py index 032bdf6c429..9e111618de2 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -475,6 +475,8 @@ def replace( return Series(block.select_column(result_col)) def interpolate(self, method: str = "linear") -> Series: + if method == "pad": + return self.ffill() result = block_ops.interpolate(self._block, method) return Series(result) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 6bb5c4755ca..006d370818d 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -273,21 +273,26 @@ def test_series_replace_list_scalar(scalars_dfs): @pytest.mark.parametrize( - ("values",), + ("method",), ( - ([None, 1, 2, None, None, 16, None],), - ([None, None, 3.6, None],), - ([403.2, None, 352.1, None, None, 111.9],), + ("linear",), + ("values",), + ("slinear",), + ("nearest",), + ("zero",), + ("pad",), ), ) -def test_series_interpolate(values): - pd_series = pd.Series(values) +def test_series_interpolate(method): + values = [None, 1, 2, None, None, 16, None] + index = [-3.2, 11.4, 3.56, 4, 4.32, 5.55, 76.8] + pd_series = pd.Series(values, index) bf_series = series.Series(pd_series) # Pandas can only interpolate on "float64" columns # https://github.com/pandas-dev/pandas/issues/40252 - pd_result = pd_series.astype("float64").interpolate() - bf_result = bf_series.interpolate().to_pandas() + pd_result = pd_series.astype("float64").interpolate(method=method) + bf_result = bf_series.interpolate(method=method).to_pandas() # pd uses non-null types, while bf uses nullable types pd.testing.assert_series_equal( diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 3e4cec284c9..10f73d25b77 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2872,17 +2872,6 @@ def interpolate(self, method: str = "linear"): """ Fill NaN values using an interpolation method. - Args: - method (str, default 'linear'): - Interpolation technique to use. Only 'linear' supported. - 'linear': Ignore the index and treat the values as equally spaced. - This is the only method supported on MultiIndexes. - - Returns: - DataFrame: - Returns the same object type as the caller, interpolated at - some or all ``NaN`` values - **Examples:** >>> import bigframes.pandas as bpd @@ -2891,17 +2880,41 @@ def interpolate(self, method: str = "linear"): >>> df = bpd.DataFrame({ ... 'A': [1, 2, 3, None, None, 6], ... 'B': [None, 6, None, 2, None, 3], - ... }) + ... }, index=[0, 0.1, 0.3, 0.7, 0.9, 1.0]) >>> df.interpolate() - A B - 0 1.0 - 1 2.0 6.0 - 2 3.0 4.0 - 3 4.0 2.0 - 4 5.0 2.5 - 5 6.0 3.0 + A B + 0.0 1.0 + 0.1 2.0 6.0 + 0.3 3.0 4.0 + 0.7 4.0 2.0 + 0.9 5.0 2.5 + 1.0 6.0 3.0 [6 rows x 2 columns] + >>> df.interpolate(method="values") + A B + 0.0 1.0 + 0.1 2.0 6.0 + 0.3 3.0 4.666667 + 0.7 4.714286 2.0 + 0.9 5.571429 2.666667 + 1.0 6.0 3.0 + + [6 rows x 2 columns] + + Args: + method (str, default 'linear'): + Interpolation technique to use. Only 'linear' supported. + 'linear': Ignore the index and treat the values as equally spaced. + This is the only method supported on MultiIndexes. + 'index', 'values': use the actual numerical values of the index. + 'pad': Fill in NaNs using existing values. + 'nearest', 'zero', 'slinear': Emulates `scipy.interpolate.interp1d` + + Returns: + DataFrame: + Returns the same object type as the caller, interpolated at + some or all ``NaN`` values """ 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 b569e5699c7..b25b0c75cff 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -920,31 +920,49 @@ def interpolate(self, method: str = "linear"): """ Fill NaN values using an interpolation method. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({ + ... 'A': [1, 2, 3, None, None, 6], + ... 'B': [None, 6, None, 2, None, 3], + ... }, index=[0, 0.1, 0.3, 0.7, 0.9, 1.0]) + >>> df.interpolate() + A B + 0.0 1.0 + 0.1 2.0 6.0 + 0.3 3.0 4.0 + 0.7 4.0 2.0 + 0.9 5.0 2.5 + 1.0 6.0 3.0 + + [6 rows x 2 columns] + >>> df.interpolate(method="values") + A B + 0.0 1.0 + 0.1 2.0 6.0 + 0.3 3.0 4.666667 + 0.7 4.714286 2.0 + 0.9 5.571429 2.666667 + 1.0 6.0 3.0 + + [6 rows x 2 columns] + + Args: method (str, default 'linear'): Interpolation technique to use. Only 'linear' supported. 'linear': Ignore the index and treat the values as equally spaced. This is the only method supported on MultiIndexes. - + 'index', 'values': use the actual numerical values of the index. + 'pad': Fill in NaNs using existing values. + 'nearest', 'zero', 'slinear': Emulates `scipy.interpolate.interp1d` Returns: Series: Returns the same object type as the caller, interpolated at some or all ``NaN`` values - - **Examples:** - - >>> import bigframes.pandas as bpd - >>> bpd.options.display.progress_bar = None - - >>> series = bpd.Series([1, 2, 3, None, None, 6]) - >>> series.interpolate() - 0 1.0 - 1 2.0 - 2 3.0 - 3 4.0 - 4 5.0 - 5 6.0 - dtype: Float64 """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 58f420c91d94ca085e9810f36513ffe772bfddcf Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 8 Nov 2023 16:04:14 -0800 Subject: [PATCH 146/197] feat: add unordered sql compilation (#156) 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/core/__init__.py | 33 +- bigframes/core/blocks.py | 9 +- bigframes/core/compile/__init__.py | 10 +- bigframes/core/compile/compiled.py | 1504 ++++++++++------- bigframes/core/compile/compiler.py | 173 +- bigframes/core/compile/concat.py | 100 ++ bigframes/core/compile/row_identity.py | 70 +- bigframes/core/compile/single_column.py | 104 +- bigframes/dataframe.py | 6 + bigframes/series.py | 7 + bigframes/session/__init__.py | 6 +- tests/system/large/ml/test_cluster.py | 4 +- tests/system/large/ml/test_pipeline.py | 4 +- tests/system/large/test_remote_function.py | 32 +- tests/system/small/ml/test_cluster.py | 4 +- tests/system/small/ml/test_core.py | 3 +- tests/system/small/ml/test_decomposition.py | 6 +- .../system/small/operations/test_datetimes.py | 22 +- tests/system/small/operations/test_strings.py | 42 +- tests/system/small/test_dataframe.py | 183 +- tests/system/small/test_dataframe_io.py | 17 +- tests/system/small/test_groupby.py | 35 +- tests/system/small/test_multiindex.py | 6 +- tests/system/small/test_pandas.py | 21 +- tests/system/small/test_remote_function.py | 24 +- tests/system/small/test_series.py | 111 +- tests/system/utils.py | 42 +- tests/unit/test_core.py | 32 +- 28 files changed, 1652 insertions(+), 958 deletions(-) create mode 100644 bigframes/core/compile/concat.py diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 9426df54f07..63f36d4ddd2 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -23,7 +23,8 @@ import ibis.expr.types as ibis_types import pandas -import bigframes.core.compile as compiled +import bigframes.core.compile.compiled as compiled +import bigframes.core.compile.compiler as compiler import bigframes.core.guid import bigframes.core.nodes as nodes from bigframes.core.ordering import OrderingColumnReference @@ -78,7 +79,7 @@ def from_pandas(cls, pd_df: pandas.DataFrame): @property def column_ids(self) -> typing.Sequence[str]: - return self.compile().column_ids + return self._compile_ordered().column_ids @property def session(self) -> Session: @@ -88,15 +89,18 @@ def session(self) -> Session: return self.node.session[0] if required_session else get_global_session() def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: - return self.compile().get_column_type(key) + return self._compile_ordered().get_column_type(key) - def compile(self) -> compiled.CompiledArrayValue: - return compiled.compile_node(self.node) + def _compile_ordered(self) -> compiled.OrderedIR: + return compiler.compile_ordered(self.node) + + def _compile_unordered(self) -> compiled.UnorderedIR: + return compiler.compile_unordered(self.node) def shape(self) -> typing.Tuple[int, int]: """Returns dimensions as (length, width) tuple.""" - width = len(self.compile().columns) - count_expr = self.compile()._to_ibis_expr("unordered").count() + width = len(self._compile_unordered().columns) + count_expr = self._compile_unordered()._to_ibis_expr().count() # Support in-memory engines for hermetic unit tests. if not self.node.session: @@ -121,11 +125,14 @@ def to_sql( col_id_overrides: typing.Mapping[str, str] = {}, sorted: bool = False, ) -> str: - return self.compile().to_sql( - offset_column=offset_column, - col_id_overrides=col_id_overrides, - sorted=sorted, - ) + if sorted or offset_column: + return self._compile_ordered().to_sql( + offset_column=offset_column, + col_id_overrides=col_id_overrides, + sorted=sorted, + ) + else: + return self._compile_unordered().to_sql(col_id_overrides=col_id_overrides) def start_query( self, @@ -154,7 +161,7 @@ def start_query( 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.""" - compiled_value = self.compile() + compiled_value = self._compile_ordered() ibis_expr = compiled_value._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True ) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index bf3b9321c3b..6358d28e2ea 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -386,6 +386,8 @@ def to_pandas( max_download_size: Optional[int] = None, sampling_method: Optional[str] = None, random_state: Optional[int] = None, + *, + ordered: bool = True, ) -> Tuple[pd.DataFrame, bigquery.QueryJob]: """Run query and download results as a pandas DataFrame.""" if max_download_size is None: @@ -412,6 +414,7 @@ def to_pandas( max_download_size=max_download_size, sampling_method=sampling_method, random_state=random_state, + ordered=ordered, ) return df, query_job @@ -446,12 +449,16 @@ def _compute_and_count( max_download_size: Optional[int] = None, sampling_method: Optional[str] = None, random_state: Optional[int] = None, + *, + ordered: bool = True, ) -> 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._apply_value_keys_to_expr(value_keys=value_keys) - results_iterator, query_job = expr.start_query(max_results=max_results) + results_iterator, query_job = expr.start_query( + max_results=max_results, sorted=ordered + ) table_size = ( expr.session._get_table_size(query_job.destination) / _BYTES_TO_MEGABYTES diff --git a/bigframes/core/compile/__init__.py b/bigframes/core/compile/__init__.py index c86f4463dc0..761fd9a4650 100644 --- a/bigframes/core/compile/__init__.py +++ b/bigframes/core/compile/__init__.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from bigframes.core.compile.compiled import CompiledArrayValue -from bigframes.core.compile.compiler import compile_node +from bigframes.core.compile.compiled import OrderedIR, UnorderedIR +from bigframes.core.compile.compiler import compile_ordered, compile_unordered __all__ = [ - "compile_node", - "CompiledArrayValue", + "compile_ordered", + "compile_unordered", + "OrderedIR", + "UnorderedIR", ] diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index 1134f1aab01..4ba5e6bd080 100644 --- a/bigframes/core/compile/compiled.py +++ b/bigframes/core/compile/compiled.py @@ -13,8 +13,8 @@ # limitations under the License. from __future__ import annotations +import abc import functools -import math import textwrap import typing from typing import Collection, Iterable, Literal, Optional, Sequence @@ -32,8 +32,6 @@ ExpressionOrdering, IntegerEncoding, OrderingColumnReference, - reencode_order_string, - StringEncoding, ) import bigframes.core.utils as utils from bigframes.core.window_spec import WindowSpec @@ -44,8 +42,568 @@ ORDER_ID_COLUMN = "bigframes_ordering_id" PREDICATE_COLUMN = "bigframes_predicate" +T = typing.TypeVar("T", bound="BaseIbisIR") -class CompiledArrayValue: + +class BaseIbisIR(abc.ABC): + """Implementation detail, contains common logic between ordered and unordered IR""" + + def __init__( + self, + table: ibis_types.Table, + columns: Sequence[ibis_types.Value], + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + self._table = table + self._predicates = tuple(predicates) if predicates is not None else () + # Allow creating a DataFrame directly from an Ibis table expression. + # TODO(swast): Validate that each column references the same table (or + # no table for literal values). + self._columns = tuple(columns) + # 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} + + @property + def columns(self) -> typing.Tuple[ibis_types.Value, ...]: + return self._columns + + @property + def column_ids(self) -> typing.Sequence[str]: + return tuple(self._column_names.keys()) + + @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 + ) + + @abc.abstractmethod + def select_columns(self: T, column_ids: typing.Sequence[str]) -> T: + """Creates a new expression based on this expression with new columns.""" + ... + + def drop_columns(self: T, columns: Iterable[str]) -> T: + return self.select_columns( + [col for col in self.column_ids if col not in columns] + ) + + @abc.abstractmethod + def filter(self: T, predicate_id: str, keep_null: bool = False) -> T: + """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" + ... + + @abc.abstractmethod + def unpivot( + self: T, + row_labels: typing.Sequence[typing.Hashable], + unpivot_columns: typing.Sequence[ + typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + ], + *, + passthrough_columns: typing.Sequence[str] = (), + index_col_ids: typing.Sequence[str] = ["index"], + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + ] = pandas.Float64Dtype(), + how="left", + ) -> T: + """ + 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 + """ + ... + + @abc.abstractmethod + def _reproject_to_table(self: T) -> T: + """ + 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. + """ + ... + + def project_unary_op( + self: T, + input_column_id: str, + op: ops.UnaryOp, + output_column_id: typing.Optional[str] = None, + ) -> T: + """Creates a new expression based on this expression with unary operation applied to one column.""" + result_id = ( + output_column_id or input_column_id + ) # overwrite input if not output id provided + value = op._as_ibis(self._get_ibis_column(input_column_id)).name(result_id) + return self._set_or_replace_by_id(result_id, value) + + def project_binary_op( + self: T, + left_column_id: str, + right_column_id: str, + op: ops.BinaryOp, + output_column_id: str, + ) -> T: + """Creates a new expression based on this expression with binary operation applied to two columns.""" + value = op( + 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) + + def project_ternary_op( + self: T, + col_id_1: str, + col_id_2: str, + col_id_3: str, + op: ops.TernaryOp, + output_column_id: str, + ) -> T: + """Creates a new expression based on this expression with ternary operation applied to three columns.""" + value = op( + 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) + + def assign(self: T, source_id: str, destination_id: str) -> T: + return self._set_or_replace_by_id( + destination_id, self._get_ibis_column(source_id) + ) + + def assign_constant( + self: T, + destination_id: str, + value: typing.Any, + dtype: typing.Optional[bigframes.dtypes.Dtype], + ) -> T: + # 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)}. {constants.FEEDBACK_LINK}" + ) + expr = self._set_or_replace_by_id(destination_id, ibis_value) + return expr._reproject_to_table() + + @abc.abstractmethod + def _set_or_replace_by_id(self: T, id: str, new_value: ibis_types.Value) -> T: + ... + + def _get_ibis_column(self, key: str) -> ibis_types.Value: + """Gets the Ibis expression for a given column.""" + if key not in self.column_ids: + raise ValueError( + "Column name {} not in set of values: {}".format(key, self.column_ids) + ) + return typing.cast(ibis_types.Value, self._column_names[key]) + + def get_column_type(self, key: str) -> bigframes.dtypes.Dtype: + ibis_type = typing.cast( + bigframes.dtypes.IbisDtype, self._get_ibis_column(key).type() + ) + return typing.cast( + bigframes.dtypes.Dtype, + bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), + ) + + +# Ibis Implementations +class UnorderedIR(BaseIbisIR): + def __init__( + self, + table: ibis_types.Table, + columns: Sequence[ibis_types.Value], + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + super().__init__(table, columns, predicates) + + def builder(self): + """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 UnorderedIR.Builder( + self._table, + columns=self._columns, + predicates=self._predicates, + ) + + def to_sql( + self, + offset_column: typing.Optional[str] = None, + col_id_overrides: typing.Mapping[str, str] = {}, + sorted: bool = False, + ) -> str: + if offset_column or sorted: + raise ValueError("Cannot produce sorted sql in unordered mode") + sql = ibis_bigquery.Backend().compile( + self._to_ibis_expr( + col_id_overrides=col_id_overrides, + ) + ) + return typing.cast(str, sql) + + def _to_ibis_expr( + self, + *, + expose_hidden_cols: bool = False, + fraction: Optional[float] = None, + col_id_overrides: typing.Mapping[str, str] = {}, + ): + """ + 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. + + * "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. + * "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: + expose_hidden_cols: + 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. + """ + columns = list(self._columns) + 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) + # 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()) + + # Special case for empty tables, since we can't create an empty + # projection. + if not columns: + return ibis.memtable([]) + + # Make sure all dtypes are the "canonical" ones for BigFrames. This is + # important for operations like UNION where the schema must match. + 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(base_table[PREDICATE_COLUMN]) + 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 select_columns(self, column_ids: typing.Sequence[str]) -> UnorderedIR: + """Creates a new expression based on this expression with new columns.""" + columns = [self._get_ibis_column(col_id) for col_id in column_ids] + builder = self.builder() + builder.columns = list(columns) + new_expr = builder.build() + return new_expr + + def filter(self, predicate_id: str, keep_null: bool = False) -> UnorderedIR: + condition = typing.cast( + ibis_types.BooleanValue, self._get_ibis_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) -> UnorderedIR: + """Filter the table on a given expression, the predicate must be a boolean series aligned with the table expression.""" + expr = self.builder() + expr.predicates = [*self._predicates, predicate_value] + return expr.build() + + def unpivot( + self, + row_labels: typing.Sequence[typing.Hashable], + unpivot_columns: typing.Sequence[ + typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + ], + *, + passthrough_columns: typing.Sequence[str] = (), + index_col_ids: typing.Sequence[str] = ["index"], + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + ] = pandas.Float64Dtype(), + how="left", + ) -> UnorderedIR: + if how not in ("left", "right"): + raise ValueError("'how' must be 'left' or 'right'") + table = self._to_ibis_expr() + 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.") + + unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") + unpivot_table = table.cross_join( + ibis.memtable({unpivot_offset_id: range(row_n)}) + ) + # Use ibis memtable to infer type of rowlabels (if possible) + # TODO: Allow caller to specify dtype + 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 + ] + + 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)): + 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 + ) + unpivot_values.append(unpivot_value.name(result_col)) + + unpivot_table = unpivot_table.select( + passthrough_columns, + *label_columns, + *unpivot_values, + 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] + return UnorderedIR( + table=unpivot_table, + columns=[ + *[unpivot_table[col_id] for col_id in index_col_ids], + *value_columns, + *passthrough_values, + ], + ) + + def aggregate( + self, + aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp, str]], + by_column_ids: typing.Sequence[str] = (), + dropna: bool = True, + ) -> OrderedIR: + """ + Apply aggregations to the expression. + Arguments: + 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() + 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( + tuple( + [ + OrderingColumnReference(column_id=column_id) + for column_id in by_column_ids + ] + ), + total_ordering_columns=frozenset(by_column_ids), + ) + columns = tuple(result[key] for key in result.columns) + expr = OrderedIR(result, columns=columns, ordering=ordering) + if dropna: + for column_id in by_column_ids: + expr = expr._filter( + 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() + 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_value_columns=tuple( + [OrderingColumnReference(ORDER_ID_COLUMN)] + ), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ) + return OrderedIR( + result, + columns=[result[col_id] for col_id in [*stats.keys()]], + hidden_ordering_columns=[result[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def corr_aggregate( + self, corr_aggregations: typing.Sequence[typing.Tuple[str, str, str]] + ) -> OrderedIR: + """ + 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() + 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=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ) + return OrderedIR( + result, + columns=[result[col_id] for col_id in [*stats.keys()]], + hidden_ordering_columns=[result[ORDER_ID_COLUMN]], + ordering=ordering, + ) + + def _uniform_sampling(self, fraction: float) -> UnorderedIR: + """Sampling the table on given fraction. + + .. warning:: + The row numbers of result is non-deterministic, avoid to use. + """ + table = self._to_ibis_expr(fraction=fraction) + columns = [table[column_name] for column_name in self._column_names] + return UnorderedIR( + table, + columns=columns, + ) + + ## Helpers + def _set_or_replace_by_id( + self, id: str, new_value: ibis_types.Value + ) -> UnorderedIR: + builder = self.builder() + if id in self.column_ids: + builder.columns = [ + val if (col_id != id) else new_value.name(id) + for col_id, val in zip(self.column_ids, self._columns) + ] + else: + builder.columns = [*self.columns, new_value.name(id)] + return builder.build() + + def _reproject_to_table(self) -> UnorderedIR: + """ + 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() + columns = [table[column_name] for column_name in self._column_names] + return UnorderedIR( + table, + columns=columns, + ) + + class Builder: + def __init__( + self, + table: ibis_types.Table, + columns: Collection[ibis_types.Value] = (), + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, + ): + self.table = table + self.columns = list(columns) + self.predicates = list(predicates) if predicates is not None else None + + def build(self) -> UnorderedIR: + return UnorderedIR( + table=self.table, + columns=self.columns, + predicates=self.predicates, + ) + + +class OrderedIR(BaseIbisIR): """Immutable BigQuery DataFrames expression tree. Note: Usage of this class is considered to be private and subject to change @@ -71,17 +629,11 @@ def __init__( ordering: ExpressionOrdering = ExpressionOrdering(), predicates: Optional[Collection[ibis_types.BooleanValue]] = None, ): - self._table = table - self._predicates = tuple(predicates) if predicates is not None else () + super().__init__(table, columns, predicates) # TODO: Validate ordering 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. - # 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) @@ -111,10 +663,10 @@ def __init__( raise ValueError(f"Illegal ordering keys: {ordering.all_ordering_columns}") @classmethod - def mem_expr_from_pandas( + def from_pandas( cls, pd_df: pandas.DataFrame, - ) -> CompiledArrayValue: + ) -> OrderedIR: """ Builds an in-memory only (SQL only) expr from a pandas dataframe. """ @@ -173,27 +725,10 @@ def mem_expr_from_pandas( hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],), ) - @property - def columns(self) -> typing.Tuple[ibis_types.Value, ...]: - return self._columns - - @property - def column_ids(self) -> typing.Sequence[str]: - return tuple(self._column_names.keys()) - @property 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]: - """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.""" @@ -202,12 +737,15 @@ def _ibis_order(self) -> Sequence[ibis_types.Value]: self._ordering.all_ordering_columns, ) - def builder(self) -> ArrayValueBuilder: + def to_unordered(self) -> UnorderedIR: + return UnorderedIR(self._table, self._columns, self._predicates) + + def builder(self) -> OrderedIR.Builder: """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( + return OrderedIR.Builder( self._table, columns=self._columns, hidden_ordering_columns=self._hidden_ordering_columns, @@ -215,160 +753,39 @@ def builder(self) -> ArrayValueBuilder: predicates=self._predicates, ) - def drop_columns(self, columns: Iterable[str]) -> CompiledArrayValue: - # 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_any_column(key).type() - ) - return typing.cast( - bigframes.dtypes.Dtype, - bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type), - ) - - def _get_ibis_column(self, key: str) -> ibis_types.Value: - """Gets the Ibis expression for a given column.""" - if key not in self.column_ids: - raise ValueError( - "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: - """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 filter(self, predicate_id: str, keep_null: bool = False) -> CompiledArrayValue: - """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_ibis_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) -> CompiledArrayValue: - """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_value] - return expr.build() - def order_by( self, by: Sequence[OrderingColumnReference], stable: bool = False - ) -> CompiledArrayValue: + ) -> OrderedIR: expr_builder = self.builder() expr_builder.ordering = self._ordering.with_ordering_columns(by, stable=stable) return expr_builder.build() - def reversed(self) -> CompiledArrayValue: + def reversed(self) -> OrderedIR: expr_builder = self.builder() expr_builder.ordering = self._ordering.with_reverse() return expr_builder.build() - def _uniform_sampling(self, fraction: float) -> CompiledArrayValue: + def _uniform_sampling(self, fraction: float) -> OrderedIR: """Sampling the table on given fraction. .. warning:: The row numbers of result is non-deterministic, avoid to use. """ table = self._to_ibis_expr( - "unordered", expose_hidden_cols=True, fraction=fraction + ordering_mode="unordered", 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 CompiledArrayValue( + return OrderedIR( table, columns=columns, hidden_ordering_columns=hidden_ordering_columns, ordering=self._ordering, ) - @property - 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." - ) - if not self._ordering.total_order_col: - raise ValueError( - "Ordering is invalid. Marked as sequential but no total order columns." - ) - column = self._get_any_column(self._ordering.total_order_col.column_id) - return typing.cast(ibis_types.IntegerColumn, column) - - def _project_offsets(self) -> CompiledArrayValue: - """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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - integer_encoding=IntegerEncoding(True, is_sequential=True), - ) - return CompiledArrayValue( - table, - columns=columns, - hidden_ordering_columns=[table[ORDER_ID_COLUMN]], - ordering=ordering, - ) - - def _hide_column(self, column_id) -> CompiledArrayValue: - """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_ibis_column(column_id).name(new_name), - ] - expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) - return expr_builder.build() - - def promote_offsets(self, col_id: str) -> CompiledArrayValue: + def promote_offsets(self, col_id: str) -> OrderedIR: """ Convenience function to promote copy of column offsets to a value column. Can be used to reset index. """ @@ -384,194 +801,21 @@ def promote_offsets(self, col_id: str) -> CompiledArrayValue: ] return expr_builder.build() - def select_columns(self, column_ids: typing.Sequence[str]) -> CompiledArrayValue: + def select_columns(self, column_ids: typing.Sequence[str]) -> OrderedIR: """Creates a new expression based on this expression with new columns.""" columns = [self._get_ibis_column(col_id) for col_id in column_ids] expr = self 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 - expr = expr._hide_column(ordering_column) - builder = expr.builder() - builder.columns = list(columns) - new_expr = builder.build() - return new_expr - - def concat(self, other: typing.Sequence[CompiledArrayValue]) -> CompiledArrayValue: - """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.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="string_encoded", 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 - + reencode_order_string( - 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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - string_encoding=StringEncoding(True, prefix_size + max_encoding_size), - ) - return CompiledArrayValue( - 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 - ) -> CompiledArrayValue: - """Creates a new expression based on this expression with unary operation applied to one column.""" - 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) - - def project_binary_op( - self, - left_column_id: str, - right_column_id: str, - op: ops.BinaryOp, - output_column_id: str, - ) -> CompiledArrayValue: - """Creates a new expression based on this expression with binary operation applied to two columns.""" - value = op( - 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) - - def project_ternary_op( - self, - col_id_1: str, - col_id_2: str, - col_id_3: str, - op: ops.TernaryOp, - output_column_id: str, - ) -> CompiledArrayValue: - """Creates a new expression based on this expression with ternary operation applied to three columns.""" - value = op( - 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) - - def aggregate( - self, - aggregations: typing.Sequence[typing.Tuple[str, agg_ops.AggregateOp, str]], - by_column_ids: typing.Sequence[str] = (), - dropna: bool = True, - ) -> CompiledArrayValue: - """ - Apply aggregations to the expression. - Arguments: - 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("unordered") - 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( - tuple( - [ - OrderingColumnReference(column_id=column_id) - for column_id in by_column_ids - ] - ), - total_ordering_columns=frozenset(by_column_ids), - ) - columns = tuple(result[key] for key in result.columns) - expr = CompiledArrayValue(result, columns=columns, ordering=ordering) - if dropna: - for column_id in by_column_ids: - expr = expr._filter( - 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() - 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_value_columns=tuple( - [OrderingColumnReference(ORDER_ID_COLUMN)] - ), - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), - ) - return CompiledArrayValue( - result, - columns=[result[col_id] for col_id in [*stats.keys()]], - hidden_ordering_columns=[result[ORDER_ID_COLUMN]], - ordering=ordering, - ) - - def corr_aggregate( - self, corr_aggregations: typing.Sequence[typing.Tuple[str, str, str]] - ) -> CompiledArrayValue: - """ - 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("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=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), - total_ordering_columns=frozenset([ORDER_ID_COLUMN]), - integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), - ) - return CompiledArrayValue( - result, - columns=[result[col_id] for col_id in [*stats.keys()]], - hidden_ordering_columns=[result[ORDER_ID_COLUMN]], - ordering=ordering, - ) + ): + # 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 + ## Methods that only work with ordering def project_window_op( self, column_name: str, @@ -581,7 +825,7 @@ def project_window_op( *, never_skip_nulls=False, skip_reproject_unsafe: bool = False, - ) -> CompiledArrayValue: + ) -> OrderedIR: """ 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 @@ -625,6 +869,168 @@ 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 unpivot( + self, + row_labels: typing.Sequence[typing.Hashable], + unpivot_columns: typing.Sequence[ + typing.Tuple[str, typing.Sequence[typing.Optional[str]]] + ], + *, + passthrough_columns: typing.Sequence[str] = (), + index_col_ids: typing.Sequence[str] = ["index"], + dtype: typing.Union[ + bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] + ] = pandas.Float64Dtype(), + how="left", + ) -> OrderedIR: + 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 + ): + raise ValueError("Columns and row labels must all be same length.") + + unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") + unpivot_table = table.cross_join( + ibis.memtable({unpivot_offset_id: range(row_n)}) + ) + # Use ibis memtable to infer type of rowlabels (if possible) + # TODO: Allow caller to specify dtype + 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 + ] + + 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)): + 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 + ) + unpivot_values.append(unpivot_value.name(result_col)) + + unpivot_table = unpivot_table.select( + 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=tuple( + [ + *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=tuple( + [ + 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 OrderedIR( + table=unpivot_table, + 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 _reproject_to_table(self) -> OrderedIR: + table = self._to_ibis_expr( + ordering_mode="unordered", + 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 + if column_name in ordering_col_ids + ] + return OrderedIR( + table, + columns=columns, + hidden_ordering_columns=hidden_ordering_columns, + ordering=self._ordering, + ) + def to_sql( self, offset_column: typing.Optional[str] = None, @@ -644,21 +1050,22 @@ def to_sql( ) if sorted: sql = textwrap.dedent( - f""" - SELECT * EXCEPT (`{offsets_id}`) - FROM ({sql}) - ORDER BY `{offsets_id}` - """ + f"SELECT * EXCEPT (`{offsets_id}`)\n" + "FROM (\n" + f"{sql}\n" + ")\n" + f"ORDER BY `{offsets_id}`\n" ) return typing.cast(str, sql) def _to_ibis_expr( self, - 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, col_id_overrides: typing.Mapping[str, str] = {}, + ordering_mode: Literal["string_encoded", "offset_col", "unordered"], + order_col_name: Optional[str] = ORDER_ID_COLUMN, ): """ Creates an Ibis table expression representing the DataFrame. @@ -677,16 +1084,16 @@ def _to_ibis_expr( column name will be 'bigframes_ordering_id' Args: + expose_hidden_cols: + If True, include the hidden ordering columns in the results. + Only compatible with `order_by` and `unordered` + ``ordering_mode``. 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``. col_id_overrides: overrides the column ids for the result Returns: @@ -723,20 +1130,115 @@ def _to_ibis_expr( if not columns: return ibis.memtable([]) - # Make sure all dtypes are the "canonical" ones for BigFrames. This is - # important for operations like UNION where the schema must match. - table = self._table.select( - bigframes.dtypes.ibis_value_to_canonical_type(column) for column in columns + # Make sure all dtypes are the "canonical" ones for BigFrames. This is + # important for operations like UNION where the schema must match. + 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(base_table[PREDICATE_COLUMN]) + 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 filter(self, predicate_id: str, keep_null: bool = False) -> OrderedIR: + condition = typing.cast( + ibis_types.BooleanValue, self._get_ibis_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) -> OrderedIR: + """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_value] + return expr.build() + + def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value) -> OrderedIR: + """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_ids: + builder.columns = [ + val if (col_id != id) else new_value.name(id) + for col_id, val in zip(self.column_ids, self._columns) + ] + else: + builder.columns = [*self.columns, new_value.name(id)] + return builder.build() + + ## Ordering specific helpers + 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 _hide_column(self, column_id) -> OrderedIR: + """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_ibis_column(column_id).name(new_name), + ] + expr_builder.ordering = self._ordering.with_column_remap({column_id: new_name}) + return expr_builder.build() + + def _project_offsets(self) -> OrderedIR: + """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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + integer_encoding=IntegerEncoding(True, is_sequential=True), + ) + return OrderedIR( + table, + columns=columns, + hidden_ordering_columns=[table[ORDER_ID_COLUMN]], + ordering=ordering, ) - base_table = table - if self._reduced_predicate is not None: - table = table.filter(base_table[PREDICATE_COLUMN]) - 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, @@ -789,34 +1291,6 @@ def _create_string_ordering_column(self) -> ibis_types.StringColumn: ) return encode_order_string(row_nums) - def _reproject_to_table(self) -> CompiledArrayValue: - """ - 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( - "unordered", - 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 - if column_name in ordering_col_ids - ] - return CompiledArrayValue( - 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] = ( [ @@ -851,229 +1325,29 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal group_by=group_by, ) - def unpivot( - self, - row_labels: typing.Sequence[typing.Hashable], - unpivot_columns: typing.Sequence[ - typing.Tuple[str, typing.Sequence[typing.Optional[str]]] - ], - *, - passthrough_columns: typing.Sequence[str] = (), - index_col_ids: typing.Sequence[str] = ["index"], - dtype: typing.Union[ - bigframes.dtypes.Dtype, typing.Sequence[bigframes.dtypes.Dtype] - ] = pandas.Float64Dtype(), - how="left", - ) -> CompiledArrayValue: - """ - 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 - """ - if how not in ("left", "right"): - raise ValueError("'how' must be 'left' or 'right'") - 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( - len(source_columns) == row_n for _, source_columns in unpivot_columns - ): - raise ValueError("Columns and row labels must all be same length.") - - unpivot_offset_id = bigframes.core.guid.generate_guid("unpivot_offsets_") - unpivot_table = table.cross_join( - ibis.memtable({unpivot_offset_id: range(row_n)}) - ) - # Use ibis memtable to infer type of rowlabels (if possible) - # TODO: Allow caller to specify dtype - 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 - ] - - label_columns = [] - for label_part, (col_id, label_dtype) in enumerate( - zip(index_col_ids, labels_dtypes) + class Builder: + def __init__( + self, + table: ibis_types.Table, + ordering: ExpressionOrdering, + columns: Collection[ibis_types.Value] = (), + hidden_ordering_columns: Collection[ibis_types.Value] = (), + predicates: Optional[Collection[ibis_types.BooleanValue]] = None, ): - # 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)): - 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 - ) - unpivot_values.append(unpivot_value.name(result_col)) - - unpivot_table = unpivot_table.select( - 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=tuple( - [ - *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=tuple( - [ - 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 CompiledArrayValue( - table=unpivot_table, - 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) -> CompiledArrayValue: - return self._set_or_replace_by_id( - destination_id, self._get_ibis_column(source_id) - ) - - def assign_constant( - self, - destination_id: str, - value: typing.Any, - dtype: typing.Optional[bigframes.dtypes.Dtype], - ) -> CompiledArrayValue: - # 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)}. {constants.FEEDBACK_LINK}" + 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) -> OrderedIR: + return OrderedIR( + table=self.table, + columns=self.columns, + hidden_ordering_columns=self.hidden_ordering_columns, + ordering=self.ordering, + predicates=self.predicates, ) - 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 - ) -> CompiledArrayValue: - """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_ids: - builder.columns = [ - val if (col_id != id) else new_value.name(id) - for col_id, val in zip(self.column_ids, self._columns) - ] - else: - builder.columns = [*self.columns, new_value.name(id)] - return builder.build() - - -class ArrayValueBuilder: - """Mutable expression class. - Use ArrayValue.builder() to create from a ArrayValue object. - """ - - def __init__( - self, - table: ibis_types.Table, - ordering: ExpressionOrdering, - columns: Collection[ibis_types.Value] = (), - hidden_ordering_columns: Collection[ibis_types.Value] = (), - predicates: Optional[Collection[ibis_types.BooleanValue]] = None, - ): - 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) -> CompiledArrayValue: - return CompiledArrayValue( - table=self.table, - columns=self.columns, - hidden_ordering_columns=self.hidden_ordering_columns, - ordering=self.ordering, - predicates=self.predicates, - ) def _reduce_predicate_list( diff --git a/bigframes/core/compile/compiler.py b/bigframes/core/compile/compiler.py index 195d830122a..662e73a4331 100644 --- a/bigframes/core/compile/compiler.py +++ b/bigframes/core/compile/compiler.py @@ -19,7 +19,8 @@ import pandas as pd -import bigframes.core.compile as compiled +import bigframes.core.compile.compiled as compiled +import bigframes.core.compile.concat as concat_impl import bigframes.core.compile.single_column import bigframes.core.nodes as nodes @@ -28,120 +29,167 @@ import bigframes.session +def compile_ordered(node: nodes.BigFrameNode) -> compiled.OrderedIR: + return typing.cast(compiled.OrderedIR, compile_node(node, True)) + + +def compile_unordered(node: nodes.BigFrameNode) -> compiled.UnorderedIR: + return typing.cast(compiled.UnorderedIR, compile_node(node, False)) + + @functools.cache -def compile_node(node: nodes.BigFrameNode) -> compiled.CompiledArrayValue: +def compile_node( + node: nodes.BigFrameNode, ordered: bool = True +) -> compiled.UnorderedIR | compiled.OrderedIR: """Compile node into CompileArrayValue. Caches result.""" - return _compile_node(node) + return _compile_node(node, ordered) @functools.singledispatch -def _compile_node(node: nodes.BigFrameNode) -> compiled.CompiledArrayValue: +def _compile_node( + node: nodes.BigFrameNode, ordered: bool = True +) -> compiled.UnorderedIR: """Defines transformation but isn't cached, always use compile_node instead""" - raise ValueError(f"Can't compile unnrecognized node: {node}") + raise ValueError(f"Can't compile unrecognized node: {node}") @_compile_node.register -def compile_join(node: nodes.JoinNode): - compiled_left = compile_node(node.left_child) - compiled_right = compile_node(node.right_child) - return bigframes.core.compile.single_column.join_by_column( - compiled_left, - node.left_column_ids, - compiled_right, - node.right_column_ids, - how=node.how, - allow_row_identity_join=node.allow_row_identity_join, - ) +def compile_join(node: nodes.JoinNode, ordered: bool = True): + if ordered: + left_ordered = compile_ordered(node.left_child) + right_ordered = compile_ordered(node.right_child) + return bigframes.core.compile.single_column.join_by_column_ordered( + left_ordered, + node.left_column_ids, + right_ordered, + node.right_column_ids, + how=node.how, + allow_row_identity_join=node.allow_row_identity_join, + ) + else: + left_unordered = compile_unordered(node.left_child) + right_unordered = compile_unordered(node.right_child) + return bigframes.core.compile.single_column.join_by_column_unordered( + left_unordered, + node.left_column_ids, + right_unordered, + node.right_column_ids, + how=node.how, + allow_row_identity_join=node.allow_row_identity_join, + ) @_compile_node.register -def compile_select(node: nodes.SelectNode): - return compile_node(node.child).select_columns(node.column_ids) +def compile_select(node: nodes.SelectNode, ordered: bool = True): + return compile_node(node.child, ordered).select_columns(node.column_ids) @_compile_node.register -def compile_drop(node: nodes.DropColumnsNode): - return compile_node(node.child).drop_columns(node.columns) +def compile_drop(node: nodes.DropColumnsNode, ordered: bool = True): + return compile_node(node.child, ordered).drop_columns(node.columns) @_compile_node.register -def compile_readlocal(node: nodes.ReadLocalNode): +def compile_readlocal(node: nodes.ReadLocalNode, ordered: bool = True): array_as_pd = pd.read_feather(io.BytesIO(node.feather_bytes)) - return compiled.CompiledArrayValue.mem_expr_from_pandas(array_as_pd) + ordered_ir = compiled.OrderedIR.from_pandas(array_as_pd) + if ordered: + return ordered_ir + else: + return ordered_ir.to_unordered() @_compile_node.register -def compile_readgbq(node: nodes.ReadGbqNode): - return compiled.CompiledArrayValue( - node.table, - node.columns, - node.hidden_ordering_columns, - node.ordering, - ) +def compile_readgbq(node: nodes.ReadGbqNode, ordered: bool = True): + if ordered: + return compiled.OrderedIR( + node.table, + node.columns, + node.hidden_ordering_columns, + node.ordering, + ) + else: + return compiled.UnorderedIR( + node.table, + node.columns, + ) @_compile_node.register -def compile_promote_offsets(node: nodes.PromoteOffsetsNode): - return compile_node(node.child).promote_offsets(node.col_id) +def compile_promote_offsets(node: nodes.PromoteOffsetsNode, ordered: bool = True): + result = compile_ordered(node.child).promote_offsets(node.col_id) + return result if ordered else result.to_unordered() @_compile_node.register -def compile_filter(node: nodes.FilterNode): - return compile_node(node.child).filter(node.predicate_id, node.keep_null) +def compile_filter(node: nodes.FilterNode, ordered: bool = True): + return compile_node(node.child, ordered).filter(node.predicate_id, node.keep_null) @_compile_node.register -def compile_orderby(node: nodes.OrderByNode): - return compile_node(node.child).order_by(node.by, node.stable) +def compile_orderby(node: nodes.OrderByNode, ordered: bool = True): + if ordered: + return compile_ordered(node.child).order_by(node.by, node.stable) + else: + return compile_unordered(node.child) @_compile_node.register -def compile_reversed(node: nodes.ReversedNode): - return compile_node(node.child).reversed() +def compile_reversed(node: nodes.ReversedNode, ordered: bool = True): + if ordered: + return compile_ordered(node.child).reversed() + else: + return compile_unordered(node.child) @_compile_node.register -def compile_project_unary(node: nodes.ProjectUnaryOpNode): - return compile_node(node.child).project_unary_op( +def compile_project_unary(node: nodes.ProjectUnaryOpNode, ordered: bool = True): + return compile_node(node.child, ordered).project_unary_op( node.input_id, node.op, node.output_id ) @_compile_node.register -def compile_project_binary(node: nodes.ProjectBinaryOpNode): - return compile_node(node.child).project_binary_op( +def compile_project_binary(node: nodes.ProjectBinaryOpNode, ordered: bool = True): + return compile_node(node.child, ordered).project_binary_op( node.left_input_id, node.right_input_id, node.op, node.output_id ) @_compile_node.register -def compile_project_ternary(node: nodes.ProjectTernaryOpNode): - return compile_node(node.child).project_ternary_op( +def compile_project_ternary(node: nodes.ProjectTernaryOpNode, ordered: bool = True): + return compile_node(node.child, ordered).project_ternary_op( node.input_id1, node.input_id2, node.input_id3, node.op, node.output_id ) @_compile_node.register -def compile_concat(node: nodes.ConcatNode): - compiled_nodes = [compile_node(node) for node in node.children] - return compiled_nodes[0].concat(compiled_nodes[1:]) +def compile_concat(node: nodes.ConcatNode, ordered: bool = True): + if ordered: + compiled_ordered = [compile_ordered(node) for node in node.children] + return concat_impl.concat_ordered(compiled_ordered) + else: + compiled_unordered = [compile_unordered(node) for node in node.children] + return concat_impl.concat_unordered(compiled_unordered) @_compile_node.register -def compile_aggregate(node: nodes.AggregateNode): - return compile_node(node.child).aggregate( +def compile_aggregate(node: nodes.AggregateNode, ordered: bool = True): + result = compile_unordered(node.child).aggregate( node.aggregations, node.by_column_ids, node.dropna ) + return result if ordered else result.to_unordered() @_compile_node.register -def compile_corr(node: nodes.CorrNode): - return compile_node(node.child).corr_aggregate(node.corr_aggregations) +def compile_corr(node: nodes.CorrNode, ordered: bool = True): + result = compile_unordered(node.child).corr_aggregate(node.corr_aggregations) + return result if ordered else result.to_unordered() @_compile_node.register -def compile_window(node: nodes.WindowOpNode): - return compile_node(node.child).project_window_op( +def compile_window(node: nodes.WindowOpNode, ordered: bool = True): + result = compile_ordered(node.child).project_window_op( node.column_name, node.op, node.window_spec, @@ -149,16 +197,17 @@ def compile_window(node: nodes.WindowOpNode): never_skip_nulls=node.never_skip_nulls, skip_reproject_unsafe=node.skip_reproject_unsafe, ) + return result if ordered else result.to_unordered() @_compile_node.register -def compile_reproject(node: nodes.ReprojectOpNode): - return compile_node(node.child)._reproject_to_table() +def compile_reproject(node: nodes.ReprojectOpNode, ordered: bool = True): + return compile_node(node.child, ordered)._reproject_to_table() @_compile_node.register -def compile_unpivot(node: nodes.UnpivotNode): - return compile_node(node.child).unpivot( +def compile_unpivot(node: nodes.UnpivotNode, ordered: bool = True): + return compile_node(node.child, ordered).unpivot( node.row_labels, node.unpivot_columns, passthrough_columns=node.passthrough_columns, @@ -169,17 +218,17 @@ def compile_unpivot(node: nodes.UnpivotNode): @_compile_node.register -def compile_assign(node: nodes.AssignNode): - return compile_node(node.child).assign(node.source_id, node.destination_id) +def compile_assign(node: nodes.AssignNode, ordered: bool = True): + return compile_node(node.child, ordered).assign(node.source_id, node.destination_id) @_compile_node.register -def compile_assign_constant(node: nodes.AssignConstantNode): - return compile_node(node.child).assign_constant( +def compile_assign_constant(node: nodes.AssignConstantNode, ordered: bool = True): + return compile_node(node.child, ordered).assign_constant( node.destination_id, node.value, node.dtype ) @_compile_node.register -def compiler_random_sample(node: nodes.RandomSampleNode): - return compile_node(node.child)._uniform_sampling(node.fraction) +def compiler_random_sample(node: nodes.RandomSampleNode, ordered: bool = True): + return compile_node(node.child, ordered)._uniform_sampling(node.fraction) diff --git a/bigframes/core/compile/concat.py b/bigframes/core/compile/concat.py new file mode 100644 index 00000000000..d39569370e8 --- /dev/null +++ b/bigframes/core/compile/concat.py @@ -0,0 +1,100 @@ +# 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 math +import typing + +import ibis + +import bigframes.core.compile.compiled as compiled +from bigframes.core.ordering import ( + ExpressionOrdering, + OrderingColumnReference, + reencode_order_string, + StringEncoding, +) + +ORDER_ID_COLUMN = "bigframes_ordering_id" + + +def concat_unordered( + items: typing.Sequence[compiled.UnorderedIR], +) -> compiled.UnorderedIR: + """Append together multiple ArrayValue objects.""" + if len(items) == 1: + return items[0] + tables = [] + for expr in items: + table = expr._to_ibis_expr() + # Rename the value columns based on horizontal offset before applying union. + table = table.select( + [table[col].name(f"column_{i}") for i, col in enumerate(table.columns)] + ) + tables.append(table) + combined_table = ibis.union(*tables) + return compiled.UnorderedIR( + combined_table, + columns=[combined_table[col] for col in combined_table.columns], + ) + + +def concat_ordered( + items: typing.Sequence[compiled.OrderedIR], +) -> compiled.OrderedIR: + """Append together multiple ArrayValue objects.""" + if len(items) == 1: + return items[0] + + tables = [] + prefix_base = 10 + prefix_size = math.ceil(math.log(len(items), prefix_base)) + # Must normalize all ids to the same encoding size + max_encoding_size = max( + *[expression._ordering.string_encoding.length for expression in items], + ) + for i, expr in enumerate(items): + ordering_prefix = str(i).zfill(prefix_size) + 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. + table = table.select( + [ + table[col].name(f"column_{i}") + if col != ORDER_ID_COLUMN + else ( + ordering_prefix + + reencode_order_string(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_value_columns=tuple([OrderingColumnReference(ORDER_ID_COLUMN)]), + total_ordering_columns=frozenset([ORDER_ID_COLUMN]), + string_encoding=StringEncoding(True, prefix_size + max_encoding_size), + ) + return compiled.OrderedIR( + 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, + ) diff --git a/bigframes/core/compile/row_identity.py b/bigframes/core/compile/row_identity.py index 2e9bc0527ca..71d53f90dc4 100644 --- a/bigframes/core/compile/row_identity.py +++ b/bigframes/core/compile/row_identity.py @@ -23,16 +23,76 @@ import ibis.expr.types as ibis_types import bigframes.constants as constants -import bigframes.core.compile as compiled +import bigframes.core.compile.compiled as compiled import bigframes.core.joins.name_resolution as naming import bigframes.core.ordering as orderings SUPPORTED_ROW_IDENTITY_HOW = {"outer", "left", "inner"} -def join_by_row_identity( - left: compiled.CompiledArrayValue, right: compiled.CompiledArrayValue, *, how: str -) -> compiled.CompiledArrayValue: +def join_by_row_identity_unordered( + left: compiled.UnorderedIR, + right: compiled.UnorderedIR, + *, + how: str, +) -> compiled.UnorderedIR: + """Compute join when we are joining by row identity not a specific column.""" + if how not in SUPPORTED_ROW_IDENTITY_HOW: + raise NotImplementedError( + f"Only how='outer','left','inner' currently supported. {constants.FEEDBACK_LINK}" + ) + + 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 + + # Public mapping must use JOIN_NAME_REMAPPER to stay in sync with consumers of join result + map_left_id, map_right_id = naming.JOIN_NAME_REMAPPER( + left.column_ids, right.column_ids + ) + joined_columns = [ + _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]) + for key in right.column_ids + ] + + joined_expr = compiled.UnorderedIR( + left._table, + columns=joined_columns, + predicates=combined_predicates, + ) + return joined_expr + + +def join_by_row_identity_ordered( + left: compiled.OrderedIR, + right: compiled.OrderedIR, + *, + how: str, +) -> compiled.OrderedIR: """Compute join when we are joining by row identity not a specific column.""" if how not in SUPPORTED_ROW_IDENTITY_HOW: raise NotImplementedError( @@ -118,7 +178,7 @@ def join_by_row_identity( if key.column_id in right._hidden_ordering_column_names.keys() ] - joined_expr = compiled.CompiledArrayValue( + joined_expr = compiled.OrderedIR( left._table, columns=joined_columns, hidden_ordering_columns=hidden_ordering_columns, diff --git a/bigframes/core/compile/single_column.py b/bigframes/core/compile/single_column.py index 93ba3f16f1c..a9088feb49b 100644 --- a/bigframes/core/compile/single_column.py +++ b/bigframes/core/compile/single_column.py @@ -23,16 +23,16 @@ import ibis.expr.datatypes as ibis_dtypes import ibis.expr.types as ibis_types -import bigframes.core.compile as compiled +import bigframes.core.compile.compiled as compiled import bigframes.core.compile.row_identity import bigframes.core.joins as joining import bigframes.core.ordering as orderings -def join_by_column( - left: compiled.CompiledArrayValue, +def join_by_column_ordered( + left: compiled.OrderedIR, left_column_ids: typing.Sequence[str], - right: compiled.CompiledArrayValue, + right: compiled.OrderedIR, right_column_ids: typing.Sequence[str], *, how: Literal[ @@ -43,7 +43,7 @@ def join_by_column( "cross", ], allow_row_identity_join: bool = True, -) -> compiled.CompiledArrayValue: +) -> compiled.OrderedIR: """Join two expressions by column equality. Arguments: @@ -68,13 +68,13 @@ def join_by_column( # 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_ibis_column(lcol) .name("index") - .equals(right._get_any_column(rcol).name("index")) + .equals(right._get_ibis_column(rcol).name("index")) for lcol, rcol in zip(left_column_ids, right_column_ids) ) ): - return bigframes.core.compile.row_identity.join_by_row_identity( + return bigframes.core.compile.row_identity.join_by_row_identity_ordered( left, right, how=how ) else: @@ -89,12 +89,12 @@ def join_by_column( r_mapping = {**r_public_mapping, **r_hidden_mapping} left_table = left._to_ibis_expr( - "unordered", + ordering_mode="unordered", expose_hidden_cols=True, col_id_overrides=l_mapping, ) right_table = right._to_ibis_expr( - "unordered", + ordering_mode="unordered", expose_hidden_cols=True, col_id_overrides=r_mapping, ) @@ -135,7 +135,7 @@ def join_by_column( for col in right._hidden_ordering_columns ], ] - return compiled.CompiledArrayValue( + return compiled.OrderedIR( combined_table, columns=columns, hidden_ordering_columns=hidden_ordering_columns, @@ -143,6 +143,88 @@ def join_by_column( ) +def join_by_column_unordered( + left: compiled.UnorderedIR, + left_column_ids: typing.Sequence[str], + right: compiled.UnorderedIR, + right_column_ids: typing.Sequence[str], + *, + how: Literal[ + "inner", + "left", + "outer", + "right", + "cross", + ], + allow_row_identity_join: bool = True, +) -> compiled.UnorderedIR: + """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. + 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. 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 + and how in bigframes.core.compile.row_identity.SUPPORTED_ROW_IDENTITY_HOW + 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_ibis_column(lcol) + .name("index") + .equals(right._get_ibis_column(rcol).name("index")) + for lcol, rcol in zip(left_column_ids, right_column_ids) + ) + ): + return bigframes.core.compile.row_identity.join_by_row_identity_unordered( + left, right, how=how + ) + else: + # Value column mapping must use JOIN_NAME_REMAPPER to stay in sync with consumers of join result + l_mapping, r_mapping = joining.JOIN_NAME_REMAPPER( + left.column_ids, right.column_ids + ) + left_table = left._to_ibis_expr( + col_id_overrides=l_mapping, + ) + right_table = right._to_ibis_expr( + col_id_overrides=r_mapping, + ) + join_conditions = [ + 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) + ] + + combined_table = ibis.join( + left_table, + right_table, + predicates=join_conditions, + how=how, # type: ignore + ) + # We could filter out the original join columns, but predicates/ordering + # might still reference them in implicit joins. + 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] + return compiled.UnorderedIR( + combined_table, + columns=columns, + ) + + 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(): diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 1e143144fe3..0a035754915 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -861,6 +861,8 @@ def to_pandas( max_download_size: Optional[int] = None, sampling_method: Optional[str] = None, random_state: Optional[int] = None, + *, + ordered: bool = True, ) -> pandas.DataFrame: """Write DataFrame to pandas DataFrame. @@ -880,6 +882,9 @@ def to_pandas( 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. + ordered (bool, default True): + Determines whether the resulting pandas dataframe will be deterministically ordered. + In some cases, unordered may result in a faster-executing query. Returns: pandas.DataFrame: A pandas DataFrame with all rows and columns of this DataFrame if the @@ -891,6 +896,7 @@ def to_pandas( max_download_size=max_download_size, sampling_method=sampling_method, random_state=random_state, + ordered=ordered, ) self._set_internal_query_job(query_job) return df.set_axis(self._block.column_labels, axis=1, copy=False) diff --git a/bigframes/series.py b/bigframes/series.py index 9e111618de2..1952acbf6d5 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -271,6 +271,8 @@ def to_pandas( max_download_size: Optional[int] = None, sampling_method: Optional[str] = None, random_state: Optional[int] = None, + *, + ordered: bool = True, ) -> pandas.Series: """Writes Series to pandas Series. @@ -290,6 +292,10 @@ def to_pandas( 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. + ordered (bool, default True): + Determines whether the resulting pandas series will be deterministically ordered. + In some cases, unordered may result in a faster-executing query. + Returns: pandas.Series: A pandas Series with all rows of this Series if the data_sampling_threshold_mb @@ -300,6 +306,7 @@ def to_pandas( max_download_size=max_download_size, sampling_method=sampling_method, random_state=random_state, + ordered=ordered, ) self._set_internal_query_job(query_job) series = df[self._value_column] diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index b49e2469a91..8f9fa377873 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -617,10 +617,8 @@ def _read_gbq_table( is_total_ordering = True ordering = orderings.ExpressionOrdering( ordering_value_columns=tuple( - [ - core.OrderingColumnReference(column_id) - for column_id in total_ordering_cols - ] + core.OrderingColumnReference(column_id) + for column_id in total_ordering_cols ), total_ordering_columns=frozenset(total_ordering_cols), ) diff --git a/tests/system/large/ml/test_cluster.py b/tests/system/large/ml/test_cluster.py index eae6896669b..f01116665f1 100644 --- a/tests/system/large/ml/test_cluster.py +++ b/tests/system/large/ml/test_cluster.py @@ -16,7 +16,7 @@ import pytest from bigframes.ml import cluster -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal @pytest.mark.flaky(retries=2, delay=120) @@ -105,7 +105,7 @@ def test_cluster_configure_fit_score_predict( index=pd.Index(["test1", "test2", "test3", "test4"], dtype="string[pyarrow]"), ) expected.index.name = "observation" - assert_pandas_df_equal_ignore_ordering(result, expected) + assert_pandas_df_equal(result, expected, ignore_order=True) # save, load, check n_clusters to ensure configuration was kept reloaded_model = model.to_gbq( diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 6874a9f301a..3e569540587 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -24,7 +24,7 @@ pipeline, preprocessing, ) -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal def test_pipeline_linear_regression_fit_score_predict( @@ -555,7 +555,7 @@ def test_pipeline_standard_scaler_kmeans_fit_score_predict( ), ) expected.index.name = "observation" - assert_pandas_df_equal_ignore_ordering(result, expected) + assert_pandas_df_equal(result, expected, ignore_order=True) def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_index): diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index c8f8f66ebaa..6ed3e6511a2 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -32,7 +32,7 @@ get_cloud_function_name, get_remote_function_locations, ) -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal # 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 @@ -357,7 +357,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets(session.bqclient, functions_client, square) @@ -401,7 +401,7 @@ def add_one(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -446,7 +446,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets(session.bqclient, functions_client, square) @@ -497,7 +497,7 @@ def sign(num): 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) + assert_pandas_df_equal(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) @@ -542,7 +542,7 @@ def circumference(radius): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -591,7 +591,7 @@ def find_team(num): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -675,7 +675,7 @@ def inner_test(): 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) + assert_pandas_df_equal(bf_result, pd_result) # Test that the remote function works as expected inner_test() @@ -765,7 +765,7 @@ def is_odd(num): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -808,7 +808,7 @@ def is_odd(num): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -852,7 +852,7 @@ def test_remote_udf_lambda( 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -909,7 +909,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -954,7 +954,7 @@ def pd_np_foo(x): # comparing for the purpose of this test pd_result.result = pd_result.result.astype(pandas.Float64Dtype()) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -998,7 +998,7 @@ def test_internal(rf, udf): 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) + assert_pandas_df_equal(bf_result, pd_result) # Create an explicit name for the remote function prefixer = test_utils.prefixer.Prefixer("foo", "") @@ -1167,7 +1167,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( @@ -1204,7 +1204,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) finally: # clean up the gcp assets created for the remote function cleanup_remote_function_assets( diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py index d95a1e1bc22..caeffa77689 100644 --- a/tests/system/small/ml/test_cluster.py +++ b/tests/system/small/ml/test_cluster.py @@ -15,7 +15,7 @@ import pandas as pd from bigframes.ml import cluster -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal _PD_NEW_PENGUINS = pd.DataFrame.from_dict( { @@ -68,7 +68,7 @@ def test_kmeans_predict(session, penguins_kmeans_model: cluster.KMeans): dtype="Int64", index=pd.Index(["test1", "test2", "test3", "test4"], dtype="string[pyarrow]"), ) - assert_pandas_df_equal_ignore_ordering(result, expected) + assert_pandas_df_equal(result, expected, ignore_order=True) def test_kmeans_score(session, penguins_kmeans_model: cluster.KMeans): diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index f911dd7eebc..ec1f351d876 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -225,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], }, ) - tests.system.utils.assert_pandas_df_equal_ignore_ordering( + tests.system.utils.assert_pandas_df_equal( result, expected, check_exact=False, @@ -233,6 +233,7 @@ def test_pca_model_principal_component_info(penguins_bqml_pca_model: core.BqmlMo # int64 Index by default in pandas versus Int64 (nullable) Index in BigQuery DataFrame check_index_type=False, check_dtype=False, + ignore_order=True, ) diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index e31681f4a09..cc4d2e5801c 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -130,13 +130,14 @@ def test_pca_explained_variance_(penguins_pca_model: decomposition.PCA): "explained_variance": [3.278657, 1.270829, 1.125354], }, ) - tests.system.utils.assert_pandas_df_equal_ignore_ordering( + tests.system.utils.assert_pandas_df_equal( result, expected, check_exact=False, rtol=0.1, check_index_type=False, check_dtype=False, + ignore_order=True, ) @@ -149,11 +150,12 @@ def test_pca_explained_variance_ratio_(penguins_pca_model: decomposition.PCA): "explained_variance_ratio": [0.469357, 0.181926, 0.1611], }, ) - tests.system.utils.assert_pandas_df_equal_ignore_ordering( + tests.system.utils.assert_pandas_df_equal( result, expected, check_exact=False, rtol=0.1, check_index_type=False, check_dtype=False, + ignore_order=True, ) diff --git a/tests/system/small/operations/test_datetimes.py b/tests/system/small/operations/test_datetimes.py index 7dc55b9367c..177194c7a8a 100644 --- a/tests/system/small/operations/test_datetimes.py +++ b/tests/system/small/operations/test_datetimes.py @@ -16,7 +16,7 @@ import pytest import bigframes.series -from tests.system.utils import assert_series_equal_ignoring_order +from tests.system.utils import assert_series_equal DATETIME_COL_NAMES = [("datetime_col",), ("timestamp_col",)] @@ -33,7 +33,7 @@ def test_day(scalars_dfs, col_name): bf_result = bf_series.dt.day.to_pandas() pd_result = scalars_pandas_df[col_name].dt.day - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -51,7 +51,7 @@ def test_date(scalars_dfs, col_name): bf_result = bf_series.dt.date.to_pandas() pd_result = scalars_pandas_df[col_name].dt.date - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -69,7 +69,7 @@ def test_dayofweek(scalars_dfs, col_name): 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) + assert_series_equal(pd_result, bf_result, check_dtype=False) @pytest.mark.parametrize( @@ -84,7 +84,7 @@ def test_hour(scalars_dfs, col_name): bf_result = bf_series.dt.hour.to_pandas() pd_result = scalars_pandas_df[col_name].dt.hour - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -102,7 +102,7 @@ def test_minute(scalars_dfs, col_name): bf_result = bf_series.dt.minute.to_pandas() pd_result = scalars_pandas_df[col_name].dt.minute - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -120,7 +120,7 @@ def test_month(scalars_dfs, col_name): bf_result = bf_series.dt.month.to_pandas() pd_result = scalars_pandas_df[col_name].dt.month - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -138,7 +138,7 @@ def test_quarter(scalars_dfs, col_name): bf_result = bf_series.dt.quarter.to_pandas() pd_result = scalars_pandas_df[col_name].dt.quarter - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -156,7 +156,7 @@ def test_second(scalars_dfs, col_name): bf_result = bf_series.dt.second.to_pandas() pd_result = scalars_pandas_df[col_name].dt.second - assert_series_equal_ignoring_order( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -174,7 +174,7 @@ def test_time(scalars_dfs, col_name): bf_result = bf_series.dt.time.to_pandas() pd_result = scalars_pandas_df[col_name].dt.time - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -192,7 +192,7 @@ def test_year(scalars_dfs, col_name): bf_result = bf_series.dt.year.to_pandas() pd_result = scalars_pandas_df[col_name].dt.year - assert_series_equal_ignoring_order( + assert_series_equal( 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 index 241cbd576b1..27a35134d4a 100644 --- a/tests/system/small/operations/test_strings.py +++ b/tests/system/small/operations/test_strings.py @@ -19,7 +19,7 @@ import bigframes.series -from ...utils import assert_series_equal_ignoring_order +from ...utils import assert_series_equal def test_find(scalars_dfs): @@ -31,7 +31,7 @@ def test_find(scalars_dfs): # 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( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -173,7 +173,7 @@ def test_len(scalars_dfs): # 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( + assert_series_equal( pd_result.astype(pd.Int64Dtype()), bf_result, ) @@ -186,7 +186,7 @@ def test_lower(scalars_dfs): bf_result = bf_series.str.lower().to_pandas() pd_result = scalars_pandas_df[col_name].str.lower() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -205,7 +205,7 @@ def test_reverse(scalars_dfs): else: pd_result.loc[i] = cell[::-1] - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -222,7 +222,7 @@ def test_slice(scalars_dfs, start, stop): pd_series = scalars_pandas_df[col_name] pd_result = pd_series.str.slice(start, stop) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -235,7 +235,7 @@ def test_strip(scalars_dfs): bf_result = bf_series.str.strip().to_pandas() pd_result = scalars_pandas_df[col_name].str.strip() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -248,7 +248,7 @@ def test_upper(scalars_dfs): bf_result = bf_series.str.upper().to_pandas() pd_result = scalars_pandas_df[col_name].str.upper() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -330,7 +330,7 @@ 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( + assert_series_equal( bf_result, pd_result.astype(pd.BooleanDtype()) # the dtype here is a case of intentional diversion from pandas @@ -342,7 +342,7 @@ 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( + assert_series_equal( bf_result, pd_result.astype(pd.BooleanDtype()) # the dtype here is a case of intentional diversion from pandas @@ -357,7 +357,7 @@ def test_rstrip(scalars_dfs): bf_result = bf_series.str.rstrip().to_pandas() pd_result = scalars_pandas_df[col_name].str.rstrip() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -370,7 +370,7 @@ def test_lstrip(scalars_dfs): bf_result = bf_series.str.lstrip().to_pandas() pd_result = scalars_pandas_df[col_name].str.lstrip() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -384,7 +384,7 @@ def test_repeat(scalars_dfs, repeats): bf_result = bf_series.str.repeat(repeats).to_pandas() pd_result = scalars_pandas_df[col_name].str.repeat(repeats) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -397,7 +397,7 @@ def test_capitalize(scalars_dfs): bf_result = bf_series.str.capitalize().to_pandas() pd_result = scalars_pandas_df[col_name].str.capitalize() - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -415,7 +415,7 @@ def test_cat_with_series(scalars_dfs): pd_right = scalars_pandas_df[col_name] pd_result = pd_left.str.cat(others=pd_right) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -429,7 +429,7 @@ def test_str_match(scalars_dfs): bf_result = bf_series.str.match(pattern).to_pandas() pd_result = scalars_pandas_df[col_name].str.match(pattern) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -443,7 +443,7 @@ def test_str_fullmatch(scalars_dfs): bf_result = bf_series.str.fullmatch(pattern).to_pandas() pd_result = scalars_pandas_df[col_name].str.fullmatch(pattern) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -456,7 +456,7 @@ 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) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -469,7 +469,7 @@ def test_str_pad(scalars_dfs): 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( + assert_series_equal( pd_result, bf_result, ) @@ -492,7 +492,7 @@ def test_str_ljust(scalars_dfs): 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( + assert_series_equal( pd_result, bf_result, ) @@ -505,7 +505,7 @@ def test_str_rjust(scalars_dfs): 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( + assert_series_equal( pd_result, bf_result, ) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 605d4abc1d0..e522878229a 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -28,10 +28,7 @@ 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, -) +from tests.system.utils import assert_pandas_df_equal, assert_series_equal def test_df_construct_copy(scalars_dfs): @@ -98,7 +95,7 @@ def test_get_column(scalars_dfs): series = scalars_df[col_name] bf_result = series.to_pandas() pd_result = scalars_pandas_df[col_name] - assert_series_equal_ignoring_order(bf_result, pd_result) + assert_series_equal(bf_result, pd_result) def test_get_column_nonstring(scalars_dfs): @@ -106,7 +103,7 @@ def test_get_column_nonstring(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) + assert_series_equal(bf_result, pd_result) def test_hasattr(scalars_dfs): @@ -116,15 +113,24 @@ def test_hasattr(scalars_dfs): assert not hasattr(scalars_df, "not_exist") -def test_head_with_custom_column_labels(scalars_df_index, scalars_pandas_df_index): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_head_with_custom_column_labels( + scalars_df_index, scalars_pandas_df_index, ordered +): rename_mapping = { "int64_col": "Integer Column", "string_col": "言語列", } bf_df = scalars_df_index.rename(columns=rename_mapping).head(3) - bf_result = bf_df.to_pandas() + bf_result = bf_df.to_pandas(ordered=ordered) pd_result = scalars_pandas_df_index.rename(columns=rename_mapping).head(3) - pandas.testing.assert_frame_equal(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=not ordered) def test_tail_with_custom_column_labels(scalars_df_index, scalars_pandas_df_index): @@ -183,7 +189,7 @@ def test_get_column_by_attr(scalars_dfs): series = scalars_df.int64_col bf_result = series.to_pandas() pd_result = scalars_pandas_df.int64_col - assert_series_equal_ignoring_order(bf_result, pd_result) + assert_series_equal(bf_result, pd_result) def test_get_columns(scalars_dfs): @@ -246,7 +252,7 @@ def test_drop_with_custom_column_labels(scalars_dfs): pd_result = scalars_pandas_df.rename(columns=rename_mapping).drop( columns=dropped_columns ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result) def test_drop_index(scalars_dfs): @@ -420,7 +426,7 @@ def test_filter_df(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_assign_new_column(scalars_dfs): @@ -433,7 +439,7 @@ def test_assign_new_column(scalars_dfs): # 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) + assert_pandas_df_equal(bf_result, pd_result) def test_assign_new_column_w_loc(scalars_dfs): @@ -564,7 +570,7 @@ def test_assign_existing_column(scalars_dfs): # 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) + assert_pandas_df_equal(bf_result, pd_result) def test_assign_listlike_to_empty_df(session): @@ -576,7 +582,7 @@ def test_assign_listlike_to_empty_df(session): pd_result["new_col"] = pd_result["new_col"].astype("Int64") pd_result.index = pd_result.index.astype("Int64") - assert_pandas_df_equal_ignore_ordering(bf_result.to_pandas(), pd_result) + assert_pandas_df_equal(bf_result.to_pandas(), pd_result) def test_assign_to_empty_df_multiindex_error(session): @@ -595,14 +601,21 @@ def test_assign_to_empty_df_multiindex_error(session): empty_pandas_df.assign(new_col=[1, 2, 3, 4, 5, 6, 7, 8, 9]) -def test_assign_series(scalars_dfs): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_assign_series(scalars_dfs, ordered): scalars_df, scalars_pandas_df = scalars_dfs column_name = "int64_col" df = scalars_df.assign(new_col=scalars_df[column_name]) - bf_result = df.to_pandas() + bf_result = df.to_pandas(ordered=ordered) pd_result = scalars_pandas_df.assign(new_col=scalars_pandas_df[column_name]) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=not ordered) def test_assign_series_overwrite(scalars_dfs): @@ -614,7 +627,7 @@ def test_assign_series_overwrite(scalars_dfs): **{column_name: scalars_pandas_df[column_name] + 3} ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result) def test_assign_sequential(scalars_dfs): @@ -629,7 +642,7 @@ def test_assign_sequential(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) # Require an index so that the self-join is consistent each time. @@ -663,7 +676,7 @@ def test_assign_different_df( new_col=scalars_pandas_df_index[column_name] ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result) def test_assign_different_df_w_loc( @@ -714,7 +727,7 @@ def test_assign_callable_lambda(scalars_dfs): # 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.parametrize( @@ -965,7 +978,9 @@ def test_df_merge(scalars_dfs, merge_how): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal( + bf_result, pd_result, ignore_order=True, check_index_type=False + ) @pytest.mark.parametrize( @@ -998,7 +1013,9 @@ def test_df_merge_multi_key(scalars_dfs, left_on, right_on): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal( + bf_result, pd_result, ignore_order=True, check_index_type=False + ) @pytest.mark.parametrize( @@ -1028,7 +1045,9 @@ def test_merge_custom_col_name(scalars_dfs, merge_how): 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) + assert_pandas_df_equal( + bf_result, pd_result, ignore_order=True, check_index_type=False + ) @pytest.mark.parametrize( @@ -1061,7 +1080,9 @@ def test_merge_left_on_right_on(scalars_dfs, merge_how): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal( + bf_result, pd_result, ignore_order=True, check_index_type=False + ) def test_get_dtypes(scalars_df_default_index): @@ -1375,7 +1396,7 @@ def test_df_abs(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_df_isnull(scalars_dfs): @@ -1392,7 +1413,7 @@ def test_df_isnull(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_df_notnull(scalars_dfs): @@ -1409,7 +1430,7 @@ def test_df_notnull(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.parametrize( @@ -1629,7 +1650,7 @@ def test_scalar_binop(scalars_dfs, op, other_scalar, reverse_operands): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.parametrize(("other_scalar"), [1, -2]) @@ -1641,7 +1662,7 @@ def test_mod(scalars_dfs, other_scalar): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_scalar_binop_str_exception(scalars_dfs): @@ -1697,7 +1718,7 @@ def test_series_binop_axis_index( 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.parametrize( @@ -1743,8 +1764,15 @@ def test_binop_df_df_binary_op( # Differnt table will only work for explicit index, since default index orders are arbitrary. +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) def test_series_binop_add_different_table( - scalars_df_index, scalars_pandas_df_index, scalars_df_2_index + scalars_df_index, scalars_pandas_df_index, scalars_df_2_index, ordered ): df_columns = ["int64_col", "float64_col"] series_column = "int64_too" @@ -1752,13 +1780,13 @@ def test_series_binop_add_different_table( bf_result = ( scalars_df_index[df_columns] .add(scalars_df_2_index[series_column], axis="index") - .to_pandas() + .to_pandas(ordered=ordered) ) 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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=not ordered) # TODO(garrettwu): Test series binop with different index @@ -1779,7 +1807,7 @@ def test_join_same_table(scalars_dfs, how): 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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) @all_joins @@ -1792,7 +1820,7 @@ def test_join_different_table( 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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) def test_join_duplicate_columns_raises_not_implemented(scalars_dfs): @@ -1821,7 +1849,7 @@ def test_join_param_on(scalars_dfs, how): 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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) @pytest.mark.parametrize( @@ -2037,7 +2065,14 @@ def test_df_describe(scalars_dfs): ).all() -def test_df_stack(scalars_dfs): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_df_stack(scalars_dfs, ordered): 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 @@ -2047,11 +2082,13 @@ def test_df_stack(scalars_dfs): # Can only stack identically-typed columns columns = ["int64_col", "int64_too", "rowindex_2"] - bf_result = scalars_df[columns].stack().to_pandas() + bf_result = scalars_df[columns].stack().to_pandas(ordered=ordered) 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) + assert_series_equal( + bf_result, pd_result, check_dtype=False, ignore_order=not ordered + ) def test_df_melt_default(scalars_dfs): @@ -2097,7 +2134,14 @@ def test_df_melt_parameterized(scalars_dfs): ) -def test_df_unstack(scalars_dfs): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_df_unstack(scalars_dfs, ordered): scalars_df, scalars_pandas_df = scalars_dfs # To match bigquery dataframes scalars_pandas_df = scalars_pandas_df.copy() @@ -2110,11 +2154,13 @@ def test_df_unstack(scalars_dfs): ] # unstack on mono-index produces series - bf_result = scalars_df[columns].unstack().to_pandas() + bf_result = scalars_df[columns].unstack().to_pandas(ordered=ordered) 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) + assert_series_equal( + bf_result, pd_result, check_dtype=False, ignore_order=not ordered + ) @pytest.mark.parametrize( @@ -2259,14 +2305,18 @@ def test_iloc_slice_zero_step(scalars_df_index): 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:].to_pandas() +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index, ordered): + bf_result = scalars_df_index.iloc[1:].iloc[1:].to_pandas(ordered=ordered) pd_result = scalars_pandas_df_index.iloc[1:].iloc[1:] - pd.testing.assert_frame_equal( - bf_result, - pd_result, - ) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=not ordered) @pytest.mark.parametrize( @@ -2457,6 +2507,13 @@ def test_loc_setitem_bool_series_scalar_type_error(scalars_dfs): pd_df.loc[pd_df["int64_too"] == 1, "string_col"] = 99 +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) @pytest.mark.parametrize( ("op"), [ @@ -2471,16 +2528,18 @@ def test_loc_setitem_bool_series_scalar_type_error(scalars_dfs): ], ids=["sum", "mean", "min", "max", "std", "var", "count", "nunique"], ) -def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op): +def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op, ordered): 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.to_pandas() + bf_result = bf_series.to_pandas(ordered=ordered) # 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) + assert_series_equal( + pd_series, bf_result, check_index_type=False, ignore_order=not ordered + ) @pytest.mark.parametrize( @@ -2571,16 +2630,25 @@ def test_df_skew_too_few_values(scalars_dfs): pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False) -def test_df_skew(scalars_dfs): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_df_skew(scalars_dfs, ordered): columns = ["float64_col", "int64_col"] scalars_df, scalars_pandas_df = scalars_dfs - bf_result = scalars_df[columns].skew().to_pandas() + bf_result = scalars_df[columns].skew().to_pandas(ordered=ordered) 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) + assert_series_equal( + pd_result, bf_result, check_index_type=False, ignore_order=not ordered + ) def test_df_kurt_too_few_values(scalars_dfs): @@ -2731,9 +2799,10 @@ 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()) # Ignore ordering as pandas order differently depending on version - assert_pandas_df_equal_ignore_ordering( + assert_pandas_df_equal( bf_result, pd_result, + ignore_order=True, check_names=False, ) diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index 8f5d706f621..d700d93be99 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -19,10 +19,7 @@ import pyarrow as pa import pytest -from tests.system.utils import ( - assert_pandas_df_equal_ignore_ordering, - convert_pandas_dtypes, -) +from tests.system.utils import assert_pandas_df_equal, convert_pandas_dtypes try: import pandas_gbq # type: ignore @@ -380,7 +377,7 @@ def test_to_sql_query_unnamed_index_included( 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) + assert_pandas_df_equal(roundtrip.to_pandas(), pd_df, check_index_type=False) def test_to_sql_query_named_index_included( @@ -397,7 +394,7 @@ def test_to_sql_query_named_index_included( 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) + assert_pandas_df_equal(roundtrip.to_pandas(), pd_df) def test_to_sql_query_unnamed_index_excluded( @@ -412,7 +409,9 @@ def test_to_sql_query_unnamed_index_excluded( 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(), pd_df) + assert_pandas_df_equal( + roundtrip.to_pandas(), pd_df, check_index_type=False, ignore_order=True + ) def test_to_sql_query_named_index_excluded( @@ -429,4 +428,6 @@ def test_to_sql_query_named_index_excluded( "rowindex_2", drop=True ).reset_index(drop=True) roundtrip = session.read_gbq(sql) - assert_pandas_df_equal_ignore_ordering(roundtrip.to_pandas(), pd_df) + assert_pandas_df_equal( + roundtrip.to_pandas(), pd_df, check_index_type=False, ignore_order=True + ) diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 05154f7ab7d..a24713c2b3f 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -16,6 +16,7 @@ import pytest import bigframes.pandas as bpd +from tests.system.utils import assert_pandas_df_equal @pytest.mark.parametrize( @@ -88,16 +89,23 @@ 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): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_dataframe_groupby_agg_string( + scalars_df_index, scalars_pandas_df_index, ordered +): 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() + bf_result_computed = bf_result.to_pandas(ordered=ordered) - pd.testing.assert_frame_equal( - pd_result, - bf_result_computed, - check_dtype=False, + assert_pandas_df_equal( + pd_result, bf_result_computed, check_dtype=False, ignore_order=not ordered ) @@ -270,13 +278,22 @@ def test_dataframe_groupby_kurt(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) -def test_dataframe_groupby_diff(scalars_df_index, scalars_pandas_df_index): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_dataframe_groupby_diff(scalars_df_index, scalars_pandas_df_index, ordered): 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() + bf_result_computed = bf_result.to_pandas(ordered=ordered) - pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + assert_pandas_df_equal( + pd_result, bf_result_computed, check_dtype=False, ignore_order=not ordered + ) def test_dataframe_groupby_getitem( diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index bc35f633fdb..e7e93849c60 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -16,7 +16,7 @@ import pytest import bigframes.pandas as bpd -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal # Row Multi-index tests @@ -429,7 +429,7 @@ def test_multi_index_dataframe_join(scalars_dfs, how): (["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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) @all_joins @@ -450,7 +450,7 @@ def test_multi_index_dataframe_join_on(scalars_dfs, how): 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) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) @pytest.mark.parametrize( diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py index b88901f3bcf..a1079288cfb 100644 --- a/tests/system/small/test_pandas.py +++ b/tests/system/small/test_pandas.py @@ -16,16 +16,23 @@ import pytest import bigframes.pandas as bpd -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal -def test_concat_dataframe(scalars_dfs): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_concat_dataframe(scalars_dfs, ordered): scalars_df, scalars_pandas_df = scalars_dfs bf_result = bpd.concat(11 * [scalars_df]) - bf_result = bf_result.to_pandas() + bf_result = bf_result.to_pandas(ordered=ordered) pd_result = pd.concat(11 * [scalars_pandas_df]) - pd.testing.assert_frame_equal(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=not ordered) def test_concat_series(scalars_dfs): @@ -252,7 +259,7 @@ def test_merge(scalars_dfs, merge_how): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) @pytest.mark.parametrize( @@ -286,7 +293,7 @@ def test_merge_left_on_right_on(scalars_dfs, merge_how): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) def test_pd_merge_cross(scalars_dfs): @@ -341,7 +348,7 @@ def test_merge_series(scalars_dfs, merge_how): sort=True, ) - assert_pandas_df_equal_ignore_ordering(bf_result, pd_result) + assert_pandas_df_equal(bf_result, pd_result, ignore_order=True) def test_cut(scalars_dfs): diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 89907a53dfc..3d8532a13b6 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -18,7 +18,7 @@ import bigframes from bigframes import remote_function as rf -from tests.system.utils import assert_pandas_df_equal_ignore_ordering +from tests.system.utils import assert_pandas_df_equal @pytest.fixture(scope="module") @@ -121,7 +121,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -170,7 +170,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -246,7 +246,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -309,7 +309,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -348,7 +348,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -387,7 +387,7 @@ def square(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -418,7 +418,7 @@ def add_one(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -447,7 +447,7 @@ def add_one(x): 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.flaky(retries=2, delay=120) @@ -535,7 +535,7 @@ def square1(x): 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()) + assert_pandas_df_equal(s1_result.to_pandas(), s2_result.to_pandas()) @pytest.mark.flaky(retries=2, delay=120) @@ -583,7 +583,9 @@ def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id): 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) + assert_pandas_df_equal( + direct_df, indirect_df, ignore_order=True, check_index_type=False + ) @pytest.mark.flaky(retries=2, delay=120) diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 006d370818d..f59d64fe06f 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -24,10 +24,7 @@ import bigframes.pandas import bigframes.series as series -from tests.system.utils import ( - assert_pandas_df_equal_ignore_ordering, - assert_series_equal_ignoring_order, -) +from tests.system.utils import assert_pandas_df_equal, assert_series_equal def test_series_construct_copy(scalars_dfs): @@ -210,7 +207,7 @@ def test_abs(scalars_dfs, col_name): 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) + assert_series_equal(pd_result, bf_result) def test_fillna(scalars_dfs): @@ -218,7 +215,7 @@ def test_fillna(scalars_dfs): col_name = "string_col" bf_result = scalars_df[col_name].fillna("Missing").to_pandas() pd_result = scalars_pandas_df[col_name].fillna("Missing") - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -496,7 +493,7 @@ def test_series_int_int_operators_scalar( 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) + assert_series_equal(pd_result, bf_result) def test_series_pow_scalar(scalars_dfs): @@ -505,7 +502,7 @@ def test_series_pow_scalar(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) + assert_series_equal(pd_result, bf_result) def test_series_pow_scalar_reverse(scalars_dfs): @@ -514,7 +511,7 @@ def test_series_pow_scalar_reverse(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) + assert_series_equal(pd_result, bf_result) @pytest.mark.parametrize( @@ -540,7 +537,7 @@ def test_series_bool_bool_operators_scalar( 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) + assert_series_equal(pd_result.astype(pd.BooleanDtype()), bf_result) @pytest.mark.parametrize( @@ -578,7 +575,7 @@ 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"]).to_pandas() pd_result = operator(scalars_pandas_df["int64_col"], scalars_pandas_df["int64_too"]) - assert_series_equal_ignoring_order(pd_result, bf_result) + assert_series_equal(pd_result, bf_result) @pytest.mark.parametrize( @@ -728,7 +725,7 @@ def test_series_add_scalar(scalars_dfs, other): 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) + assert_series_equal(pd_result, bf_result) @pytest.mark.parametrize( @@ -744,7 +741,7 @@ def test_series_add_bigframes_series(scalars_dfs, left_col, right_col): 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) + assert_series_equal(pd_result, bf_result) @pytest.mark.parametrize( @@ -766,7 +763,7 @@ def test_series_add_bigframes_series_nested( scalars_pandas_df[left_col] + scalars_pandas_df[right_col] ) + scalars_pandas_df[righter_col] - assert_series_equal_ignoring_order(pd_result, bf_result) + assert_series_equal(pd_result, bf_result) def test_series_add_different_table_default_index( @@ -924,7 +921,7 @@ def test_isnull(scalars_dfs): # 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) + assert_series_equal(pd_series.astype(pd.BooleanDtype()), bf_series) def test_notnull(scalars_dfs): @@ -935,7 +932,7 @@ def test_notnull(scalars_dfs): # 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) + assert_series_equal(pd_series.astype(pd.BooleanDtype()), bf_series) def test_round(scalars_dfs): @@ -944,7 +941,7 @@ def test_round(scalars_dfs): 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) + assert_series_equal(pd_result, bf_result) def test_eq_scalar(scalars_dfs): @@ -953,7 +950,7 @@ def test_eq_scalar(scalars_dfs): 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) + assert_series_equal(pd_result, bf_result) def test_eq_wider_type_scalar(scalars_dfs): @@ -962,7 +959,7 @@ def test_eq_wider_type_scalar(scalars_dfs): 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) + assert_series_equal(pd_result, bf_result) def test_ne_scalar(scalars_dfs): @@ -971,7 +968,7 @@ def test_ne_scalar(scalars_dfs): 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) + assert_series_equal(pd_result, bf_result) def test_eq_int_scalar(scalars_dfs): @@ -980,7 +977,7 @@ def test_eq_int_scalar(scalars_dfs): 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) + assert_series_equal(pd_result, bf_result) @pytest.mark.parametrize( @@ -999,7 +996,7 @@ def test_eq_same_type_series(scalars_dfs, 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) + assert_series_equal(pd_result.astype(pd.BooleanDtype()), bf_result) def test_loc_setitem_cell(scalars_df_index, scalars_pandas_df_index): @@ -1036,7 +1033,7 @@ def test_ne_obj_series(scalars_dfs): # 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) + assert_series_equal(pd_result.astype(pd.BooleanDtype()), bf_result) def test_indexing_using_unselected_series(scalars_dfs): @@ -1045,7 +1042,7 @@ def test_indexing_using_unselected_series(scalars_dfs): 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( + assert_series_equal( pd_result, bf_result, ) @@ -1061,7 +1058,7 @@ def test_indexing_using_selected_series(scalars_dfs): scalars_pandas_df["string_col"].eq("Hello, World!") ] - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -1083,7 +1080,7 @@ def test_nested_filter(scalars_dfs): ) # 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( + assert_series_equal( pd_result, bf_result, ) @@ -1102,7 +1099,7 @@ def test_binop_repeated_application_does_row_identity_joins(scalars_dfs): bf_result = bf_series.to_pandas() pd_result = pd_series - assert_series_equal_ignoring_order( + assert_series_equal( bf_result, pd_result, ) @@ -1124,10 +1121,9 @@ def test_binop_opposite_filters(scalars_dfs): 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, - ) + # Passes with ignore_order=False only with some dependency sets + # TODO: Determine desired behavior and make test more strict + assert_series_equal(bf_result, pd_result, ignore_order=True) def test_binop_left_filtered(scalars_dfs): @@ -1142,10 +1138,9 @@ def test_binop_left_filtered(scalars_dfs): 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, - ) + # Passes with ignore_order=False only with some dependency sets + # TODO: Determine desired behavior and make test more strict + assert_series_equal(bf_result, pd_result, ignore_order=True) def test_binop_right_filtered(scalars_dfs): @@ -1160,7 +1155,7 @@ def test_binop_right_filtered(scalars_dfs): pd_bool_col = scalars_pandas_df["bool_col"] pd_result = pd_float64_col + pd_int64_col[pd_bool_col] - assert_series_equal_ignoring_order( + assert_series_equal( bf_result, pd_result, ) @@ -1265,7 +1260,7 @@ def test_groupby_sum(scalars_dfs): ) # TODO(swast): Update groupby to use index based on group by key(s). bf_result = bf_series.to_pandas() - assert_series_equal_ignoring_order( + assert_series_equal( pd_series, bf_result, check_exact=False, @@ -1283,7 +1278,7 @@ def test_groupby_std(scalars_dfs): .astype(pd.Float64Dtype()) ) bf_result = bf_series.to_pandas() - assert_series_equal_ignoring_order( + assert_series_equal( pd_series, bf_result, check_exact=False, @@ -1298,7 +1293,7 @@ def test_groupby_var(scalars_dfs): scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).var() ) bf_result = bf_series.to_pandas() - assert_series_equal_ignoring_order( + assert_series_equal( pd_series, bf_result, check_exact=False, @@ -1350,7 +1345,7 @@ def test_groupby_mean(scalars_dfs): ) # TODO(swast): Update groupby to use index based on group by key(s). bf_result = bf_series.to_pandas() - assert_series_equal_ignoring_order( + assert_series_equal( pd_series, bf_result, ) @@ -1388,7 +1383,7 @@ def test_groupby_prod(scalars_dfs): ) # TODO(swast): Update groupby to use index based on group by key(s). bf_result = bf_series.to_pandas() - assert_series_equal_ignoring_order( + assert_series_equal( pd_series, bf_result, ) @@ -1598,7 +1593,7 @@ def test_head(scalars_dfs): bf_result = scalars_df["string_col"].head(2).to_pandas() pd_result = scalars_pandas_df["string_col"].head(2) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -1613,7 +1608,7 @@ def test_tail(scalars_dfs): bf_result = scalars_df["string_col"].tail(2).to_pandas() pd_result = scalars_pandas_df["string_col"].tail(2) - assert_series_equal_ignoring_order( + assert_series_equal( pd_result, bf_result, ) @@ -2081,11 +2076,7 @@ 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()) # Ignore ordering as pandas order differently depending on version - assert_series_equal_ignoring_order( - bf_result, - pd_result, - check_names=False, - ) + assert_series_equal(bf_result, pd_result, check_names=False, ignore_order=True) def test_series_filter_like(scalars_df_index, scalars_pandas_df_index): @@ -2213,21 +2204,25 @@ def test_where_with_default(scalars_df_index, scalars_pandas_df_index): ) -def test_clip(scalars_df_index, scalars_pandas_df_index): +@pytest.mark.parametrize( + ("ordered"), + [ + (True), + (False), + ], +) +def test_clip(scalars_df_index, scalars_pandas_df_index, ordered): 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).to_pandas() + bf_result = col_bf.clip(lower_bf, upper_bf).to_pandas(ordered=ordered) 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, - ) + assert_series_equal(bf_result, pd_result, ignore_order=not ordered) def test_clip_filtered_two_sided(scalars_df_index, scalars_pandas_df_index): @@ -2298,7 +2293,7 @@ def test_to_frame(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_to_json(scalars_df_index, scalars_pandas_df_index): @@ -2466,7 +2461,7 @@ def test_mask_default_value(scalars_dfs): 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) + assert_pandas_df_equal(bf_result, pd_result) def test_mask_custom_value(scalars_dfs): @@ -2484,7 +2479,7 @@ def test_mask_custom_value(scalars_dfs): # 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) + assert_pandas_df_equal(bf_result, pd_result) @pytest.mark.parametrize( @@ -2581,7 +2576,7 @@ def test_loc_bool_series_default_index( scalars_pandas_df_default_index.bool_col ] - assert_pandas_df_equal_ignore_ordering( + assert_pandas_df_equal( bf_result.to_frame(), pd_result.to_frame(), ) diff --git a/tests/system/utils.py b/tests/system/utils.py index e2daf3b8bf0..f7831972b83 100644 --- a/tests/system/utils.py +++ b/tests/system/utils.py @@ -21,29 +21,33 @@ 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() +def assert_pandas_df_equal(df0, df1, ignore_order: bool = False, **kwargs): + if ignore_order: + # 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() +def assert_series_equal( + left: pd.Series, right: pd.Series, ignore_order: bool = False, **kwargs +): + if ignore_order: + 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) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index d9672b2635b..623448b3aa0 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -49,7 +49,7 @@ def test_arrayvalue_constructor_from_ibis_table_adds_all_columns(): ordering=ordering, hidden_ordering_columns=(), ) - assert actual.compile()._table is ibis_table + assert actual._compile_ordered()._table is ibis_table assert len(actual.column_ids) == 3 @@ -83,7 +83,7 @@ def test_arrayvalue_with_get_column(): ), total_ordering_columns=["col1"], ) - col1 = value.compile()._get_ibis_column("col1") + col1 = value._compile_ordered()._get_ibis_column("col1") assert isinstance(col1, ibis_types.Value) assert col1.get_name() == "col1" assert col1.type().is_int64() @@ -100,7 +100,7 @@ def test_arrayvalues_to_ibis_expr_with_get_column(): ), total_ordering_columns=["col1"], ) - expr = value.compile()._get_ibis_column("col1") + expr = value._compile_ordered()._get_ibis_column("col1") assert expr.get_name() == "col1" assert expr.type().is_int64() @@ -117,7 +117,7 @@ def test_arrayvalues_to_ibis_expr_with_concat(): total_ordering_columns=["col1"], ) expr = value.concat([value]) - actual = expr.compile()._to_ibis_expr("unordered") + actual = expr._compile_ordered()._to_ibis_expr(ordering_mode="unordered") assert len(actual.columns) == 3 # TODO(ashleyxu, b/299631930): test out the union expression assert actual.columns[0] == "column_0" @@ -136,8 +136,8 @@ def test_arrayvalues_to_ibis_expr_with_project_unary_op(): ), total_ordering_columns=["col1"], ) - expr = value.project_unary_op("col1", ops.AsTypeOp("string")).compile() - assert value.compile().columns[0].type().is_int64() + expr = value.project_unary_op("col1", ops.AsTypeOp("string"))._compile_ordered() + assert value._compile_ordered().columns[0].type().is_int64() assert expr.columns[0].type().is_string() @@ -152,9 +152,11 @@ def test_arrayvalues_to_ibis_expr_with_project_binary_op(): ), total_ordering_columns=["col1"], ) - expr = value.project_binary_op("col2", "col3", ops.add_op, "col4").compile() + expr = value.project_binary_op( + "col2", "col3", ops.add_op, "col4" + )._compile_ordered() assert expr.columns[3].type().is_float64() - actual = expr._to_ibis_expr("unordered") + actual = expr._to_ibis_expr(ordering_mode="unordered") assert len(expr.columns) == 4 assert actual.columns[3] == "col4" @@ -173,9 +175,9 @@ def test_arrayvalues_to_ibis_expr_with_project_ternary_op(): ) expr = value.project_ternary_op( "col2", "col3", "col4", ops.where_op, "col5" - ).compile() + )._compile_ordered() assert expr.columns[4].type().is_float64() - actual = expr._to_ibis_expr("unordered") + actual = expr._to_ibis_expr(ordering_mode="unordered") assert len(expr.columns) == 5 assert actual.columns[4] == "col5" @@ -195,8 +197,8 @@ def test_arrayvalue_to_ibis_expr_with_aggregate(): aggregations=(("col1", agg_ops.sum_op, "col4"),), by_column_ids=["col1"], dropna=False, - ).compile() - actual = expr._to_ibis_expr("unordered") + )._compile_ordered() + actual = expr._to_ibis_expr(ordering_mode="unordered") assert len(expr.columns) == 2 assert actual.columns[0] == "col1" assert actual.columns[1] == "col4" @@ -214,8 +216,10 @@ def test_arrayvalue_to_ibis_expr_with_corr_aggregate(): ), total_ordering_columns=["col1"], ) - expr = value.corr_aggregate(corr_aggregations=[("col1", "col3", "col4")]).compile() - actual = expr._to_ibis_expr("unordered") + expr = value.corr_aggregate( + corr_aggregations=[("col1", "col3", "col4")] + )._compile_ordered() + actual = expr._to_ibis_expr(ordering_mode="unordered") assert len(expr.columns) == 1 assert actual.columns[0] == "col4" assert expr.columns[0].type().is_float64() From 30b8ba2409f212631189bb2494baa8a34086bd6e Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Thu, 9 Nov 2023 00:50:16 +0000 Subject: [PATCH 147/197] build: bump dependencies in .kokoro/requirements.txt (#189) Source-Link: https://togithub.com/googleapis/synthtool/commit/febacccc98d6d224aff9d0bd0373bb5a4cd5969c Post-Processor: gcr.io/cloud-devrel-public-resources/owlbot-python:latest@sha256:caffe0a9277daeccc4d1de5c9b55ebba0901b57c2f713ec9c876b0d4ec064f61 --- .github/.OwlBot.lock.yaml | 4 +- .kokoro/requirements.txt | 532 ++++++++++++++++++++------------------ 2 files changed, 277 insertions(+), 259 deletions(-) diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index ec696b558c3..453b540c1e5 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:30470597773378105e239b59fce8eb27cc97375580d592699206d17d117143d0 -# created: 2023-11-03T00:57:07.335914631Z + digest: sha256:caffe0a9277daeccc4d1de5c9b55ebba0901b57c2f713ec9c876b0d4ec064f61 +# created: 2023-11-08T19:46:45.022803742Z diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index 16170d0ca7b..8957e21104e 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -4,91 +4,75 @@ # # pip-compile --allow-unsafe --generate-hashes requirements.in # -argcomplete==2.0.0 \ - --hash=sha256:6372ad78c89d662035101418ae253668445b391755cfe94ea52f1b9d22425b20 \ - --hash=sha256:cffa11ea77999bb0dd27bb25ff6dc142a6796142f68d45b1a26b11f58724561e +argcomplete==3.1.4 \ + --hash=sha256:72558ba729e4c468572609817226fb0a6e7e9a0a7d477b882be168c0b4a62b94 \ + --hash=sha256:fbe56f8cda08aa9a04b307d8482ea703e96a6a801611acb4be9bf3942017989f # via nox -attrs==22.1.0 \ - --hash=sha256:29adc2665447e5191d0e7c568fde78b21f9672d344281d0c6e1ab085429b22b6 \ - --hash=sha256:86efa402f67bf2df34f51a335487cf46b1ec130d02b8d39fd248abfd30da551c +attrs==23.1.0 \ + --hash=sha256:1f28b4522cdc2fb4256ac1a020c78acf9cba2c6b461ccd2c126f3aa8e8335d04 \ + --hash=sha256:6279836d581513a26f1bf235f9acd333bc9115683f14f7e8fae46c98fc50e015 # via gcp-releasetool -bleach==5.0.1 \ - --hash=sha256:085f7f33c15bd408dd9b17a4ad77c577db66d76203e5984b1bd59baeee948b2a \ - --hash=sha256:0d03255c47eb9bd2f26aa9bb7f2107732e7e8fe195ca2f64709fcf3b0a4a085c - # via readme-renderer -cachetools==5.2.0 \ - --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \ - --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db +cachetools==5.3.2 \ + --hash=sha256:086ee420196f7b2ab9ca2db2520aca326318b68fe5ba8bc4d49cca91add450f2 \ + --hash=sha256:861f35a13a451f94e301ce2bec7cac63e881232ccce7ed67fab9b5df4d3beaa1 # via google-auth certifi==2023.7.22 \ --hash=sha256:539cc1d13202e33ca466e88b2807e29f4c13049d6d87031a3c110744495cb082 \ --hash=sha256:92d6037539857d8206b8f6ae472e8b77db8058fec5937a1ef3f54304089edbb9 # 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 +cffi==1.16.0 \ + --hash=sha256:0c9ef6ff37e974b73c25eecc13952c55bceed9112be2d9d938ded8e856138bcc \ + --hash=sha256:131fd094d1065b19540c3d72594260f118b231090295d8c34e19a7bbcf2e860a \ + --hash=sha256:1b8ebc27c014c59692bb2664c7d13ce7a6e9a629be20e54e7271fa696ff2b417 \ + --hash=sha256:2c56b361916f390cd758a57f2e16233eb4f64bcbeee88a4881ea90fca14dc6ab \ + --hash=sha256:2d92b25dbf6cae33f65005baf472d2c245c050b1ce709cc4588cdcdd5495b520 \ + --hash=sha256:31d13b0f99e0836b7ff893d37af07366ebc90b678b6664c955b54561fc36ef36 \ + --hash=sha256:32c68ef735dbe5857c810328cb2481e24722a59a2003018885514d4c09af9743 \ + --hash=sha256:3686dffb02459559c74dd3d81748269ffb0eb027c39a6fc99502de37d501faa8 \ + --hash=sha256:582215a0e9adbe0e379761260553ba11c58943e4bbe9c36430c4ca6ac74b15ed \ + --hash=sha256:5b50bf3f55561dac5438f8e70bfcdfd74543fd60df5fa5f62d94e5867deca684 \ + --hash=sha256:5bf44d66cdf9e893637896c7faa22298baebcd18d1ddb6d2626a6e39793a1d56 \ + --hash=sha256:6602bc8dc6f3a9e02b6c22c4fc1e47aa50f8f8e6d3f78a5e16ac33ef5fefa324 \ + --hash=sha256:673739cb539f8cdaa07d92d02efa93c9ccf87e345b9a0b556e3ecc666718468d \ + --hash=sha256:68678abf380b42ce21a5f2abde8efee05c114c2fdb2e9eef2efdb0257fba1235 \ + --hash=sha256:68e7c44931cc171c54ccb702482e9fc723192e88d25a0e133edd7aff8fcd1f6e \ + --hash=sha256:6b3d6606d369fc1da4fd8c357d026317fbb9c9b75d36dc16e90e84c26854b088 \ + --hash=sha256:748dcd1e3d3d7cd5443ef03ce8685043294ad6bd7c02a38d1bd367cfd968e000 \ + --hash=sha256:7651c50c8c5ef7bdb41108b7b8c5a83013bfaa8a935590c5d74627c047a583c7 \ + --hash=sha256:7b78010e7b97fef4bee1e896df8a4bbb6712b7f05b7ef630f9d1da00f6444d2e \ + --hash=sha256:7e61e3e4fa664a8588aa25c883eab612a188c725755afff6289454d6362b9673 \ + --hash=sha256:80876338e19c951fdfed6198e70bc88f1c9758b94578d5a7c4c91a87af3cf31c \ + --hash=sha256:8895613bcc094d4a1b2dbe179d88d7fb4a15cee43c052e8885783fac397d91fe \ + --hash=sha256:88e2b3c14bdb32e440be531ade29d3c50a1a59cd4e51b1dd8b0865c54ea5d2e2 \ + --hash=sha256:8f8e709127c6c77446a8c0a8c8bf3c8ee706a06cd44b1e827c3e6a2ee6b8c098 \ + --hash=sha256:9cb4a35b3642fc5c005a6755a5d17c6c8b6bcb6981baf81cea8bfbc8903e8ba8 \ + --hash=sha256:9f90389693731ff1f659e55c7d1640e2ec43ff725cc61b04b2f9c6d8d017df6a \ + --hash=sha256:a09582f178759ee8128d9270cd1344154fd473bb77d94ce0aeb2a93ebf0feaf0 \ + --hash=sha256:a6a14b17d7e17fa0d207ac08642c8820f84f25ce17a442fd15e27ea18d67c59b \ + --hash=sha256:a72e8961a86d19bdb45851d8f1f08b041ea37d2bd8d4fd19903bc3083d80c896 \ + --hash=sha256:abd808f9c129ba2beda4cfc53bde801e5bcf9d6e0f22f095e45327c038bfe68e \ + --hash=sha256:ac0f5edd2360eea2f1daa9e26a41db02dd4b0451b48f7c318e217ee092a213e9 \ + --hash=sha256:b29ebffcf550f9da55bec9e02ad430c992a87e5f512cd63388abb76f1036d8d2 \ + --hash=sha256:b2ca4e77f9f47c55c194982e10f058db063937845bb2b7a86c84a6cfe0aefa8b \ + --hash=sha256:b7be2d771cdba2942e13215c4e340bfd76398e9227ad10402a8767ab1865d2e6 \ + --hash=sha256:b84834d0cf97e7d27dd5b7f3aca7b6e9263c56308ab9dc8aae9784abb774d404 \ + --hash=sha256:b86851a328eedc692acf81fb05444bdf1891747c25af7529e39ddafaf68a4f3f \ + --hash=sha256:bcb3ef43e58665bbda2fb198698fcae6776483e0c4a631aa5647806c25e02cc0 \ + --hash=sha256:c0f31130ebc2d37cdd8e44605fb5fa7ad59049298b3f745c74fa74c62fbfcfc4 \ + --hash=sha256:c6a164aa47843fb1b01e941d385aab7215563bb8816d80ff3a363a9f8448a8dc \ + --hash=sha256:d8a9d3ebe49f084ad71f9269834ceccbf398253c9fac910c4fd7053ff1386936 \ + --hash=sha256:db8e577c19c0fda0beb7e0d4e09e0ba74b1e4c092e0e40bfa12fe05b6f6d75ba \ + --hash=sha256:dc9b18bf40cc75f66f40a7379f6a9513244fe33c0e8aa72e2d56b0196a7ef872 \ + --hash=sha256:e09f3ff613345df5e8c3667da1d918f9149bd623cd9070c983c013792a9a62eb \ + --hash=sha256:e4108df7fe9b707191e55f33efbcb2d81928e10cea45527879a4749cbe472614 \ + --hash=sha256:e6024675e67af929088fda399b2094574609396b1decb609c55fa58b028a32a1 \ + --hash=sha256:e70f54f1796669ef691ca07d046cd81a29cb4deb1e5f942003f401c0c4a2695d \ + --hash=sha256:e715596e683d2ce000574bae5d07bd522c781a822866c20495e52520564f0969 \ + --hash=sha256:e760191dd42581e023a68b758769e2da259b5d52e3103c6060ddc02c9edb8d7b \ + --hash=sha256:ed86a35631f7bfbb28e108dd96773b9d5a6ce4811cf6ea468bb6a359b256b1e4 \ + --hash=sha256:ee07e47c12890ef248766a6e55bd38ebfb2bb8edd4142d56db91b21ea68b7627 \ + --hash=sha256:fa3a0128b152627161ce47201262d3140edb5a5c3da88d73a1b790a959126956 \ + --hash=sha256:fcc8eb6d5902bb1cf6dc4f187ee3ea80a1eba0a89aba40a5cb20a5087d961357 # via cryptography charset-normalizer==2.1.1 \ --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845 \ @@ -109,78 +93,74 @@ colorlog==6.7.0 \ # via # gcp-docuploader # nox -commonmark==0.9.1 \ - --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ - --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 - # via rich -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 +cryptography==41.0.5 \ + --hash=sha256:0c327cac00f082013c7c9fb6c46b7cc9fa3c288ca702c74773968173bda421bf \ + --hash=sha256:0d2a6a598847c46e3e321a7aef8af1436f11c27f1254933746304ff014664d84 \ + --hash=sha256:227ec057cd32a41c6651701abc0328135e472ed450f47c2766f23267b792a88e \ + --hash=sha256:22892cc830d8b2c89ea60148227631bb96a7da0c1b722f2aac8824b1b7c0b6b8 \ + --hash=sha256:392cb88b597247177172e02da6b7a63deeff1937fa6fec3bbf902ebd75d97ec7 \ + --hash=sha256:3be3ca726e1572517d2bef99a818378bbcf7d7799d5372a46c79c29eb8d166c1 \ + --hash=sha256:573eb7128cbca75f9157dcde974781209463ce56b5804983e11a1c462f0f4e88 \ + --hash=sha256:580afc7b7216deeb87a098ef0674d6ee34ab55993140838b14c9b83312b37b86 \ + --hash=sha256:5a70187954ba7292c7876734183e810b728b4f3965fbe571421cb2434d279179 \ + --hash=sha256:73801ac9736741f220e20435f84ecec75ed70eda90f781a148f1bad546963d81 \ + --hash=sha256:7d208c21e47940369accfc9e85f0de7693d9a5d843c2509b3846b2db170dfd20 \ + --hash=sha256:8254962e6ba1f4d2090c44daf50a547cd5f0bf446dc658a8e5f8156cae0d8548 \ + --hash=sha256:88417bff20162f635f24f849ab182b092697922088b477a7abd6664ddd82291d \ + --hash=sha256:a48e74dad1fb349f3dc1d449ed88e0017d792997a7ad2ec9587ed17405667e6d \ + --hash=sha256:b948e09fe5fb18517d99994184854ebd50b57248736fd4c720ad540560174ec5 \ + --hash=sha256:c707f7afd813478e2019ae32a7c49cd932dd60ab2d2a93e796f68236b7e1fbf1 \ + --hash=sha256:d38e6031e113b7421db1de0c1b1f7739564a88f1684c6b89234fbf6c11b75147 \ + --hash=sha256:d3977f0e276f6f5bf245c403156673db103283266601405376f075c849a0b936 \ + --hash=sha256:da6a0ff8f1016ccc7477e6339e1d50ce5f59b88905585f77193ebd5068f1e797 \ + --hash=sha256:e270c04f4d9b5671ebcc792b3ba5d4488bf7c42c3c241a3748e2599776f29696 \ + --hash=sha256:e886098619d3815e0ad5790c973afeee2c0e6e04b4da90b88e6bd06e2a0b1b72 \ + --hash=sha256:ec3b055ff8f1dce8e6ef28f626e0972981475173d7973d63f271b29c8a2897da \ + --hash=sha256:fba1e91467c65fe64a82c689dc6cf58151158993b13eb7a7f3f4b7f395636723 # via # gcp-releasetool # secretstorage -distlib==0.3.6 \ - --hash=sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46 \ - --hash=sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e +distlib==0.3.7 \ + --hash=sha256:2e24928bc811348f0feb63014e97aaae3037f2cf48712d51ae61df7fd6075057 \ + --hash=sha256:9dafe54b34a028eafd95039d5e5d4851a13734540f1331060d31c9916e7147a8 # via virtualenv -docutils==0.19 \ - --hash=sha256:33995a6753c30b7f577febfc2c50411fec6aac7f7ffeb7c4cfe5991072dcf9e6 \ - --hash=sha256:5e1de4d849fee02c63b040a4a3fd567f4ab104defd8a5511fbbc24a8a017efbc +docutils==0.20.1 \ + --hash=sha256:96f387a2c5562db4476f09f13bbab2192e764cac08ebbf3a34a95d9b1e4a59d6 \ + --hash=sha256:f08a4e276c3a1583a86dce3e34aba3fe04d02bba2dd51ed16106244e8a923e3b # via readme-renderer -filelock==3.8.0 \ - --hash=sha256:55447caa666f2198c5b6b13a26d2084d26fa5b115c00d065664b2124680c4edc \ - --hash=sha256:617eb4e5eedc82fc5f47b6d61e4d11cb837c56cb4544e39081099fa17ad109d4 +filelock==3.13.1 \ + --hash=sha256:521f5f56c50f8426f5e03ad3b281b490a87ef15bc6c526f168290f0c7148d44e \ + --hash=sha256:57dbda9b35157b05fb3e58ee91448612eb674172fab98ee235ccb0b5bee19a1c # via virtualenv -gcp-docuploader==0.6.4 \ - --hash=sha256:01486419e24633af78fd0167db74a2763974765ee8078ca6eb6964d0ebd388af \ - --hash=sha256:70861190c123d907b3b067da896265ead2eeb9263969d6955c9e0bb091b5ccbf +gcp-docuploader==0.6.5 \ + --hash=sha256:30221d4ac3e5a2b9c69aa52fdbef68cc3f27d0e6d0d90e220fc024584b8d2318 \ + --hash=sha256:b7458ef93f605b9d46a4bf3a8dc1755dad1f31d030c8679edf304e343b347eea # via -r requirements.in -gcp-releasetool==1.10.5 \ - --hash=sha256:174b7b102d704b254f2a26a3eda2c684fd3543320ec239baf771542a2e58e109 \ - --hash=sha256:e29d29927fe2ca493105a82958c6873bb2b90d503acac56be2c229e74de0eec9 +gcp-releasetool==1.16.0 \ + --hash=sha256:27bf19d2e87aaa884096ff941aa3c592c482be3d6a2bfe6f06afafa6af2353e3 \ + --hash=sha256:a316b197a543fd036209d0caba7a8eb4d236d8e65381c80cbc6d7efaa7606d63 # via -r requirements.in -google-api-core==2.10.2 \ - --hash=sha256:10c06f7739fe57781f87523375e8e1a3a4674bf6392cd6131a3222182b971320 \ - --hash=sha256:34f24bd1d5f72a8c4519773d99ca6bf080a6c4e041b4e9f024fe230191dda62e +google-api-core==2.12.0 \ + --hash=sha256:c22e01b1e3c4dcd90998494879612c38d0a3411d1f7b679eb89e2abe3ce1f553 \ + --hash=sha256:ec6054f7d64ad13b41e43d96f735acbd763b0f3b695dabaa2d579673f6a6e160 # via # google-cloud-core # google-cloud-storage -google-auth==2.14.1 \ - --hash=sha256:ccaa901f31ad5cbb562615eb8b664b3dd0bf5404a67618e642307f00613eda4d \ - --hash=sha256:f5d8701633bebc12e0deea4df8abd8aff31c28b355360597f7f2ee60f2e4d016 +google-auth==2.23.4 \ + --hash=sha256:79905d6b1652187def79d491d6e23d0cbb3a21d3c7ba0dbaa9c8a01906b13ff3 \ + --hash=sha256:d4bbc92fe4b8bfd2f3e8d88e5ba7085935da208ee38a134fc280e7ce682a05f2 # via # gcp-releasetool # google-api-core # google-cloud-core # google-cloud-storage -google-cloud-core==2.3.2 \ - --hash=sha256:8417acf6466be2fa85123441696c4badda48db314c607cf1e5d543fa8bdc22fe \ - --hash=sha256:b9529ee7047fd8d4bf4a2182de619154240df17fbe60ead399078c1ae152af9a +google-cloud-core==2.3.3 \ + --hash=sha256:37b80273c8d7eee1ae816b3a20ae43585ea50506cb0e60f3cf5be5f87f1373cb \ + --hash=sha256:fbd11cad3e98a7e5b0343dc07cb1039a5ffd7a5bb96e1f1e27cee4bda4a90863 # via google-cloud-storage -google-cloud-storage==2.6.0 \ - --hash=sha256:104ca28ae61243b637f2f01455cc8a05e8f15a2a18ced96cb587241cdd3820f5 \ - --hash=sha256:4ad0415ff61abdd8bb2ae81c1f8f7ec7d91a1011613f2db87c614c550f97bfe9 +google-cloud-storage==2.13.0 \ + --hash=sha256:ab0bf2e1780a1b74cf17fccb13788070b729f50c252f0c94ada2aae0ca95437d \ + --hash=sha256:f62dc4c7b6cd4360d072e3deb28035fbdad491ac3d9b0b1815a12daea10f37c7 # via gcp-docuploader google-crc32c==1.5.0 \ --hash=sha256:024894d9d3cfbc5943f8f230e23950cd4906b2fe004c72e29b209420a1e6b05a \ @@ -251,29 +231,31 @@ google-crc32c==1.5.0 \ --hash=sha256:f583edb943cf2e09c60441b910d6a20b4d9d626c75a36c8fcac01a6c96c01183 \ --hash=sha256:fd8536e902db7e365f49e7d9029283403974ccf29b13fc7028b97e2295b33556 \ --hash=sha256:fe70e325aa68fa4b5edf7d1a4b6f691eb04bbccac0ace68e34820d283b5f80d4 - # via google-resumable-media -google-resumable-media==2.4.0 \ - --hash=sha256:2aa004c16d295c8f6c33b2b4788ba59d366677c0a25ae7382436cb30f776deaa \ - --hash=sha256:8d5518502f92b9ecc84ac46779bd4f09694ecb3ba38a3e7ca737a86d15cbca1f + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.6.0 \ + --hash=sha256:972852f6c65f933e15a4a210c2b96930763b47197cdf4aa5f5bea435efb626e7 \ + --hash=sha256:fc03d344381970f79eebb632a3c18bb1828593a2dc5572b5f90115ef7d11e81b # via google-cloud-storage -googleapis-common-protos==1.57.0 \ - --hash=sha256:27a849d6205838fb6cc3c1c21cb9800707a661bb21c6ce7fb13e99eb1f8a0c46 \ - --hash=sha256:a9f4a1d7f6d9809657b7f1316a1aa527f6664891531bcfcc13b6696e685f443c +googleapis-common-protos==1.61.0 \ + --hash=sha256:22f1915393bb3245343f6efe87f6fe868532efc12aa26b391b15132e1279f1c0 \ + --hash=sha256:8a64866a97f6304a7179873a465d6eee97b7a24ec6cfd78e0f575e96b821240b # via google-api-core idna==3.4 \ --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \ --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2 # via requests -importlib-metadata==5.0.0 \ - --hash=sha256:da31db32b304314d044d3c12c79bd59e307889b287ad12ff387b3500835fc2ab \ - --hash=sha256:ddb0e35065e8938f867ed4928d0ae5bf2a53b7773871bfe6bcc7e4fcdc7dea43 +importlib-metadata==6.8.0 \ + --hash=sha256:3ebb78df84a805d7698245025b975d9d67053cd94c79245ba4b3eb694abe68bb \ + --hash=sha256:dbace7892d8c0c4ac1ad096662232f831d4e64f4c4545bd53016a3e9d4654743 # via # -r requirements.in # keyring # twine -jaraco-classes==3.2.3 \ - --hash=sha256:2353de3288bc6b82120752201c6b1c1a14b058267fa424ed5ce5984e3b922158 \ - --hash=sha256:89559fa5c1d3c34eff6f631ad80bb21f378dbcbb35dd161fd2c6b93f5be2f98a +jaraco-classes==3.3.0 \ + --hash=sha256:10afa92b6743f25c0cf5f37c6bb6e18e2c5bb84a16527ccfc0040ea377e7aaeb \ + --hash=sha256:c063dd08e89217cee02c8d5e5ec560f2c8ce6cdc2fcdc2e68f7b2e5547ed3621 # via keyring jeepney==0.8.0 \ --hash=sha256:5efe48d255973902f6badc3ce55e2aa6c5c3b3bc642059ef3a91247bcfcc5806 \ @@ -285,75 +267,121 @@ jinja2==3.1.2 \ --hash=sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852 \ --hash=sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61 # via gcp-releasetool -keyring==23.11.0 \ - --hash=sha256:3dd30011d555f1345dec2c262f0153f2f0ca6bca041fb1dc4588349bb4c0ac1e \ - --hash=sha256:ad192263e2cdd5f12875dedc2da13534359a7e760e77f8d04b50968a821c2361 +keyring==24.2.0 \ + --hash=sha256:4901caaf597bfd3bbd78c9a0c7c4c29fcd8310dab2cffefe749e916b6527acd6 \ + --hash=sha256:ca0746a19ec421219f4d713f848fa297a661a8a8c1504867e55bfb5e09091509 # via # gcp-releasetool # twine -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 +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:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c \ + --hash=sha256:1577735524cdad32f9f694208aa75e422adba74f1baee7551620e43a3141f559 \ + --hash=sha256:1b40069d487e7edb2676d3fbdb2b0829ffa2cd63a2ec26c4938b2d34391b4ecc \ + --hash=sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb \ + --hash=sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939 \ + --hash=sha256:282c2cb35b5b673bbcadb33a585408104df04f14b2d9b01d4c345a3b92861c2c \ + --hash=sha256:2c1b19b3aaacc6e57b7e25710ff571c24d6c3613a45e905b1fde04d691b98ee0 \ + --hash=sha256:2ef12179d3a291be237280175b542c07a36e7f60718296278d8593d21ca937d4 \ + --hash=sha256:338ae27d6b8745585f87218a3f23f1512dbf52c26c28e322dbe54bcede54ccb9 \ + --hash=sha256:3c0fae6c3be832a0a0473ac912810b2877c8cb9d76ca48de1ed31e1c68386575 \ + --hash=sha256:3fd4abcb888d15a94f32b75d8fd18ee162ca0c064f35b11134be77050296d6ba \ + --hash=sha256:42de32b22b6b804f42c5d98be4f7e5e977ecdd9ee9b660fda1a3edf03b11792d \ + --hash=sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd \ + --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:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007 \ + --hash=sha256:787003c0ddb00500e49a10f2844fac87aa6ce977b90b0feaaf9de23c22508b24 \ + --hash=sha256:7ef3cb2ebbf91e330e3bb937efada0edd9003683db6b57bb108c4001f37a02ea \ + --hash=sha256:8023faf4e01efadfa183e863fefde0046de576c6f14659e8782065bcece22198 \ + --hash=sha256:8758846a7e80910096950b67071243da3e5a20ed2546e6392603c096778d48e0 \ + --hash=sha256:8afafd99945ead6e075b973fefa56379c5b5c53fd8937dad92c662da5d8fd5ee \ + --hash=sha256:8c41976a29d078bb235fea9b2ecd3da465df42a562910f9022f1a03107bd02be \ + --hash=sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2 \ + --hash=sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1 \ + --hash=sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707 \ + --hash=sha256:962f82a3086483f5e5f64dbad880d31038b698494799b097bc59c2edf392fce6 \ + --hash=sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c \ + --hash=sha256:9dcdfd0eaf283af041973bff14a2e143b8bd64e069f4c383416ecd79a81aab58 \ + --hash=sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823 \ + --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:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc \ + --hash=sha256:fec21693218efe39aa7f8599346e90c705afa52c5b31ae019b2e57e8f6542bb2 \ + --hash=sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11 # via jinja2 -more-itertools==9.0.0 \ - --hash=sha256:250e83d7e81d0c87ca6bd942e6aeab8cc9daa6096d12c5308f3f92fa5e5c1f41 \ - --hash=sha256:5a6257e40878ef0520b1803990e3e22303a41b5714006c32a3fd8304b26ea1ab +mdurl==0.1.2 \ + --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \ + --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba + # via markdown-it-py +more-itertools==10.1.0 \ + --hash=sha256:626c369fa0eb37bac0291bce8259b332fd59ac792fa5497b59837309cd5b114a \ + --hash=sha256:64e0735fcfdc6f3464ea133afe8ea4483b1c5fe3a3d69852e6503b43a0b222e6 # via jaraco-classes -nox==2022.11.21 \ - --hash=sha256:0e41a990e290e274cb205a976c4c97ee3c5234441a8132c8c3fd9ea3c22149eb \ - --hash=sha256:e21c31de0711d1274ca585a2c5fde36b1aa962005ba8e9322bf5eeed16dcd684 +nh3==0.2.14 \ + --hash=sha256:116c9515937f94f0057ef50ebcbcc10600860065953ba56f14473ff706371873 \ + --hash=sha256:18415df36db9b001f71a42a3a5395db79cf23d556996090d293764436e98e8ad \ + --hash=sha256:203cac86e313cf6486704d0ec620a992c8bc164c86d3a4fd3d761dd552d839b5 \ + --hash=sha256:2b0be5c792bd43d0abef8ca39dd8acb3c0611052ce466d0401d51ea0d9aa7525 \ + --hash=sha256:377aaf6a9e7c63962f367158d808c6a1344e2b4f83d071c43fbd631b75c4f0b2 \ + --hash=sha256:525846c56c2bcd376f5eaee76063ebf33cf1e620c1498b2a40107f60cfc6054e \ + --hash=sha256:5529a3bf99402c34056576d80ae5547123f1078da76aa99e8ed79e44fa67282d \ + --hash=sha256:7771d43222b639a4cd9e341f870cee336b9d886de1ad9bec8dddab22fe1de450 \ + --hash=sha256:88c753efbcdfc2644a5012938c6b9753f1c64a5723a67f0301ca43e7b85dcf0e \ + --hash=sha256:93a943cfd3e33bd03f77b97baa11990148687877b74193bf777956b67054dcc6 \ + --hash=sha256:9be2f68fb9a40d8440cbf34cbf40758aa7f6093160bfc7fb018cce8e424f0c3a \ + --hash=sha256:a0c509894fd4dccdff557068e5074999ae3b75f4c5a2d6fb5415e782e25679c4 \ + --hash=sha256:ac8056e937f264995a82bf0053ca898a1cb1c9efc7cd68fa07fe0060734df7e4 \ + --hash=sha256:aed56a86daa43966dd790ba86d4b810b219f75b4bb737461b6886ce2bde38fd6 \ + --hash=sha256:e8986f1dd3221d1e741fda0a12eaa4a273f1d80a35e31a1ffe579e7c621d069e \ + --hash=sha256:f99212a81c62b5f22f9e7c3e347aa00491114a5647e1f13bbebd79c3e5f08d75 + # via readme-renderer +nox==2023.4.22 \ + --hash=sha256:0b1adc619c58ab4fa57d6ab2e7823fe47a32e70202f287d78474adcc7bda1891 \ + --hash=sha256:46c0560b0dc609d7d967dc99e22cb463d3c4caf54a5fda735d6c11b5177e3a9f # via -r requirements.in -packaging==21.3 \ - --hash=sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb \ - --hash=sha256:ef103e05f519cdc783ae24ea4e2e0f508a9c99b2d4969652eed6a2e1ea5bd522 +packaging==23.2 \ + --hash=sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5 \ + --hash=sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7 # via # gcp-releasetool # nox -pkginfo==1.8.3 \ - --hash=sha256:848865108ec99d4901b2f7e84058b6e7660aae8ae10164e015a6dcf5b242a594 \ - --hash=sha256:a84da4318dd86f870a9447a8c98340aa06216bfc6f2b7bdc4b8766984ae1867c +pkginfo==1.9.6 \ + --hash=sha256:4b7a555a6d5a22169fcc9cf7bfd78d296b0361adad412a346c1226849af5e546 \ + --hash=sha256:8fd5896e8718a4372f0ea9cc9d96f6417c9b986e23a4d116dda26b62cc29d046 # via twine -platformdirs==2.5.4 \ - --hash=sha256:1006647646d80f16130f052404c6b901e80ee4ed6bef6792e1f238a8969106f7 \ - --hash=sha256:af0276409f9a02373d540bf8480021a048711d572745aef4b7842dad245eba10 +platformdirs==3.11.0 \ + --hash=sha256:cf8ee52a3afdb965072dcc652433e0c7e3e40cf5ea1477cd4b3b1d2eb75495b3 \ + --hash=sha256:e9d171d00af68be50e9202731309c4e658fd8bc76f55c11c7dd760d023bda68e # via virtualenv protobuf==3.20.3 \ --hash=sha256:03038ac1cfbc41aa21f6afcbcd357281d7521b4157926f30ebecc8d4ea59dcb7 \ @@ -383,34 +411,30 @@ protobuf==3.20.3 \ # gcp-releasetool # google-api-core # googleapis-common-protos -pyasn1==0.4.8 \ - --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \ - --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba +pyasn1==0.5.0 \ + --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \ + --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde # via # pyasn1-modules # rsa -pyasn1-modules==0.2.8 \ - --hash=sha256:905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e \ - --hash=sha256:a50b808ffeb97cb3601dd25981f6b016cbb3d31fbf57a8b8a87428e6158d0c74 +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.0 \ - --hash=sha256:77a3299119af881904cd5ecd1ac6a66214b6e9bed1f2db16993b54adede64094 \ - --hash=sha256:f7e36cffc4c517fbc252861b9a6e4644ca0e5abadf9a113c72d1358ad09b9500 +pygments==2.16.1 \ + --hash=sha256:13fc09fa63bc8d8671a6d247e1eb303c4b343eaee81d861f3404db2935653692 \ + --hash=sha256:1daff0494820c69bc8941e407aa20f577374ee88364ee10a98fdbe0aece96e29 # via # readme-renderer # rich -pyjwt==2.6.0 \ - --hash=sha256:69285c7e31fc44f68a1feb309e948e0df53259d579295e6cfe2b1792329f05fd \ - --hash=sha256:d83c3d892a77bbb74d3e1a2cfa90afaadb60945205d1095d9221f04466f64c14 +pyjwt==2.8.0 \ + --hash=sha256:57e28d156e3d5c10088e0c68abb90bfac3df82b40a71bd0daa20c65ccd5c23de \ + --hash=sha256:59127c392cc44c2da5bb3192169a91f429924e17aff6534d70fdc02ab3e04320 # 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 @@ -418,9 +442,9 @@ python-dateutil==2.8.2 \ --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \ --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9 # via gcp-releasetool -readme-renderer==37.3 \ - --hash=sha256:cd653186dfc73055656f090f227f5cb22a046d7f71a841dfa305f55c9a513273 \ - --hash=sha256:f67a16caedfa71eef48a31b39708637a6f4664c4394801a7b0d6432d13907343 +readme-renderer==42.0 \ + --hash=sha256:13d039515c1f24de668e2c93f2e877b9dbe6c6c32328b90a40a49d8b2b85f36d \ + --hash=sha256:2d55489f83be4992fe4454939d1a051c33edbab778e82761d060c9fc6b308cd1 # via twine requests==2.31.0 \ --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ @@ -431,17 +455,17 @@ requests==2.31.0 \ # google-cloud-storage # requests-toolbelt # twine -requests-toolbelt==0.10.1 \ - --hash=sha256:18565aa58116d9951ac39baa288d3adb5b3ff975c4f25eee78555d89e8f247f7 \ - --hash=sha256:62e09f7ff5ccbda92772a29f394a49c3ad6cb181d568b1337626b2abb628a63d +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==12.6.0 \ - --hash=sha256:a4eb26484f2c82589bd9a17c73d32a010b1e29d89f1604cd9bf3a2097b81bb5e \ - --hash=sha256:ba3a3775974105c221d31141f2c116f4fd65c5ceb0698657a11e9f295ec93fd0 +rich==13.6.0 \ + --hash=sha256:2b38e2fe9ca72c9a00170a1a2d20c63c790d0e10ef1fe35eba76e1e7b1d7d245 \ + --hash=sha256:5c14d22737e6d5084ef4771b62d5d4363165b403455a30a1c8ca39dc7b644bef # via twine rsa==4.9 \ --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ @@ -455,43 +479,37 @@ six==1.16.0 \ --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 # via - # bleach # gcp-docuploader - # google-auth # python-dateutil -twine==4.0.1 \ - --hash=sha256:42026c18e394eac3e06693ee52010baa5313e4811d5a11050e7d48436cf41b9e \ - --hash=sha256:96b1cf12f7ae611a4a40b6ae8e9570215daff0611828f5fe1f37a16255ab24a0 +twine==4.0.2 \ + --hash=sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8 \ + --hash=sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8 # via -r requirements.in -typing-extensions==4.4.0 \ - --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \ - --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e +typing-extensions==4.8.0 \ + --hash=sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0 \ + --hash=sha256:df8e4339e9cb77357558cbdbceca33c303714cf861d1eef15e1070055ae8b7ef # via -r requirements.in -urllib3==1.26.18 \ - --hash=sha256:34b97092d7e0a3a8cf7cd10e386f401b3737364026c45e622aa02903dffe0f07 \ - --hash=sha256:f8ecc1bba5667413457c529ab955bf8c67b45db799d159066261719e328580a0 +urllib3==2.0.7 \ + --hash=sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84 \ + --hash=sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e # via # requests # twine -virtualenv==20.16.7 \ - --hash=sha256:8691e3ff9387f743e00f6bb20f70121f5e4f596cae754531f2b3b3a1b1ac696e \ - --hash=sha256:efd66b00386fdb7dbe4822d172303f40cd05e50e01740b19ea42425cbe653e29 +virtualenv==20.24.6 \ + --hash=sha256:02ece4f56fbf939dbbc33c0715159951d6bf14aaf5457b092e4548e1382455af \ + --hash=sha256:520d056652454c5098a00c0f073611ccbea4c79089331f60bf9d7ba247bb7381 # via nox -webencodings==0.5.1 \ - --hash=sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78 \ - --hash=sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923 - # via bleach -wheel==0.38.4 \ - --hash=sha256:965f5259b566725405b05e7cf774052044b1ed30119b5d586b2703aafe8719ac \ - --hash=sha256:b60533f3f5d530e971d6737ca6d58681ee434818fab630c83a734bb10c083ce8 +wheel==0.41.3 \ + --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \ + --hash=sha256:4d4987ce51a49370ea65c0bfd2234e8ce80a12780820d9dc462597a6e60d0841 # via -r requirements.in -zipp==3.10.0 \ - --hash=sha256:4fcb6f278987a6605757302a6e40e896257570d11c51628968ccb2a47e80c6c1 \ - --hash=sha256:7a7262fd930bd3e36c50b9a64897aec3fafff3dfdeec9623ae22b40e93f99bb8 +zipp==3.17.0 \ + --hash=sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31 \ + --hash=sha256:84e64a1c28cf7e91ed2078bb8cc8c259cb19b76942096c8d7b84947690cabaf0 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: -setuptools==65.5.1 \ - --hash=sha256:d0b9a8433464d5800cbe05094acf5c6d52a91bfac9b52bcfc4d41382be5d5d31 \ - --hash=sha256:e197a19aa8ec9722928f2206f8de752def0e4c9fc6953527360d1c36d94ddb2f +setuptools==68.2.2 \ + --hash=sha256:4ac1475276d2f1c48684874089fefcd83bd7162ddaafb81fac866ba0db282a87 \ + --hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a # via -r requirements.in From d49ae42a379fafd601cc94227e7f8f14b3d5f8c3 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Wed, 8 Nov 2023 17:31:50 -0800 Subject: [PATCH 148/197] docs: add llm kmeans notebook as an included example (#177) * docs: add llm kmeans notebook as an included example --------- Co-authored-by: Henry J Solberg Co-authored-by: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> --- .../bq_dataframes_llm_kmeans.ipynb | 690 ++++++++++++++++++ noxfile.py | 1 + 2 files changed, 691 insertions(+) create mode 100644 notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb new file mode 100644 index 00000000000..46c49552889 --- /dev/null +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -0,0 +1,690 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "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." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Use BigQuery DataFrames to cluster and characterize complaints\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", + "
" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Overview\n", + "\n", + "The goal of this notebook is to demonstrate a comment characterization algorithm for an online business. We will accomplish this using [Google's PaLM 2](https://ai.google/discover/palm2/) and [KMeans clustering](https://en.wikipedia.org/wiki/K-means_clustering) in three steps:\n", + "\n", + "1. Use PaLM2TextEmbeddingGenerator to [generate text embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings) for each of 10000 complaints sent to an online bank. If you're not familiar with what a text embedding is, it's a list of numbers that are like coordinates in an imaginary \"meaning space\" for sentences. (It's like [word embeddings](https://en.wikipedia.org/wiki/Word_embedding), but for more general text.) The important point for our purposes is that similar sentences are close to each other in this imaginary space.\n", + "2. Use KMeans clustering to group together complaints whose text embeddings are near to eachother. This will give us sets of similar complaints, but we don't yet know _why_ these complaints are similar.\n", + "3. Prompt PaLM2TextGenerator in English asking what the difference is between the groups of complaints that we got. Thanks to the power of modern LLMs, the response might give us a very good idea of what these complaints are all about, but remember to [\"understand the limits of your dataset and model.\"](https://ai.google/responsibility/responsible-ai-practices/#:~:text=Understand%20the%20limitations%20of%20your%20dataset%20and%20model)\n", + "\n", + "We will tie these pieces together in Python using BigQuery DataFrames. [Click here](https://cloud.google.com/bigquery/docs/dataframes-quickstart) to learn more about BigQuery DataFrames!" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Dataset\n", + "\n", + "This notebook uses the [CFPB Consumer Complaint Database](https://console.cloud.google.com/marketplace/product/cfpb/complaint-database)." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery (compute)\n", + "* BigQuery ML\n", + "* Generative AI support on Vertex AI\n", + "\n", + "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models), [Generative AI support on Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing#generative_ai_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." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Before you begin\n", + "\n", + "Complete the tasks in this section to set up your environment." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "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,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 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)." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# set your project ID below\n", + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set the project id in gcloud\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "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": {}, + "outputs": [], + "source": [ + "REGION = \"US\" # @param {type: \"string\"}" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you might have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Vertex AI Workbench**\n", + "\n", + "Do nothing, you are already authenticated." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Local JupyterLab instance**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Colab**\n", + "\n", + "Uncomment and run the following cell:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Connect to Vertex AI\n", + "\n", + "In order to use PaLM2TextGenerator, we will need to set up a [cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from google.cloud import bigquery_connection_v1 as bq_connection\n", + "\n", + "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)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "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'" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we are ready to use BigQuery DataFrames!" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "xckgWno6ouHY" + }, + "source": [ + "## Step 1: Text embedding " + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Project Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "R7STCS8xB5d2" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "\n", + "bf.options.bigquery.project = PROJECT_ID\n", + "bf.options.bigquery.location = REGION" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "v6FGschEowht" + }, + "source": [ + "Data Input - read the data from a publicly available BigQuery dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "zDSwoBo1CU3G" + }, + "outputs": [], + "source": [ + "input_df = bf.read_gbq(\"bigquery-public-data.cfpb_complaints.complaint_database\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tYDoaKgJChiq" + }, + "outputs": [], + "source": [ + "issues_df = input_df[[\"consumer_complaint_narrative\"]].dropna()\n", + "issues_df.head(n=5) # View the first five complaints" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Download 10000 complaints to use with PaLM2TextEmbeddingGenerator" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OltYSUEcsSOW" + }, + "outputs": [], + "source": [ + "# Choose 10,000 complaints randomly and store them in a column in a DataFrame\n", + "downsampled_issues_df = issues_df.sample(n=10000)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "Wl2o-NYMoygb" + }, + "source": [ + "Generate the text embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "li38q8FzDDMu" + }, + "outputs": [], + "source": [ + "from bigframes.ml.llm import PaLM2TextEmbeddingGenerator\n", + "\n", + "model = PaLM2TextEmbeddingGenerator() # No connection id needed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cOuSOQ5FDewD" + }, + "outputs": [], + "source": [ + "# Will take ~3 minutes to compute the embeddings\n", + "predicted_embeddings = model.predict(downsampled_issues_df)\n", + "# Notice the lists of numbers that are our text embeddings for each complaint\n", + "predicted_embeddings.head() " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4H_etYfsEOFP" + }, + "outputs": [], + "source": [ + "# Join the complaints with their embeddings in the same DataFrame\n", + "combined_df = downsampled_issues_df.join(predicted_embeddings)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now have the complaints and their text embeddings as two columns in our combined_df. Recall that complaints with numerically similar text embeddings should have similar meanings semantically. We will now group similar complaints together." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "OUZ3NNbzo1Tb" + }, + "source": [ + "## Step 2: KMeans clustering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "AhNTnEC5FRz2" + }, + "outputs": [], + "source": [ + "from bigframes.ml.cluster import KMeans\n", + "\n", + "cluster_model = KMeans(n_clusters=10) # We will divide our complaints into 10 groups" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Perform KMeans clustering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6poSxh-fGJF7" + }, + "outputs": [], + "source": [ + "# Use KMeans clustering to calculate our groups. Will take ~3 minutes.\n", + "cluster_model.fit(combined_df[[\"text_embedding\"]])\n", + "clustered_result = cluster_model.predict(combined_df[[\"text_embedding\"]])\n", + "# Notice the CENTROID_ID column, which is the ID number of the group that\n", + "# each complaint belongs to.\n", + "clustered_result.head(n=5)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Join the group number to the complaints and their text embeddings\n", + "combined_clustered_result = combined_df.join(clustered_result)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Our dataframe combined_clustered_result now has three columns: the complaints, their text embeddings, and an ID from 1-10 (inclusive) indicating which semantically similar group they belong to." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "21rNsFMHo8hO" + }, + "source": [ + "## Step 3: Summarize the complaints" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Build prompts - we will choose just two of our categories and prompt PaLM2TextGenerator to identify their salient characteristics. The prompt is natural language in a python string." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2E7wXM_jGqo6" + }, + "outputs": [], + "source": [ + "# Using bigframes, with syntax identical to pandas,\n", + "# filter out the first and second groups\n", + "cluster_1_result = combined_clustered_result[\n", + " combined_clustered_result[\"CENTROID_ID\"] == 1\n", + "][[\"consumer_complaint_narrative\"]]\n", + "cluster_1_result_pandas = cluster_1_result.head(5).to_pandas()\n", + "\n", + "cluster_2_result = combined_clustered_result[\n", + " combined_clustered_result[\"CENTROID_ID\"] == 2\n", + "][[\"consumer_complaint_narrative\"]]\n", + "cluster_2_result_pandas = cluster_2_result.head(5).to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZNDiueI9IP5e" + }, + "outputs": [], + "source": [ + "# Build plain-text prompts to send to PaLM 2. Use only 5 complaints from each group.\n", + "prompt1 = 'comment list 1:\\n'\n", + "for i in range(5):\n", + " prompt1 += str(i + 1) + '. ' + \\\n", + " cluster_1_result_pandas[\"consumer_complaint_narrative\"].iloc[i] + '\\n'\n", + "\n", + "prompt2 = 'comment list 2:\\n'\n", + "for i in range(5):\n", + " prompt2 += str(i + 1) + '. ' + \\\n", + " cluster_2_result_pandas[\"consumer_complaint_narrative\"].iloc[i] + '\\n'\n", + "\n", + "print(prompt1)\n", + "print(prompt2)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "BfHGJLirzSvH" + }, + "outputs": [], + "source": [ + "# The plain English request we will make of PaLM 2\n", + "prompt = (\n", + " \"Please highlight the most obvious difference between\"\n", + " \"the two lists of comments:\\n\" + prompt1 + prompt2\n", + ")\n", + "print(prompt)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Get a response from PaLM 2 LLM by making a call to Vertex AI using our connection." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "mL5P0_3X04dE" + }, + "outputs": [], + "source": [ + "from bigframes.ml.llm import PaLM2TextGenerator\n", + "\n", + "session = bf.get_global_session()\n", + "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", + "q_a_model = PaLM2TextGenerator(session=session, connection_name=connection)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ICWHsqAW1FNk" + }, + "outputs": [], + "source": [ + "# Make a DataFrame containing only a single row with our prompt for PaLM 2\n", + "df = bf.DataFrame({\"prompt\": [prompt]})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "gB7e1LXU1pst" + }, + "outputs": [], + "source": [ + "# Send the request for PaLM 2 to generate a response to our prompt\n", + "major_difference = q_a_model.predict(df)\n", + "# PaLM 2's response is the only row in the dataframe result \n", + "major_difference[\"ml_generate_text_llm_result\"].iloc[0]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now see PaLM2TextGenerator's characterization of the different comment groups. Thanks for using BigQuery DataFrames!" + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "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.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/noxfile.py b/noxfile.py index 34b055de445..3dd23ba04f6 100644 --- a/noxfile.py +++ b/noxfile.py @@ -609,6 +609,7 @@ def notebook(session): # our test infrastructure. "notebooks/getting_started/getting_started_bq_dataframes.ipynb", "notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb", + "notebooks/generative_ai/bq_dataframes_llm_kmeans.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", From 37914a4077c681881491f5c36d1a9c9f4255e18f Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Wed, 8 Nov 2023 18:08:13 -0800 Subject: [PATCH 149/197] feat: add series.sample (identical to existing dataframe.sample) (#187) We're duplicating some arg-parsing logic here. Discussed briefly with Trevor. This is the case for other methods as well- we might want to add a sharing mechanism for dataframe/series (superclass like pandas?) in the future. The documentation already exists in third_party/core/generic.py, which is actually what prompted this feat/fix. --- bigframes/series.py | 16 ++++++++++++++++ tests/system/small/test_series.py | 27 +++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/bigframes/series.py b/bigframes/series.py index 1952acbf6d5..4fab1fe9439 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1447,6 +1447,22 @@ def map( result_df = self_df.join(map_df, on="series") return result_df[self.name] + def sample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + *, + random_state: Optional[int] = None, + ) -> Series: + 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 Series( + self._block._split(ns=ns, fracs=fracs, random_state=random_state)[0] + ) + def __array_ufunc__( self, ufunc: numpy.ufunc, method: str, *inputs, **kwargs ) -> Series: diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index f59d64fe06f..d9fc23fad01 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2922,3 +2922,30 @@ def test_map_series_input_duplicates_error(scalars_dfs): 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) + + +@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.int64_col.sample(frac=frac, n=n, random_state=random_state) + 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 + assert bf_result.shape[0] == expected_sample_size From c816d843e6f3c5a944cd4395ed0e1e91cec49812 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 9 Nov 2023 22:19:24 +0000 Subject: [PATCH 150/197] docs: Add docstring code samples for `Series.apply` and `DataFrame.map` (#185) * docs: Add docstring code samples for `Series.apply` and `DataFrame.map` * improved docstring with concurrency-safe code samples * Correct indentation of text in code samples --- .../bigframes_vendored/pandas/core/frame.py | 62 ++++++++++++++++- .../bigframes_vendored/pandas/core/series.py | 68 +++++++++++++++++-- 2 files changed, 123 insertions(+), 7 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 10f73d25b77..b35d0f3b2e8 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2159,8 +2159,68 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: In pandas 2.1.0, DataFrame.applymap is deprecated and renamed to DataFrame.map. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Let's use ``reuse=False`` flag to make sure a new ``remote_function`` + is created every time we run the following code, but you can skip it + to potentially reuse a previously deployed ``remote_function`` from + the same user defined function. + + >>> @bpd.remote_function([int], float, reuse=False) + ... def minutes_to_hours(x): + ... return x/60 + + >>> df_minutes = bpd.DataFrame( + ... {"system_minutes" : [0, 30, 60, 90, 120], + ... "user_minutes" : [0, 15, 75, 90, 6]}) + >>> df_minutes + system_minutes user_minutes + 0 0 0 + 1 30 15 + 2 60 75 + 3 90 90 + 4 120 6 + + [5 rows x 2 columns] + + >>> df_hours = df_minutes.map(minutes_to_hours) + >>> df_hours + system_minutes user_minutes + 0 0.0 0.0 + 1 0.5 0.25 + 2 1.0 1.25 + 3 1.5 1.5 + 4 2.0 0.1 + + [5 rows x 2 columns] + + If there are ``NA``/``None`` values in the data, you can ignore + applying the remote function on such values by specifying + ``na_action='ignore'``. + + >>> df_minutes = bpd.DataFrame( + ... { + ... "system_minutes" : [0, 30, 60, None, 90, 120, bpd.NA], + ... "user_minutes" : [0, 15, 75, 90, 6, None, bpd.NA] + ... }, dtype="Int64") + >>> df_hours = df_minutes.map(minutes_to_hours, na_action='ignore') + >>> df_hours + system_minutes user_minutes + 0 0.0 0.0 + 1 0.5 0.25 + 2 1.0 1.25 + 3 1.5 + 4 1.5 0.1 + 5 2.0 + 6 + + [7 rows x 2 columns] + Args: - func: + func (function): Python function wrapped by ``remote_function`` decorator, returns a single value from a single value. na_action (Optional[str], default None): diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index b25b0c75cff..c6d98075f57 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -728,18 +728,74 @@ def apply( func, ) -> DataFrame | Series: """ - Invoke function on values of Series. + Invoke function on values of a Series. - Can be ufunc (a NumPy function that applies to the entire Series) - or a Python function that only works on single values. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Let's use ``reuse=False`` flag to make sure a new ``remote_function`` + is created every time we run the following code, but you can skip it + to potentially reuse a previously deployed ``remote_function`` from + the same user defined function. + + >>> @bpd.remote_function([int], float, reuse=False) + ... def minutes_to_hours(x): + ... return x/60 + + >>> minutes = bpd.Series([0, 30, 60, 90, 120]) + >>> minutes + 0 0 + 1 30 + 2 60 + 3 90 + 4 120 + dtype: Int64 + + >>> hours = minutes.apply(minutes_to_hours) + >>> hours + 0 0.0 + 1 0.5 + 2 1.0 + 3 1.5 + 4 2.0 + dtype: Float64 + + You could turn a user defined function with external package + dependencies into a BigQuery DataFrames remote function. You would + provide the names of the packages via ``packages`` param. + + >>> @bpd.remote_function( + ... [str], + ... str, + ... reuse=False, + ... packages=["cryptography"], + ... ) + ... def get_hash(input): + ... from cryptography.fernet import Fernet + ... + ... # handle missing value + ... if input is None: + ... input = "" + ... + ... key = Fernet.generate_key() + ... f = Fernet(key) + ... return f.encrypt(input.encode()).decode() + + >>> names = bpd.Series(["Alice", "Bob"]) + >>> hashes = names.apply(get_hash) Args: func (function): - Python function or NumPy ufunc to apply. + BigFrames DataFrames ``remote_function`` to apply. The function + should take a scalar and return a scalar. It will be applied to + every element in the ``Series``. Returns: - bigframes.series.Series: If func returns a Series object the result - will be a DataFrame. + bigframes.series.Series: A new Series with values representing the + return value of the ``func`` applied to each element of the original + Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From 9d2e6dc1ae4e11e80da4aabe0daa3a6044137cc6 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 9 Nov 2023 15:00:23 -0800 Subject: [PATCH 151/197] fix: use random table when loading data for `read_csv`, `read_json`, `read_parquet` (#175) 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/session/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 8f9fa377873..bd5845631b2 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -806,7 +806,7 @@ def _read_gbq_with_ordering( def _read_bigquery_load_job( self, filepath_or_buffer: str | IO["bytes"], - table: bigquery.Table, + table: Union[bigquery.Table, bigquery.TableReference], *, job_config: bigquery.LoadJobConfig, index_col: Iterable[str] | str = (), @@ -1042,7 +1042,7 @@ def read_csv( encoding: Optional[str] = None, **kwargs, ) -> dataframe.DataFrame: - table = bigquery.Table(self._create_session_table()) + table = bigframes_io.random_table(self._anonymous_dataset) if engine is not None and engine == "bigquery": if any(param is not None for param in (dtype, names)): @@ -1156,7 +1156,7 @@ def read_parquet( # 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()) + table = bigframes_io.random_table(self._anonymous_dataset) job_config = bigquery.LoadJobConfig() job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED @@ -1179,7 +1179,7 @@ def read_json( engine: Literal["ujson", "pyarrow", "bigquery"] = "ujson", **kwargs, ) -> dataframe.DataFrame: - table = bigquery.Table(self._create_session_table()) + table = bigframes_io.random_table(self._anonymous_dataset) if engine == "bigquery": From 03606cda30eb7645bfd4534460112dcca56b0ab0 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 10 Nov 2023 12:50:14 -0600 Subject: [PATCH 152/197] fix: default to 7 days expiration for `read_csv`, `read_json`, `read_parquet` (#193) 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) Follow-up to https://togithub.com/googleapis/python-bigquery-dataframes/pull/175/files#r1389686556 🦕 --- bigframes/session/__init__.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index bd5845631b2..27616de08de 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -836,11 +836,20 @@ def _read_bigquery_load_job( ) self._start_generic_job(load_job) + table_id = f"{table.project}.{table.dataset_id}.{table.table_id}" + + # Update the table expiration so we aren't limited to the default 24 + # hours of the anonymous dataset. + table_expiration = bigquery.Table(table_id) + table_expiration.expires = ( + datetime.datetime.now(datetime.timezone.utc) + constants.DEFAULT_EXPIRATION + ) + self.bqclient.update_table(table_expiration, ["expires"]) # 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}", + table_id, index_col=index_col, col_order=col_order, ) From 741c75e5797e26a1487ff3da76a07953d9537f3f Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 10 Nov 2023 12:20:12 -0800 Subject: [PATCH 153/197] fix: use random table for `read_pandas` (#192) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: use random table for read_pandas * fix: default to 7 days expiration for `read_csv`, `read_json`, `read_parquet` (#193) 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) Follow-up to https://togithub.com/googleapis/python-bigquery-dataframes/pull/175/files#r1389686556 🦕 * fix: use random table for read_pandas * fix: use the real table expression --------- Co-authored-by: Tim Swast --- bigframes/core/__init__.py | 2 +- bigframes/session/__init__.py | 19 ++++++------------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 63f36d4ddd2..b640692bc85 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -165,7 +165,7 @@ def cached(self, cluster_cols: typing.Sequence[str]) -> ArrayValue: ibis_expr = compiled_value._to_ibis_expr( ordering_mode="unordered", expose_hidden_cols=True ) - tmp_table = self.session._ibis_to_session_table( + tmp_table = self.session._ibis_to_temp_table( ibis_expr, cluster_cols=cluster_cols, api_name="cached" ) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 27616de08de..9b881de9a02 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -36,7 +36,6 @@ Tuple, Union, ) -import uuid import warnings import google.api_core.client_info @@ -986,7 +985,7 @@ def _read_pandas( job_config.clustering_fields = cluster_cols job_config.labels = {"bigframes-api": api_name} - load_table_destination = self._create_session_table() + load_table_destination = bigframes_io.random_table(self._anonymous_dataset) load_job = self.bqclient.load_table_from_dataframe( pandas_dataframe_copy, load_table_destination, @@ -999,8 +998,9 @@ def _read_pandas( 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}`" + table_expression = self.ibis_client.table( + load_table_destination.table_id, + database=f"{load_table_destination.project}.{load_table_destination.dataset_id}", ) # 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. @@ -1278,13 +1278,6 @@ def _check_file_size(self, filepath: str): "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_empty_temp_table( self, schema: Iterable[bigquery.SchemaField], @@ -1319,7 +1312,7 @@ def _create_sequential_ordering( ibis.row_number().cast(ibis_dtypes.int64).name(default_ordering_name) ) table = table.mutate(**{default_ordering_name: default_ordering_col}) - table_ref = self._ibis_to_session_table( + table_ref = self._ibis_to_temp_table( table, cluster_cols=list(index_cols) + [default_ordering_name], api_name=api_name, @@ -1335,7 +1328,7 @@ def _create_sequential_ordering( ) return table, ordering - def _ibis_to_session_table( + def _ibis_to_temp_table( self, table: ibis_types.Table, cluster_cols: Iterable[str], From 8ab81dee4d0eee499094f2dd576550f0c59d7551 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 10 Nov 2023 16:12:14 -0800 Subject: [PATCH 154/197] feat: read_gbq creates order deterministically without table copy (#191) 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/dataframe.py | 3 +- bigframes/ml/core.py | 20 +- bigframes/series.py | 3 +- bigframes/session/__init__.py | 418 +++++++----------- tests/system/small/ml/test_cluster.py | 110 ++--- tests/system/small/ml/test_core.py | 208 ++++----- tests/system/small/ml/test_decomposition.py | 104 ++--- tests/system/small/ml/test_forecasting.py | 1 + tests/system/small/test_dataframe_io.py | 18 + tests/system/small/test_pandas_options.py | 10 +- tests/system/small/test_progress_bar.py | 8 - tests/unit/ml/test_golden_sql.py | 2 + .../ibis/backends/bigquery/registry.py | 12 +- .../ibis/expr/operations/__init__.py | 5 +- .../ibis/expr/operations/json.py | 9 + .../bigframes_vendored/pandas/io/gbq.py | 10 - .../bigframes_vendored/pandas/io/parquet.py | 6 - .../bigframes_vendored/pandas/io/pickle.py | 10 - 18 files changed, 438 insertions(+), 519 deletions(-) create mode 100644 third_party/bigframes_vendored/ibis/expr/operations/json.py diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 0a035754915..bdbc00e6207 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2719,7 +2719,8 @@ def _get_block(self) -> blocks.Block: return self._block def _cached(self) -> DataFrame: - return DataFrame(self._block.cached()) + self._set_block(self._block.cached()) + return self _DataFrameOrSeries = typing.TypeVar("_DataFrameOrSeries") diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 4c5a48cf625..39d01fca8d7 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -126,7 +126,7 @@ def generate_text_embedding( def forecast(self) -> bpd.DataFrame: sql = self._model_manipulation_sql_generator.ml_forecast() - return self._session.read_gbq(sql) + return self._session.read_gbq(sql, index_col="forecast_timestamp").reset_index() def evaluate(self, input_data: Optional[bpd.DataFrame] = None): # TODO: validate input data schema @@ -139,14 +139,18 @@ def centroids(self) -> bpd.DataFrame: sql = self._model_manipulation_sql_generator.ml_centroids() - return self._session.read_gbq(sql) + return self._session.read_gbq( + sql, index_col=["centroid_id", "feature"] + ).reset_index() def principal_components(self) -> bpd.DataFrame: assert self._model.model_type == "PCA" sql = self._model_manipulation_sql_generator.ml_principal_components() - return self._session.read_gbq(sql) + return self._session.read_gbq( + sql, index_col=["principal_component_id", "feature"] + ).reset_index() def principal_component_info(self) -> bpd.DataFrame: assert self._model.model_type == "PCA" @@ -228,10 +232,12 @@ def create_model( Returns: a BqmlModel, wrapping a trained model in BigQuery """ options = dict(options) + # Cache dataframes to make sure base table is not a snapshot + # cached dataframe creates a full copy, never uses snapshot if y_train is None: - input_data = X_train + input_data = X_train._cached() else: - input_data = X_train.join(y_train, how="outer") + input_data = X_train._cached().join(y_train._cached(), how="outer") options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) session = X_train._session @@ -259,7 +265,9 @@ def create_time_series_model( ), "Time stamp data input must only contain 1 column." options = dict(options) - input_data = X_train.join(y_train, how="outer") + # Cache dataframes to make sure base table is not a snapshot + # cached dataframe creates a full copy, never uses snapshot + input_data = X_train._cached().join(y_train._cached(), how="outer") options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]}) options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) diff --git a/bigframes/series.py b/bigframes/series.py index 4fab1fe9439..28290d591e0 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1503,7 +1503,8 @@ def _slice( ) def _cached(self) -> Series: - return Series(self._block.cached()) + self._set_block(self._block.cached()) + return self def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 9b881de9a02..82c5a1c8d0d 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -17,10 +17,10 @@ from __future__ import annotations import datetime +import itertools import logging import os import re -import textwrap import typing from typing import ( Any, @@ -81,6 +81,7 @@ # 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.ibis.expr.operations as vendored_ibis_ops 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 @@ -378,12 +379,6 @@ def read_gbq_query( ... pitchSpeed, ... FROM `bigquery-public-data.baseball.games_wide` ... ''') - >>> df.head(2) - pitcherFirstName pitcherLastName pitchSpeed - 0 0 - 1 0 - - [2 rows x 3 columns] Preserve ordering in a query input. @@ -480,16 +475,6 @@ def read_gbq_table( Read a whole table, with arbitrary ordering or ordering corresponding to the primary key(s). >>> df = bpd.read_gbq_table("bigquery-public-data.ml_datasets.penguins") - >>> df.head(2) - species island culmen_length_mm \\ - 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 - 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 - - culmen_depth_mm flipper_length_mm body_mass_g sex - 0 18.4 184.0 3475.0 FEMALE - 1 19.1 184.0 4650.0 MALE - - [2 rows x 7 columns] See also: :meth:`Session.read_gbq`. """ @@ -503,7 +488,7 @@ def read_gbq_table( api_name="read_gbq_table", ) - def _read_gbq_table_to_ibis_with_total_ordering( + def _get_snapshot_sql_and_primary_key( self, table_ref: bigquery.table.TableReference, *, @@ -523,7 +508,6 @@ def _read_gbq_table_to_ibis_with_total_ordering( ), None, ) - table_expression = self.ibis_client.table( table_ref.table_id, database=f"{table_ref.project}.{table_ref.dataset_id}", @@ -534,6 +518,11 @@ def _read_gbq_table_to_ibis_with_total_ordering( # the same assumption and use these columns as the total ordering keys. table = self.bqclient.get_table(table_ref) + if table.location.casefold() != self._location.casefold(): + raise ValueError( + f"Current session is in {self._location} but dataset '{table.project}.{table.dataset_id}' is located in {table.location}" + ) + # TODO(b/305264153): Use public properties to fetch primary keys once # added to google-cloud-bigquery. primary_keys = ( @@ -542,22 +531,18 @@ def _read_gbq_table_to_ibis_with_total_ordering( .get("columns") ) - if not primary_keys: - return table_expression, None - else: - # Read from a snapshot since we won't have to copy the table data to create a total ordering. - job_config = bigquery.QueryJobConfig() - job_config.labels["bigframes-api"] = api_name - current_timestamp = list( - self.bqclient.query( - "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", - job_config=job_config, - ).result() - )[0][0] - table_expression = self.ibis_client.sql( - bigframes_io.create_snapshot_sql(table_ref, current_timestamp) - ) - return table_expression, primary_keys + job_config = bigquery.QueryJobConfig() + job_config.labels["bigframes-api"] = api_name + current_timestamp = list( + self.bqclient.query( + "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", + job_config=job_config, + ).result() + )[0][0] + table_expression = self.ibis_client.sql( + bigframes_io.create_snapshot_sql(table_ref, current_timestamp) + ) + return table_expression, primary_keys def _read_gbq_table( self, @@ -580,10 +565,7 @@ def _read_gbq_table( ( table_expression, total_ordering_cols, - ) = self._read_gbq_table_to_ibis_with_total_ordering( - table_ref, - api_name=api_name, - ) + ) = self._get_snapshot_sql_and_primary_key(table_ref, api_name=api_name) for key in col_order: if key not in table_expression.columns: @@ -596,24 +578,22 @@ 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( f"Column `{key}` of `index_col` not found in this table." ) + if col_order: + table_expression = table_expression.select([*index_cols, *col_order]) + # If the index is unique and sortable, then we don't need to generate # an ordering column. ordering = None - is_total_ordering = False - if total_ordering_cols is not None: - # Note: currently, this a table has a total ordering only when the + # Note: currently, a table has a total ordering only when the # primary key(s) are set on a table. The query engine assumes such # columns are unique, even if not enforced. - is_total_ordering = True ordering = orderings.ExpressionOrdering( ordering_value_columns=tuple( core.OrderingColumnReference(column_id) @@ -621,41 +601,17 @@ def _read_gbq_table( ), total_ordering_columns=frozenset(total_ordering_cols), ) - - if len(index_cols) != 0: - index_labels = typing.cast(List[Optional[str]], index_cols) - else: - # Use the total_ordering_cols to project offsets to use as the default index. - table_expression = table_expression.order_by(index_cols) - default_index_id = guid.generate_guid("bigframes_index_") - default_index_col = ( - ibis.row_number().cast(ibis_dtypes.int64).name(default_index_id) - ) - table_expression = table_expression.mutate( - **{default_index_id: default_index_col} - ) - index_cols = [default_index_id] - index_labels = [None] - elif 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)} + column_values = [table_expression[col] for col in table_expression.columns] + array_value = core.ArrayValue.from_ibis( + self, + table_expression, + columns=column_values, + hidden_ordering_columns=[], + ordering=ordering, ) - 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)) - - total_count = row["total_count"] - distinct_count = row["distinct_count"] - is_total_ordering = total_count == distinct_count - + elif len(index_cols) != 0: + # We have index columns, lets see if those are actually total_order_columns ordering = orderings.ExpressionOrdering( ordering_value_columns=tuple( [ @@ -665,142 +621,61 @@ def _read_gbq_table( ), total_ordering_columns=frozenset(index_cols), ) - - # We have a total ordering, so query via "time travel" so that - # the underlying data doesn't mutate. + is_total_ordering = self._check_index_uniqueness( + table_expression, index_cols + ) 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) + column_values = [ + table_expression[col] for col in table_expression.columns + ] + array_value = core.ArrayValue.from_ibis( + self, + table_expression, + columns=column_values, + hidden_ordering_columns=[], + ordering=ordering, ) else: - # 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. - """, - ) - ) - - # 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) + array_value = self._create_total_ordering(table_expression) 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=table_expression, - api_name=api_name, - ) - 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 - # 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, - api_name=api_name, + array_value = self._create_total_ordering(table_expression) + + value_columns = [col for col in array_value.column_ids if col not in index_cols] + block = blocks.Block( + array_value, + index_columns=index_cols, + column_labels=value_columns, + index_labels=index_cols, ) + if max_results: + block = block.slice(stop=max_results) + df = dataframe.DataFrame(block) - def _read_gbq_with_ordering( - self, - table_expression: ibis_types.Table, - *, - col_order: Iterable[str] = (), - col_labels: Iterable[Optional[str]] = (), - index_cols: Iterable[str] = (), - index_labels: Iterable[Optional[str]] = (), - hidden_cols: Iterable[str] = (), - ordering: orderings.ExpressionOrdering, - is_total_ordering: bool = False, - api_name: str, - ) -> dataframe.DataFrame: - """Internal helper method that loads DataFrame from Google BigQuery given an ordering column. + # If user provided index columns, should sort over it + if len(index_cols) > 0: + df.sort_index() + return df - Args: - table_expression: - an ibis table expression to be executed in BigQuery. - col_order: - 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 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: - Column name to be used for ordering. If not supplied, a default ordering is generated. - api_name: - The name of the API method. + def _check_index_uniqueness( + self, table: ibis_types.Table, index_cols: List[str] + ) -> bool: + distinct_table = table.select(*index_cols).distinct() + is_unique_sql = f"""WITH full_table AS ( + {self.ibis_client.compile(table)} + ), + distinct_table AS ( + {self.ibis_client.compile(distinct_table)} + ) - Returns: - A DataFrame representing results of the query or table. + SELECT (SELECT COUNT(*) FROM full_table) AS `total_count`, + (SELECT COUNT(*) FROM distinct_table) AS `distinct_count` """ - 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. " - f"Got {len(index_labels)}, expected {len(index_cols)}." - ) - - # 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 + results, _ = self._start_query(is_unique_sql) + row = next(iter(results)) - # 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=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: - col_labels = column_keys - return self._read_ibis( - table_expression, - index_col_values, - index_labels, - column_keys, - col_labels, - ordering=ordering, - ) + total_count = row["total_count"] + distinct_count = row["distinct_count"] + return total_count == distinct_count def _read_bigquery_load_job( self, @@ -853,40 +728,6 @@ def _read_bigquery_load_job( col_order=col_order, ) - def _read_ibis( - self, - table_expression: ibis_types.Table, - index_cols: Iterable[ibis_types.Value], - index_labels: Iterable[blocks.Label], - column_keys: Iterable[str], - column_labels: Iterable[blocks.Label], - ordering: orderings.ExpressionOrdering, - ) -> dataframe.DataFrame: - """Turns a table expression (plus index column) into a DataFrame.""" - - 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]) - - 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.from_ibis( - self, table_expression, columns, hidden_ordering_columns, ordering - ), - index_columns=[index_col.get_name() for index_col in index_cols], - column_labels=column_labels, - index_labels=index_labels, - ) - - return dataframe.DataFrame(block) - def read_gbq_model(self, model_name: str): """Loads a BigQuery ML model from BigQuery. @@ -1009,17 +850,26 @@ def _read_pandas( ): new_idx_ids, idx_labels = [], [] - df = self._read_gbq_with_ordering( - table_expression=table_expression, - col_labels=col_labels, - index_cols=new_idx_ids, - index_labels=idx_labels, - hidden_cols=(ordering_col,), + column_values = [ + table_expression[col] + for col in table_expression.columns + if col != ordering_col + ] + array_value = core.ArrayValue.from_ibis( + self, + table_expression, + columns=column_values, + hidden_ordering_columns=[table_expression[ordering_col]], ordering=ordering, - is_total_ordering=True, - api_name=api_name, ) - return df + + block = blocks.Block( + array_value, + index_columns=new_idx_ids, + column_labels=col_labels, + index_labels=idx_labels, + ) + return dataframe.DataFrame(block) def read_csv( self, @@ -1299,34 +1149,50 @@ def _create_empty_temp_table( ) return bigquery.TableReference.from_string(table) - def _create_sequential_ordering( + def _create_total_ordering( self, table: ibis_types.Table, - index_cols: Iterable[str] = (), - api_name: str = "", - ) -> Tuple[ibis_types.Table, orderings.ExpressionOrdering]: + ) -> core.ArrayValue: # 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) + ordering_hash_part = guid.generate_guid("bigframes_ordering_") + ordering_rand_part = guid.generate_guid("bigframes_ordering_") + + str_values = list( + map(lambda col: _convert_to_string(table[col]), table.columns) ) - table = table.mutate(**{default_ordering_name: default_ordering_col}) - table_ref = self._ibis_to_temp_table( - table, - cluster_cols=list(index_cols) + [default_ordering_name], - api_name=api_name, + full_row_str = ( + str_values[0].concat(*str_values[1:]) + if len(str_values) > 1 + else str_values[0] ) - table = self.ibis_client.table( - f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}" + full_row_hash = full_row_str.hash().name(ordering_hash_part) + # Used to disambiguate between identical rows (which will have identical hash) + random_value = ibis.random().name(ordering_rand_part) + + original_column_ids = table.columns + table_with_ordering = table.select( + itertools.chain(original_column_ids, [full_row_hash, random_value]) ) - ordering_reference = core.OrderingColumnReference(default_ordering_name) + + ordering_ref1 = core.OrderingColumnReference(ordering_hash_part) + ordering_ref2 = core.OrderingColumnReference(ordering_rand_part) ordering = orderings.ExpressionOrdering( - ordering_value_columns=tuple([ordering_reference]), - total_ordering_columns=frozenset([default_ordering_name]), - integer_encoding=IntegerEncoding(is_encoded=True, is_sequential=True), + ordering_value_columns=(ordering_ref1, ordering_ref2), + total_ordering_columns=frozenset([ordering_hash_part, ordering_rand_part]), + ) + columns = [table_with_ordering[col] for col in original_column_ids] + hidden_columns = [ + table_with_ordering[ordering_hash_part], + table_with_ordering[ordering_rand_part], + ] + return core.ArrayValue.from_ibis( + self, + table_with_ordering, + columns, + hidden_ordering_columns=hidden_columns, + ordering=ordering, ) - return table, ordering def _ibis_to_temp_table( self, @@ -1561,3 +1427,23 @@ def _can_cluster_bq(field: bigquery.SchemaField): "BOOL", "BOOLEAN", ) + + +def _convert_to_string(column: ibis_types.Column) -> ibis_types.StringColumn: + col_type = column.type() + if ( + col_type.is_numeric() + or col_type.is_boolean() + or col_type.is_binary() + or col_type.is_temporal() + ): + result = column.cast(ibis_dtypes.String(nullable=True)) + elif col_type.is_geospatial(): + result = typing.cast(ibis_types.GeoSpatialColumn, column).as_text() + elif col_type.is_string(): + result = column + else: + # TO_JSON_STRING works with all data types, but isn't the most efficient + # Needed for JSON, STRUCT and ARRAY datatypes + result = vendored_ibis_ops.ToJsonString(column).to_expr() # type: ignore + return typing.cast(ibis_types.StringColumn, result) diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py index caeffa77689..266a38e3ee1 100644 --- a/tests/system/small/ml/test_cluster.py +++ b/tests/system/small/ml/test_cluster.py @@ -89,59 +89,67 @@ 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}, + result = ( + penguins_kmeans_model.cluster_centers_.to_pandas() + .sort_values(["centroid_id", "feature"]) + .reset_index(drop=True) + ) + 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, ], - [], - [], - [], - [ - {"category": "MALE", "value": 0.6847826086956522}, - {"category": "FEMALE", "value": 0.2826086956521739}, - {"category": "_null_filler", "value": 0.03260869565217391}, + "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}, + ], ], - ], - }, + }, + ) + .sort_values(["centroid_id", "feature"]) + .reset_index(drop=True) ) pd.testing.assert_frame_equal( result, diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index ec1f351d876..be34a4871c6 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -78,58 +78,62 @@ def test_model_eval_with_data(penguins_bqml_linear_model, penguins_df_default_in 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}, + 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, ], - [], - [], - [], - [ - {"category": "MALE", "value": 0.34375}, - {"category": "FEMALE", "value": 0.625}, - {"category": "_null_filler", "value": 0.03125}, + "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}, + ], ], - [], - [], - [], - [ - {"category": "MALE", "value": 0.6847826086956522}, - {"category": "FEMALE", "value": 0.2826086956521739}, - {"category": "_null_filler", "value": 0.03260869565217391}, - ], - ], - }, + }, + ) + .sort_values(["centroid_id", "feature"]) + .reset_index(drop=True) ) pd.testing.assert_frame_equal( result, @@ -148,59 +152,63 @@ def test_pca_model_principal_components(penguins_bqml_pca_model: core.BqmlModel) # 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, - }, + expected = ( + pd.DataFrame( + { + "principal_component_id": [0] * 7, + "feature": [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + "sex", ], - [ - {"category": "Biscoe", "value": 0.19761120114410635}, - {"category": "Dream", "value": -0.11264736305259061}, - {"category": "Torgersen", "value": -0.07065913511418596}, + "numerical_value": [ + pd.NA, + pd.NA, + 0.401489, + -0.377482, + 0.524052, + 0.501174, + pd.NA, ], - [], - [], - [], - [], - [ - {"category": ".", "value": 0.0015916894448071784}, - {"category": "MALE", "value": 0.06869704739750442}, - {"category": "FEMALE", "value": -0.052521171596813174}, - {"category": "_null_filler", "value": -0.0034628622681684906}, + "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}, + ], ], - ], - }, + }, + ) + .sort_values(["principal_component_id", "feature"]) + .reset_index(drop=True) ) pd.testing.assert_frame_equal( result, diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py index cc4d2e5801c..42fea66cf86 100644 --- a/tests/system/small/ml/test_decomposition.py +++ b/tests/system/small/ml/test_decomposition.py @@ -57,59 +57,63 @@ def test_pca_components_(penguins_pca_model: decomposition.PCA): # 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, - }, + expected = ( + pd.DataFrame( + { + "principal_component_id": [0] * 7, + "feature": [ + "species", + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "body_mass_g", + "sex", ], - [ - {"category": "Biscoe", "value": 0.19761120114410635}, - {"category": "Dream", "value": -0.11264736305259061}, - {"category": "Torgersen", "value": -0.07065913511418596}, + "numerical_value": [ + pd.NA, + pd.NA, + 0.401489, + -0.377482, + 0.524052, + 0.501174, + pd.NA, ], - [], - [], - [], - [], - [ - {"category": ".", "value": 0.0015916894448071784}, - {"category": "MALE", "value": 0.06869704739750442}, - {"category": "FEMALE", "value": -0.052521171596813174}, - {"category": "_null_filler", "value": -0.0034628622681684906}, + "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}, + ], ], - ], - }, + }, + ) + .sort_values(["principal_component_id", "feature"]) + .reset_index(drop=True) ) pd.testing.assert_frame_equal( result, diff --git a/tests/system/small/ml/test_forecasting.py b/tests/system/small/ml/test_forecasting.py index cb27dd388c3..55079c94cfd 100644 --- a/tests/system/small/ml/test_forecasting.py +++ b/tests/system/small/ml/test_forecasting.py @@ -36,6 +36,7 @@ def test_model_predict(time_series_arima_plus_model): expected["forecast_timestamp"] = expected["forecast_timestamp"].astype( pd.ArrowDtype(pa.timestamp("us", tz="UTC")) ) + pd.testing.assert_frame_equal( predictions, expected, diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py index d700d93be99..fb9fb7bb898 100644 --- a/tests/system/small/test_dataframe_io.py +++ b/tests/system/small/test_dataframe_io.py @@ -80,6 +80,24 @@ def test_to_pandas_array_struct_correct_result(session): ) +def test_load_json(session): + df = session.read_gbq( + """SELECT + JSON_OBJECT('foo', 10, 'bar', TRUE) AS json_column + """ + ) + + result = df.to_pandas() + expected = pd.DataFrame( + { + "json_column": ['{"bar":true,"foo":10}'], + } + ) + expected.index = expected.index.astype("Int64") + pd.testing.assert_series_equal(result.dtypes, expected.dtypes) + pd.testing.assert_series_equal(result["json_column"], expected["json_column"]) + + def test_to_pandas_batches_w_correct_dtypes(scalars_df_default_index): """Verify to_pandas_batches() APIs returns the expected dtypes.""" expected = scalars_df_default_index.dtypes diff --git a/tests/system/small/test_pandas_options.py b/tests/system/small/test_pandas_options.py index ca67710d4ea..c410d70fe73 100644 --- a/tests/system/small/test_pandas_options.py +++ b/tests/system/small/test_pandas_options.py @@ -74,7 +74,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, + (google.api_core.exceptions.NotFound, ValueError), match=dataset_id_permanent, ): read_method(query) @@ -99,7 +99,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, + (google.api_core.exceptions.NotFound, ValueError), match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -145,7 +145,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, + (google.api_core.exceptions.NotFound, ValueError), match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -193,7 +193,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, + (google.api_core.exceptions.NotFound, ValueError), match=dataset_id_permanent_tokyo, ): read_method(query_tokyo) @@ -243,7 +243,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, + (google.api_core.exceptions.NotFound, ValueError), match=dataset_id_permanent, ): read_method(query) diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index 30ea63b483a..c6eee820530 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -52,14 +52,6 @@ def test_progress_bar_scalar(penguins_df_default_index: bf.dataframe.DataFrame, assert_loading_msg_exist(capsys.readouterr().out) -def test_progress_bar_read_gbq(session: bf.Session, penguins_table_id: str, capsys): - bf.options.display.progress_bar = "terminal" - capsys.readouterr() # clear output - session.read_gbq(penguins_table_id) - - assert_loading_msg_exist(capsys.readouterr().out) - - def test_progress_bar_extract_jobs( penguins_df_default_index: bf.dataframe.DataFrame, gcs_folder, capsys ): diff --git a/tests/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py index 3ca7e144a53..700eb500ffd 100644 --- a/tests/unit/ml/test_golden_sql.py +++ b/tests/unit/ml/test_golden_sql.py @@ -38,6 +38,7 @@ def mock_session(): def mock_y(): mock_y = mock.create_autospec(spec=bpd.DataFrame) mock_y.columns = pd.Index(["input_column_label"]) + mock_y._cached.return_value = mock_y return mock_y @@ -57,6 +58,7 @@ def mock_X(mock_y, mock_session): ["index_column_id"], ["index_column_label"], ) + mock_X._cached.return_value = mock_X return mock_X diff --git a/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py index a4e61ca0f91..e1b28690d7c 100644 --- a/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py +++ b/third_party/bigframes_vendored/ibis/backends/bigquery/registry.py @@ -22,10 +22,16 @@ def _last_non_null_value(translator, op: vendored_ibis_ops.LastNonNullValue): return f"LAST_VALUE({arg} IGNORE NULLS)" +def _to_json_string(translator, op: vendored_ibis_ops.ToJsonString): + arg = translator.translate(op.arg) + return f"TO_JSON_STRING({arg})" + + patched_ops = { - vendored_ibis_ops.ApproximateMultiQuantile: _approx_quantiles, - vendored_ibis_ops.FirstNonNullValue: _first_non_null_value, - vendored_ibis_ops.LastNonNullValue: _last_non_null_value, + vendored_ibis_ops.ApproximateMultiQuantile: _approx_quantiles, # type:ignore + vendored_ibis_ops.FirstNonNullValue: _first_non_null_value, # type:ignore + vendored_ibis_ops.LastNonNullValue: _last_non_null_value, # type:ignore + vendored_ibis_ops.ToJsonString: _to_json_string, # type:ignore } 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 1612d9c12e5..82197013923 100644 --- a/third_party/bigframes_vendored/ibis/expr/operations/__init__.py +++ b/third_party/bigframes_vendored/ibis/expr/operations/__init__.py @@ -1,5 +1,6 @@ # 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 +from third_party.bigframes_vendored.ibis.expr.operations.analytic import * # noqa: F401 F403 +from third_party.bigframes_vendored.ibis.expr.operations.json import * # noqa: F401 F403 +from third_party.bigframes_vendored.ibis.expr.operations.reductions import * # noqa: F401 F403 diff --git a/third_party/bigframes_vendored/ibis/expr/operations/json.py b/third_party/bigframes_vendored/ibis/expr/operations/json.py new file mode 100644 index 00000000000..dbb3fa3066c --- /dev/null +++ b/third_party/bigframes_vendored/ibis/expr/operations/json.py @@ -0,0 +1,9 @@ +# Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/json.py +from __future__ import annotations + +import ibis.expr.datatypes as dt +from ibis.expr.operations.core import Unary + + +class ToJsonString(Unary): + output_dtype = dt.string diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 575c5016187..2161310b07c 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -45,16 +45,6 @@ def read_gbq( If the input is a table ID: >>> df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") - >>> df.head(2) - species island culmen_length_mm \\ - 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 - 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 - - culmen_depth_mm flipper_length_mm body_mass_g sex - 0 18.4 184.0 3475.0 FEMALE - 1 19.1 184.0 4650.0 MALE - - [2 rows x 7 columns] Preserve ordering in a query input. diff --git a/third_party/bigframes_vendored/pandas/io/parquet.py b/third_party/bigframes_vendored/pandas/io/parquet.py index f97bd386a40..0f664e70fcb 100644 --- a/third_party/bigframes_vendored/pandas/io/parquet.py +++ b/third_party/bigframes_vendored/pandas/io/parquet.py @@ -24,12 +24,6 @@ def read_parquet( >>> gcs_path = "gs://cloud-samples-data/bigquery/us-states/us-states.parquet" >>> df = bpd.read_parquet(path=gcs_path) - >>> df.head(2) - name post_abbr - 0 Alabama AL - 1 Alaska AK - - [2 rows x 2 columns] Args: path (str): diff --git a/third_party/bigframes_vendored/pandas/io/pickle.py b/third_party/bigframes_vendored/pandas/io/pickle.py index 053ba4871c1..096d9b13d68 100644 --- a/third_party/bigframes_vendored/pandas/io/pickle.py +++ b/third_party/bigframes_vendored/pandas/io/pickle.py @@ -32,16 +32,6 @@ def read_pickle( >>> gcs_path = "gs://bigframes-dev-testing/test_pickle.pkl" >>> df = bpd.read_pickle(filepath_or_buffer=gcs_path) - >>> df.head(2) - species island culmen_length_mm \\ - 0 Adelie Penguin (Pygoscelis adeliae) Dream 36.6 - 1 Adelie Penguin (Pygoscelis adeliae) Dream 39.8 - - culmen_depth_mm flipper_length_mm body_mass_g sex - 0 18.4 184.0 3475.0 FEMALE - 1 19.1 184.0 4650.0 MALE - - [2 rows x 7 columns] Args: filepath_or_buffer (str, path object, or file-like object): From 7d2be003a16075ba2cd5327e049aaa7ca1a34f14 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 13 Nov 2023 09:12:08 -0600 Subject: [PATCH 155/197] chore: update formatter in CONTRIBUTING.rst (#178) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: update formatter in CONTRIBUTING.rst * 🦉 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 * use owlbot to update the files * remove for loop to format code samples, these are managed by the root noxfile. * 🦉 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 --- CONTRIBUTING.md | 33 --------------------------------- CONTRIBUTING.rst | 4 ++-- owlbot.py | 9 +++++++-- 3 files changed, 9 insertions(+), 37 deletions(-) delete mode 100644 CONTRIBUTING.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md deleted file mode 100644 index b16bd944285..00000000000 --- a/CONTRIBUTING.md +++ /dev/null @@ -1,33 +0,0 @@ -# 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/CONTRIBUTING.rst b/CONTRIBUTING.rst index 3933152cf78..f9103bfa72c 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -96,9 +96,9 @@ On Debian/Ubuntu:: 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:: + the nox session ``format``. This will eliminate many lint errors. Run via:: - $ nox -s blacken + $ nox -s format - PEP8 compliance is required, with exceptions defined in the linter configuration. If you have ``nox`` installed, you can test that you have not introduced diff --git a/owlbot.py b/owlbot.py index be30eea5c2c..082970018dc 100644 --- a/owlbot.py +++ b/owlbot.py @@ -99,6 +99,13 @@ "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine", ) +# Update the contributing guide to reflect some differences in this repo. +s.replace( + ["CONTRIBUTING.rst"], + re.escape("blacken"), + "format", +) + # ---------------------------------------------------------------------------- # Samples templates # ---------------------------------------------------------------------------- @@ -110,5 +117,3 @@ # ---------------------------------------------------------------------------- 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) From 87f84c9e58e7d0ea521ac386c9f02791cdddd19f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 13 Nov 2023 11:26:14 -0600 Subject: [PATCH 156/197] docs: use `head()` to get top `n` results, not to preview results (#190) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit head() requires ordering. Just peeking at the whole DataFrame or Series is actually more efficient since it doesn't require ordering and still only downloads a fraction of the results. 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 # 🦕 --- samples/snippets/pandas_methods_test.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/samples/snippets/pandas_methods_test.py b/samples/snippets/pandas_methods_test.py index 1f472d63466..bd8e29c0030 100644 --- a/samples/snippets/pandas_methods_test.py +++ b/samples/snippets/pandas_methods_test.py @@ -22,13 +22,20 @@ def test_bigquery_dataframes_pandas_methods(): bq_df = bpd.read_gbq(query_or_table) # Inspect one of the columns (or series) of the DataFrame: - bq_df["body_mass_g"].head(10) + bq_df["body_mass_g"] # Compute the mean of this series: average_body_mass = bq_df["body_mass_g"].mean() print(f"average_body_mass: {average_body_mass}") - # Calculate the mean body_mass_g by species using the groupby operation: - bq_df["body_mass_g"].groupby(by=bq_df["species"]).mean().head() + # Find the heaviest species using the groupby operation to calculate the + # mean body_mass_g: + ( + bq_df["body_mass_g"] + .groupby(by=bq_df["species"]) + .mean() + .sort_values(ascending=False) + .head(10) + ) # [END bigquery_dataframes_pandas_methods] assert average_body_mass is not None From 7506eabf2e58159507809e36abfe90c417dfe92f Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Mon, 13 Nov 2023 11:58:35 -0800 Subject: [PATCH 157/197] docs: add code samples for `read_gbq_function` using community UDFs (#188) * docs: add code samples for read_gbq_function using community UDFs --- bigframes/session/__init__.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 82c5a1c8d0d..67c429eda59 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1322,15 +1322,37 @@ def read_gbq_function( The return type of the function must be explicitly specified in the function's original definition even if not otherwise required. + BigQuery Utils provides many public functions under the ``bqutil`` project on Google Cloud Platform project + (See: https://github.com/GoogleCloudPlatform/bigquery-utils/tree/master/udfs#using-the-udfs). + You can checkout Community UDFs to use community-contributed functions. + (See: https://github.com/GoogleCloudPlatform/bigquery-utils/tree/master/udfs/community#community-udfs). + **Examples:** + Use the ``cw_lower_case_ascii_only`` function from Community UDFs. + (https://github.com/GoogleCloudPlatform/bigquery-utils/blob/master/udfs/community/cw_lower_case_ascii_only.sqlx) + >>> import bigframes.pandas as bpd >>> bpd.options.display.progress_bar = None - >>> function_name = "bqutil.fn.cw_lower_case_ascii_only" - >>> func = bpd.read_gbq_function(function_name=function_name) - >>> func.bigframes_remote_function - 'bqutil.fn.cw_lower_case_ascii_only' + >>> df = bpd.DataFrame({'id': [1, 2, 3], 'name': ['AURÉLIE', 'CÉLESTINE', 'DAPHNÉ']}) + >>> df + id name + 0 1 AURÉLIE + 1 2 CÉLESTINE + 2 3 DAPHNÉ + + [3 rows x 2 columns] + + >>> func = bpd.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only") + >>> df1 = df.assign(new_name=df['name'].apply(func)) + >>> df1 + id name new_name + 0 1 AURÉLIE aurÉlie + 1 2 CÉLESTINE cÉlestine + 2 3 DAPHNÉ daphnÉ + + [3 rows x 3 columns] Args: function_name (str): From 4a02cac88c7d7b46bed1fa813a862fc2ef9ef084 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Tue, 14 Nov 2023 11:47:55 -0600 Subject: [PATCH 158/197] feat: temporary resources no longer use BigQuery Sessions (#194) * feat: temporary resources no longer use BigQuery Sessions This allows multiple queries to run in parallel. * use anonymous dataset for models * remove reference to bq session * avoid 'model already exists' error * remove global _model_id --- bigframes/ml/core.py | 39 ++++++++++++------ bigframes/ml/sql.py | 24 +++++++---- bigframes/session/__init__.py | 53 ++++--------------------- tests/system/small/test_progress_bar.py | 6 --- tests/system/small/test_session.py | 24 ----------- tests/unit/ml/test_golden_sql.py | 47 ++++++++++++++-------- tests/unit/ml/test_sql.py | 23 ++++++++--- tests/unit/resources.py | 1 - tests/unit/test_pandas.py | 38 ------------------ 9 files changed, 97 insertions(+), 158 deletions(-) diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index 39d01fca8d7..d8135f7085f 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -16,12 +16,14 @@ from __future__ import annotations +import datetime from typing import Callable, cast, Iterable, Mapping, Optional, Union import uuid 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 @@ -192,24 +194,27 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel: class BqmlModelFactory: def __init__(self): - model_id = self._create_temp_model_id() - self._model_creation_sql_generator = ml_sql.ModelCreationSqlGenerator(model_id) + self._model_creation_sql_generator = ml_sql.ModelCreationSqlGenerator() - def _create_temp_model_id(self) -> str: - return uuid.uuid4().hex - - def _reset_model_id(self): - self._model_creation_sql_generator._model_id = self._create_temp_model_id() + def _create_model_ref( + self, dataset: bigquery.DatasetReference + ) -> bigquery.ModelReference: + return bigquery.ModelReference.from_string( + f"{dataset.project}.{dataset.dataset_id}.{uuid.uuid4().hex}" + ) 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) + model_name_full = f"{job.destination.project}.{job.destination.dataset_id}.{job.destination.table_id}" + model = bigquery.Model(model_name_full) + model.expires = ( + datetime.datetime.now(datetime.timezone.utc) + constants.DEFAULT_EXPIRATION + ) + model = session.bqclient.update_model(model, ["expires"]) - self._reset_model_id() return BqmlModel(session, model) def create_model( @@ -219,7 +224,7 @@ def create_model( 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 + """Create a session-temporary BQML model with the CREATE OR REPLACE MODEL statement Args: X_train: features columns for training @@ -241,9 +246,11 @@ def create_model( options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()}) session = X_train._session + model_ref = self._create_model_ref(session._anonymous_dataset) sql = self._model_creation_sql_generator.create_model( source_df=input_data, + model_ref=model_ref, transforms=transforms, options=options, ) @@ -272,9 +279,11 @@ def create_time_series_model( options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]}) session = X_train._session + model_ref = self._create_model_ref(session._anonymous_dataset) sql = self._model_creation_sql_generator.create_model( source_df=input_data, + model_ref=model_ref, transforms=transforms, options=options, ) @@ -287,7 +296,7 @@ def create_remote_model( 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 + """Create a session-temporary BQML remote model with the CREATE OR REPLACE MODEL statement Args: connection_name: @@ -298,8 +307,10 @@ def create_remote_model( Returns: BqmlModel: a BqmlModel wrapping a trained model in BigQuery """ + model_ref = self._create_model_ref(session._anonymous_dataset) sql = self._model_creation_sql_generator.create_remote_model( connection_name=connection_name, + model_ref=model_ref, options=options, ) @@ -310,7 +321,7 @@ def create_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 + """Create a session-temporary BQML imported model with the CREATE OR REPLACE MODEL statement Args: options: a dict of options to configure the model. Generates a BQML OPTIONS @@ -318,7 +329,9 @@ def create_imported_model( Returns: a BqmlModel, wrapping a trained model in BigQuery """ + model_ref = self._create_model_ref(session._anonymous_dataset) sql = self._model_creation_sql_generator.create_imported_model( + model_ref=model_ref, options=options, ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index 601b2710999..ab051231fb0 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -18,6 +18,8 @@ from typing import Iterable, Mapping, Optional, Union +import google.cloud.bigquery + import bigframes.constants as constants import bigframes.pandas as bpd @@ -121,22 +123,26 @@ def ml_label_encoder( 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 + def _model_id_sql( + self, + model_ref: google.cloud.bigquery.ModelReference, + ): + return f"`{model_ref.project}`.`{model_ref.dataset_id}`.`{model_ref.model_id}`" # Model create and alter def create_model( self, source_df: bpd.DataFrame, + model_ref: google.cloud.bigquery.ModelReference, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, transforms: Optional[Iterable[str]] = None, ) -> str: - """Encode the CREATE TEMP MODEL statement for BQML""" + """Encode the CREATE OR REPLACE MODEL statement for BQML""" source_sql = source_df.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}`"] + parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"] if transform_sql: parts.append(transform_sql) if options_sql: @@ -147,12 +153,13 @@ def create_model( def create_remote_model( self, connection_name: str, + model_ref: google.cloud.bigquery.ModelReference, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: - """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + """Encode the CREATE OR REPLACE MODEL statement for BQML remote model.""" options_sql = self.options(**options) - parts = [f"CREATE TEMP MODEL `{self._model_id}`"] + parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"] parts.append(self.connection(connection_name)) if options_sql: parts.append(options_sql) @@ -160,12 +167,13 @@ def create_remote_model( def create_imported_model( self, + model_ref: google.cloud.bigquery.ModelReference, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: - """Encode the CREATE TEMP MODEL statement for BQML remote model.""" + """Encode the CREATE OR REPLACE MODEL statement for BQML remote model.""" options_sql = self.options(**options) - parts = [f"CREATE TEMP MODEL `{self._model_id}`"] + parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"] if options_sql: parts.append(options_sql) return "\n".join(parts) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 67c429eda59..1b12e5deeb8 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -161,7 +161,7 @@ def __init__( application_name=context.application_name, ) - self._create_and_bind_bq_session() + self._create_bq_datasets() self.ibis_client = typing.cast( ibis_bigquery.Backend, ibis.bigquery.connect( @@ -210,19 +210,12 @@ def _project(self): def __hash__(self): # Stable hash needed to use in expression tree - return hash(self._session_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) - # 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 - ) + return hash(str(self._anonymous_dataset)) + + def _create_bq_datasets(self): + """Create and identify dataset(s) for temporary BQ resources.""" + query_job = self.bqclient.query("SELECT 1", location=self._location) query_job.result() # blocks until finished - self._session_id = query_job.session_info.session_id # The anonymous dataset is used by BigQuery to write query results and # session tables. BigQuery DataFrames also writes temp tables directly @@ -235,17 +228,6 @@ def _create_and_bind_bq_session(self): query_destination.dataset_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 remote functions, which don't yet # support proper session temporary storage yet self._session_dataset = bigquery.Dataset( @@ -254,28 +236,7 @@ def _create_and_bind_bq_session(self): self._session_dataset.location = self._location 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) - try: - query_job = self.bqclient.query(abort_session_query) - query_job.result() # blocks until finished - 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 exc.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 + """No-op. Temporary resources are deleted after 7 days.""" def read_gbq( self, diff --git a/tests/system/small/test_progress_bar.py b/tests/system/small/test_progress_bar.py index c6eee820530..bd13ac22401 100644 --- a/tests/system/small/test_progress_bar.py +++ b/tests/system/small/test_progress_bar.py @@ -90,9 +90,6 @@ def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex): def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): bf.options.display.progress_bar = "terminal" - 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_html( penguins_df_default_index.query_job @@ -109,9 +106,6 @@ def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame): 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 diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index bf72e444eba..7cd9f1dd598 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -19,7 +19,6 @@ import typing from typing import List -import google.api_core.exceptions import google.cloud.bigquery as bigquery import numpy as np import pandas as pd @@ -985,26 +984,3 @@ def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder): 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 - - # 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. - - -@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/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py index 700eb500ffd..017c96d46d2 100644 --- a/tests/unit/ml/test_golden_sql.py +++ b/tests/unit/ml/test_golden_sql.py @@ -23,17 +23,41 @@ from bigframes.ml import core, linear_model import bigframes.pandas as bpd +TEMP_MODEL_ID = bigquery.ModelReference.from_string( + "test-project._anon123.temp_model_id" +) + @pytest.fixture def mock_session(): mock_session = mock.create_autospec(spec=bigframes.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()) + mock_session._anonymous_dataset = bigquery.DatasetReference( + TEMP_MODEL_ID.project, TEMP_MODEL_ID.dataset_id + ) + + query_job = mock.create_autospec(bigquery.QueryJob) + type(query_job).destination = mock.PropertyMock( + return_value=bigquery.TableReference( + mock_session._anonymous_dataset, TEMP_MODEL_ID.model_id + ) + ) + mock_session._start_query.return_value = (None, query_job) return mock_session +@pytest.fixture +def bqml_model_factory(mocker: pytest_mock.MockerFixture): + mocker.patch( + "bigframes.ml.core.BqmlModelFactory._create_model_ref", + return_value=TEMP_MODEL_ID, + ) + bqml_model_factory = core.BqmlModelFactory() + + return bqml_model_factory + + @pytest.fixture def mock_y(): mock_y = mock.create_autospec(spec=bpd.DataFrame) @@ -63,17 +87,6 @@ def mock_X(mock_y, mock_session): return mock_X -@pytest.fixture -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_factory - - @pytest.fixture def bqml_model(mock_session): bqml_model = core.BqmlModel( @@ -91,7 +104,7 @@ def test_linear_regression_default_fit( model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - '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' + 'CREATE OR REPLACE MODEL `test-project`.`_anon123`.`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' ) @@ -101,7 +114,7 @@ def test_linear_regression_params_fit(bqml_model_factory, mock_session, mock_X, model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - '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' + 'CREATE OR REPLACE MODEL `test-project`.`_anon123`.`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' ) @@ -134,7 +147,7 @@ def test_logistic_regression_default_fit( model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - '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' + 'CREATE OR REPLACE MODEL `test-project`.`_anon123`.`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' ) @@ -148,7 +161,7 @@ def test_logistic_regression_params_fit( model.fit(mock_X, mock_y) mock_session._start_query.assert_called_once_with( - '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' + 'CREATE OR REPLACE MODEL `test-project`.`_anon123`.`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' ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 34a02edd424..ea16722393d 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -14,6 +14,7 @@ from unittest import mock +import google.cloud.bigquery as bigquery import pytest import bigframes.ml.sql as ml_sql @@ -27,7 +28,7 @@ def base_sql_generator() -> ml_sql.BaseSqlGenerator: @pytest.fixture(scope="session") def model_creation_sql_generator() -> ml_sql.ModelCreationSqlGenerator: - return ml_sql.ModelCreationSqlGenerator(model_id="my_model_id") + return ml_sql.ModelCreationSqlGenerator() @pytest.fixture(scope="session") @@ -126,11 +127,14 @@ def test_create_model_produces_correct_sql( ): sql = model_creation_sql_generator.create_model( source_df=mock_df, + model_ref=bigquery.ModelReference.from_string( + "test-proj._anonXYZ.create_model_correct_sql" + ), options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql - == """CREATE TEMP MODEL `my_model_id` + == """CREATE OR REPLACE MODEL `test-proj`.`_anonXYZ`.`create_model_correct_sql` OPTIONS( option_key1="option_value1", option_key2=2) @@ -144,6 +148,9 @@ def test_create_model_transform_produces_correct_sql( ): sql = model_creation_sql_generator.create_model( source_df=mock_df, + model_ref=bigquery.ModelReference.from_string( + "test-proj._anonXYZ.create_model_transform" + ), options={"option_key1": "option_value1", "option_key2": 2}, transforms=[ "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a", @@ -152,7 +159,7 @@ def test_create_model_transform_produces_correct_sql( ) assert ( sql - == """CREATE TEMP MODEL `my_model_id` + == """CREATE OR REPLACE MODEL `test-proj`.`_anonXYZ`.`create_model_transform` 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) @@ -168,11 +175,14 @@ def test_create_remote_model_produces_correct_sql( ): sql = model_creation_sql_generator.create_remote_model( connection_name="my_project.us.my_connection", + model_ref=bigquery.ModelReference.from_string( + "test-proj._anonXYZ.create_remote_model" + ), options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql - == """CREATE TEMP MODEL `my_model_id` + == """CREATE OR REPLACE MODEL `test-proj`.`_anonXYZ`.`create_remote_model` REMOTE WITH CONNECTION `my_project.us.my_connection` OPTIONS( option_key1="option_value1", @@ -184,11 +194,14 @@ def test_create_imported_model_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, ): sql = model_creation_sql_generator.create_imported_model( + model_ref=bigquery.ModelReference.from_string( + "test-proj._anonXYZ.create_imported_model" + ), options={"option_key1": "option_value1", "option_key2": 2}, ) assert ( sql - == """CREATE TEMP MODEL `my_model_id` + == """CREATE OR REPLACE MODEL `test-proj`.`_anonXYZ`.`create_imported_model` OPTIONS( option_key1="option_value1", option_key2=2)""" diff --git a/tests/unit/resources.py b/tests/unit/resources.py index 8fc8acd1759..8ba321d122c 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -66,7 +66,6 @@ def create_bigquery_session( credentials=credentials, location="test-region" ) session = bigframes.Session(context=bqoptions, clients_provider=clients_provider) - session._session_id = session_id return session diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index 70c5441c685..4835a24dc7d 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -17,8 +17,6 @@ import sys import unittest.mock as mock -import google.api_core.exceptions -import google.cloud.bigquery import pandas as pd import pytest @@ -26,8 +24,6 @@ import bigframes.pandas as bpd import bigframes.session -from . import resources - leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE) @@ -114,37 +110,3 @@ def test_pandas_attribute(): assert bpd.Int64Dtype is pd.Int64Dtype assert bpd.StringDtype is pd.StringDtype assert bpd.ArrowDtype is pd.ArrowDtype - - -def test_close_session_after_bq_session_ended(monkeypatch: pytest.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 'ABC'") - - # Even though the query to stop the session raises an exception, we should - # 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 3a2761f3c38d0de8b8eda47fffa15b8412aa84b0 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 14 Nov 2023 13:35:49 -0800 Subject: [PATCH 159/197] fix: all sort operation are now stable (#195) --- bigframes/core/__init__.py | 8 ++--- bigframes/core/block_transforms.py | 4 +-- bigframes/core/blocks.py | 4 +-- bigframes/core/compile/compiled.py | 6 ++-- bigframes/core/compile/compiler.py | 2 +- bigframes/core/groupby/__init__.py | 4 --- bigframes/core/nodes.py | 1 - bigframes/core/ordering.py | 47 ++++++++++++++---------------- bigframes/dataframe.py | 4 +-- bigframes/series.py | 7 +---- 10 files changed, 32 insertions(+), 55 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index b640692bc85..b476961bdc6 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -200,12 +200,8 @@ def filter(self, predicate_id: str, keep_null: bool = False) -> ArrayValue: ) ) - def order_by( - self, by: Sequence[OrderingColumnReference], stable: bool = False - ) -> ArrayValue: - return ArrayValue( - nodes.OrderByNode(child=self.node, by=tuple(by), stable=stable) - ) + def order_by(self, by: Sequence[OrderingColumnReference]) -> ArrayValue: + return ArrayValue(nodes.OrderByNode(child=self.node, by=tuple(by))) def reversed(self) -> ArrayValue: return ArrayValue(nodes.ReversedNode(child=self.node)) diff --git a/bigframes/core/block_transforms.py b/bigframes/core/block_transforms.py index e095f21f6b7..ce0fdd219a4 100644 --- a/bigframes/core/block_transforms.py +++ b/bigframes/core/block_transforms.py @@ -509,7 +509,7 @@ def nsmallest( ) for col_id in column_ids ] - block = block.order_by(order_refs, stable=True) + block = block.order_by(order_refs) if keep in ("first", "last"): return block.slice(0, n) else: # keep == "all": @@ -541,7 +541,7 @@ def nlargest( ) for col_id in column_ids ] - block = block.order_by(order_refs, stable=True) + block = block.order_by(order_refs) if keep in ("first", "last"): return block.slice(0, n) else: # keep == "all": diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 6358d28e2ea..f1113d938ec 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -235,10 +235,9 @@ def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]: def order_by( self, by: typing.Sequence[ordering.OrderingColumnReference], - stable: bool = False, ) -> Block: return Block( - self._expr.order_by(by, stable=stable), + self._expr.order_by(by), index_columns=self.index_columns, column_labels=self.column_labels, index_labels=self.index.names, @@ -1596,7 +1595,6 @@ def merge( # sort uses coalesced join keys always joined_expr = joined_expr.order_by( [ordering.OrderingColumnReference(col_id) for col_id in coalesced_ids], - stable=True, ) joined_expr = joined_expr.select_columns(result_columns) diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index 4ba5e6bd080..78050ed4f0c 100644 --- a/bigframes/core/compile/compiled.py +++ b/bigframes/core/compile/compiled.py @@ -753,11 +753,9 @@ def builder(self) -> OrderedIR.Builder: predicates=self._predicates, ) - def order_by( - self, by: Sequence[OrderingColumnReference], stable: bool = False - ) -> OrderedIR: + def order_by(self, by: Sequence[OrderingColumnReference]) -> OrderedIR: expr_builder = self.builder() - expr_builder.ordering = self._ordering.with_ordering_columns(by, stable=stable) + expr_builder.ordering = self._ordering.with_ordering_columns(by) return expr_builder.build() def reversed(self) -> OrderedIR: diff --git a/bigframes/core/compile/compiler.py b/bigframes/core/compile/compiler.py index 662e73a4331..39892635f1e 100644 --- a/bigframes/core/compile/compiler.py +++ b/bigframes/core/compile/compiler.py @@ -129,7 +129,7 @@ def compile_filter(node: nodes.FilterNode, ordered: bool = True): @_compile_node.register def compile_orderby(node: nodes.OrderByNode, ordered: bool = True): if ordered: - return compile_ordered(node.child).order_by(node.by, node.stable) + return compile_ordered(node.child).order_by(node.by) else: return compile_unordered(node.child) diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 2a19a83dd5f..22ef11dd194 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -217,7 +217,6 @@ def rolling(self, window: int, min_periods=None) -> windows.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 @@ -231,7 +230,6 @@ def expanding(self, min_periods: int = 1) -> windows.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 @@ -552,7 +550,6 @@ def rolling(self, window: int, min_periods=None) -> windows.Window: ) block = self._block.order_by( [order.OrderingColumnReference(col) for col in self._by_col_ids], - stable=True, ) return windows.Window( block, @@ -570,7 +567,6 @@ def expanding(self, min_periods: int = 1) -> windows.Window: ) block = self._block.order_by( [order.OrderingColumnReference(col) for col in self._by_col_ids], - stable=True, ) return windows.Window( block, diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py index 8f1e2e5e736..050d3562393 100644 --- a/bigframes/core/nodes.py +++ b/bigframes/core/nodes.py @@ -145,7 +145,6 @@ class FilterNode(UnaryNode): @dataclass(frozen=True) class OrderByNode(UnaryNode): by: Tuple[OrderingColumnReference, ...] - stable: bool = False @dataclass(frozen=True) diff --git a/bigframes/core/ordering.py b/bigframes/core/ordering.py index 2cecd2fe7b2..3ab89e0213c 100644 --- a/bigframes/core/ordering.py +++ b/bigframes/core/ordering.py @@ -28,8 +28,6 @@ # 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 @@ -113,17 +111,12 @@ def with_non_sequential(self): def with_ordering_columns( self, ordering_value_columns: Sequence[OrderingColumnReference] = (), - stable: bool = False, ) -> 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 @@ -131,29 +124,33 @@ def with_ordering_columns( col_ids_new = [ ordering_ref.column_id for ordering_ref in ordering_value_columns ] - if stable: - # 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 - ] - 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) + old_ordering_keep = [ + ordering_ref + for ordering_ref in self.ordering_value_columns + if ordering_ref.column_id not in col_ids_new + ] + + # Truncate to remove any unneded col references after all total order cols included + new_ordering = self._truncate_ordering( + (*ordering_value_columns, *old_ordering_keep) + ) return ExpressionOrdering( new_ordering, total_ordering_columns=self.total_ordering_columns, ) + def _truncate_ordering( + self, order_refs: tuple[OrderingColumnReference, ...] + ) -> tuple[OrderingColumnReference, ...]: + total_order_cols_remaining = set(self.total_ordering_columns) + for i in range(len(order_refs)): + column = order_refs[i].column_id + if column in total_order_cols_remaining: + total_order_cols_remaining.remove(column) + if len(total_order_cols_remaining) == 0: + return order_refs[: i + 1] + raise ValueError("Ordering did not contain all total_order_cols") + def with_reverse(self): """Reverses the ordering.""" return ExpressionOrdering( diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index bdbc00e6207..4716de48d6e 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1262,9 +1262,7 @@ def sort_values( column_id, direction=direction, na_last=na_last ) ) - return DataFrame( - self._block.order_by(ordering, stable=kind in order.STABLE_SORTS) - ) + return DataFrame(self._block.order_by(ordering)) def value_counts( self, diff --git a/bigframes/series.py b/bigframes/series.py index 28290d591e0..2cffdd5fcec 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -35,11 +35,7 @@ import bigframes.core.groupby as groupby import bigframes.core.indexers import bigframes.core.indexes as indexes -from bigframes.core.ordering import ( - OrderingColumnReference, - OrderingDirection, - STABLE_SORTS, -) +from bigframes.core.ordering import OrderingColumnReference, OrderingDirection import bigframes.core.scalar as scalars import bigframes.core.utils as utils import bigframes.core.window @@ -1067,7 +1063,6 @@ def sort_values( na_last=(na_position == "last"), ) ], - stable=kind in STABLE_SORTS, ) return Series(block) From 4ea33b7433532ae3a386a6ffa9eb57360ea39526 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 14 Nov 2023 15:28:07 -0800 Subject: [PATCH 160/197] feat: log most recent API calls as `recent-bigframes-api-xx` labels on BigQuery jobs (#145) --- bigframes/core/groupby/__init__.py | 3 + bigframes/core/log_adapter.py | 61 ++++++++++++++ bigframes/core/nodes.py | 1 + bigframes/core/window/__init__.py | 2 + bigframes/dataframe.py | 2 + bigframes/operations/datetimes.py | 2 + bigframes/operations/strings.py | 2 + bigframes/operations/structs.py | 2 + bigframes/series.py | 2 + bigframes/session/__init__.py | 7 ++ bigframes/session/_io/bigquery.py | 21 ++++- tests/unit/core/test_log_adapter.py | 60 ++++++++++++++ tests/unit/session/test_io_bigquery.py | 110 ++++++++++++++++++++++++- tests/unit/test_compute_options.py | 12 +-- 14 files changed, 276 insertions(+), 11 deletions(-) create mode 100644 bigframes/core/log_adapter.py create mode 100644 tests/unit/core/test_log_adapter.py diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 22ef11dd194..18cb83fa186 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -19,6 +19,7 @@ import pandas as pd import bigframes.constants as constants +from bigframes.core import log_adapter import bigframes.core as core import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks @@ -33,6 +34,7 @@ import third_party.bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby +@log_adapter.class_logger class DataFrameGroupBy(vendored_pandas_groupby.DataFrameGroupBy): __doc__ = vendored_pandas_groupby.GroupBy.__doc__ @@ -406,6 +408,7 @@ def _resolve_label(self, label: blocks.Label) -> str: return col_ids[0] +@log_adapter.class_logger class SeriesGroupBy(vendored_pandas_groupby.SeriesGroupBy): __doc__ = vendored_pandas_groupby.GroupBy.__doc__ diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py new file mode 100644 index 00000000000..b790d195623 --- /dev/null +++ b/bigframes/core/log_adapter.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. + +import functools +import threading +from typing import List + +_lock = threading.Lock() +MAX_LABELS_COUNT = 64 +_api_methods: List = [] + + +def class_logger(decorated_cls): + """Decorator that adds logging functionality to each method of the class.""" + for attr_name, attr_value in decorated_cls.__dict__.items(): + if callable(attr_value): + setattr(decorated_cls, attr_name, method_logger(attr_value)) + return decorated_cls + + +def method_logger(method): + """Decorator that adds logging functionality to a method.""" + + @functools.wraps(method) + def wrapper(*args, **kwargs): + api_method_name = str(method.__name__) + # Track regular and "dunder" methods + if api_method_name.startswith("__") or not api_method_name.startswith("_"): + add_api_method(api_method_name) + return method(*args, **kwargs) + + return wrapper + + +def add_api_method(api_method_name): + global _lock + global _api_methods + with _lock: + # Push the method to the front of the _api_methods list + _api_methods.insert(0, api_method_name) + # Keep the list length within the maximum limit (adjust MAX_LABELS_COUNT as needed) + _api_methods = _api_methods[:MAX_LABELS_COUNT] + + +def get_and_reset_api_methods(): + global _lock + with _lock: + previous_api_methods = list(_api_methods) + _api_methods.clear() + return previous_api_methods diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py index 050d3562393..44a8d808ff5 100644 --- a/bigframes/core/nodes.py +++ b/bigframes/core/nodes.py @@ -11,6 +11,7 @@ # 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, field diff --git a/bigframes/core/window/__init__.py b/bigframes/core/window/__init__.py index d3d081124e9..240715b6df4 100644 --- a/bigframes/core/window/__init__.py +++ b/bigframes/core/window/__init__.py @@ -16,12 +16,14 @@ import typing +from bigframes.core import log_adapter 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 +@log_adapter.class_logger class Window(vendored_pandas_rolling.Window): __doc__ = vendored_pandas_rolling.Window.__doc__ diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 4716de48d6e..57b4ca42cfc 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -41,6 +41,7 @@ import bigframes._config.display_options as display_options import bigframes.constants as constants import bigframes.core +from bigframes.core import log_adapter import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks import bigframes.core.groupby as groupby @@ -81,6 +82,7 @@ # Inherits from pandas DataFrame so that we can use the same docstrings. +@log_adapter.class_logger class DataFrame(vendored_pandas_frame.DataFrame): __doc__ = vendored_pandas_frame.DataFrame.__doc__ diff --git a/bigframes/operations/datetimes.py b/bigframes/operations/datetimes.py index 1b20c2d593d..a8a33beb577 100644 --- a/bigframes/operations/datetimes.py +++ b/bigframes/operations/datetimes.py @@ -14,12 +14,14 @@ from __future__ import annotations +from bigframes.core import log_adapter 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 +@log_adapter.class_logger class DatetimeMethods( bigframes.operations.base.SeriesMethods, vendordt.DatetimeProperties ): diff --git a/bigframes/operations/strings.py b/bigframes/operations/strings.py index 0545ea34d6a..201b19abe80 100644 --- a/bigframes/operations/strings.py +++ b/bigframes/operations/strings.py @@ -18,6 +18,7 @@ from typing import cast, Literal, Optional, Union import bigframes.constants as constants +from bigframes.core import log_adapter import bigframes.dataframe as df import bigframes.operations as ops import bigframes.operations.base @@ -32,6 +33,7 @@ } +@log_adapter.class_logger class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods): __doc__ = vendorstr.StringMethods.__doc__ diff --git a/bigframes/operations/structs.py b/bigframes/operations/structs.py index 506a5577094..b2ae98f3789 100644 --- a/bigframes/operations/structs.py +++ b/bigframes/operations/structs.py @@ -18,6 +18,7 @@ import ibis.expr.types as ibis_types +from bigframes.core import log_adapter import bigframes.dataframe import bigframes.operations import bigframes.operations.base @@ -38,6 +39,7 @@ def _as_ibis(self, x: ibis_types.Value): return struct_value[name].name(name) +@log_adapter.class_logger class StructAccessor( bigframes.operations.base.SeriesMethods, vendoracessors.StructAccessor ): diff --git a/bigframes/series.py b/bigframes/series.py index 2cffdd5fcec..c929775a000 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -30,6 +30,7 @@ import bigframes.constants as constants import bigframes.core +from bigframes.core import log_adapter import bigframes.core.block_transforms as block_ops import bigframes.core.blocks as blocks import bigframes.core.groupby as groupby @@ -55,6 +56,7 @@ LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]] +@log_adapter.class_logger class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series): def __init__(self, *args, **kwargs): self._query_job: Optional[bigquery.QueryJob] = None diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 1b12e5deeb8..069bd5d2602 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -64,6 +64,7 @@ import bigframes._config.bigquery_options as bigquery_options import bigframes.constants as constants +from bigframes.core import log_adapter import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.guid as guid @@ -1347,6 +1348,10 @@ def _start_query( Starts query job and waits for results. """ job_config = self._prepare_job_config(job_config) + api_methods = log_adapter.get_and_reset_api_methods() + job_config.labels = bigframes_io.create_job_configs_labels( + job_configs_labels=job_config.labels, api_methods=api_methods + ) query_job = self.bqclient.query(sql, job_config=job_config) opts = bigframes.options.display @@ -1381,6 +1386,8 @@ def _prepare_job_config( ) -> bigquery.QueryJobConfig: if job_config is None: job_config = self.bqclient.default_query_job_config + if job_config is None: + job_config = bigquery.QueryJobConfig() if bigframes.options.compute.maximum_bytes_billed is not None: job_config.maximum_bytes_billed = ( bigframes.options.compute.maximum_bytes_billed diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index badc91e3565..dae73301e7b 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -17,17 +17,36 @@ from __future__ import annotations import datetime +import itertools import textwrap import types -from typing import Dict, Iterable, Optional, Union +from typing import Dict, Iterable, Optional, Sequence, Union import uuid import google.cloud.bigquery as bigquery IO_ORDERING_ID = "bqdf_row_nums" +MAX_LABELS_COUNT = 64 TEMP_TABLE_PREFIX = "bqdf{date}_{random_id}" +def create_job_configs_labels( + job_configs_labels: Optional[Dict[str, str]], + api_methods: Sequence[str], +) -> Dict[str, str]: + if job_configs_labels is None: + job_configs_labels = {} + + labels = list( + itertools.chain( + job_configs_labels.keys(), + (f"recent-bigframes-api-{i}" for i in range(len(api_methods))), + ) + ) + values = list(itertools.chain(job_configs_labels.values(), api_methods)) + return dict(zip(labels[:MAX_LABELS_COUNT], values[:MAX_LABELS_COUNT])) + + def create_export_csv_statement( table_id: str, uri: str, field_delimiter: str, header: bool ) -> str: diff --git a/tests/unit/core/test_log_adapter.py b/tests/unit/core/test_log_adapter.py new file mode 100644 index 00000000000..376b7f20758 --- /dev/null +++ b/tests/unit/core/test_log_adapter.py @@ -0,0 +1,60 @@ +# 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.core import log_adapter + +MAX_LABELS_COUNT = 64 + + +@pytest.fixture +def test_instance(): + # Create a simple class for testing + @log_adapter.class_logger + class TestClass: + def method1(self): + pass + + def method2(self): + pass + + return TestClass() + + +def test_method_logging(test_instance): + test_instance.method1() + test_instance.method2() + + # Check if the methods were added to the _api_methods list + api_methods = log_adapter.get_and_reset_api_methods() + assert api_methods is not None + assert "method1" in api_methods + assert "method2" in api_methods + + +def test_add_api_method_limit(test_instance): + # Ensure that add_api_method correctly adds a method to _api_methods + for i in range(70): + test_instance.method2() + assert len(log_adapter._api_methods) == MAX_LABELS_COUNT + + +def test_get_and_reset_api_methods(test_instance): + # Ensure that get_and_reset_api_methods returns a copy and resets the list + test_instance.method1() + test_instance.method2() + previous_methods = log_adapter.get_and_reset_api_methods() + assert previous_methods is not None + assert log_adapter._api_methods == [] diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index 03470208e42..e1481d3f05c 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -19,7 +19,113 @@ import google.cloud.bigquery as bigquery import pytest -import bigframes.session._io.bigquery +import bigframes +from bigframes.core import log_adapter +import bigframes.pandas as bpd +import bigframes.session._io.bigquery as io_bq + + +def test_create_job_configs_labels_is_none(): + api_methods = ["agg", "series-mode"] + labels = io_bq.create_job_configs_labels( + job_configs_labels=None, api_methods=api_methods + ) + expected_dict = { + "recent-bigframes-api-0": "agg", + "recent-bigframes-api-1": "series-mode", + } + assert labels is not None + assert labels == expected_dict + + +def test_create_job_configs_labels_length_limit_not_met(): + cur_labels = { + "bigframes-api": "read_pandas", + "source": "bigquery-dataframes-temp", + } + api_methods = ["agg", "series-mode"] + labels = io_bq.create_job_configs_labels( + job_configs_labels=cur_labels, api_methods=api_methods + ) + expected_dict = { + "bigframes-api": "read_pandas", + "source": "bigquery-dataframes-temp", + "recent-bigframes-api-0": "agg", + "recent-bigframes-api-1": "series-mode", + } + assert labels is not None + assert len(labels) == 4 + assert labels == expected_dict + + +def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit(): + cur_labels = { + "bigframes-api": "read_pandas", + "source": "bigquery-dataframes-temp", + } + df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]}) + # Test running two methods + df.head() + df.max() + api_methods = log_adapter._api_methods + + labels = io_bq.create_job_configs_labels( + job_configs_labels=cur_labels, api_methods=api_methods + ) + expected_dict = { + "bigframes-api": "read_pandas", + "source": "bigquery-dataframes-temp", + "recent-bigframes-api-0": "__init__", + "recent-bigframes-api-1": "max", + "recent-bigframes-api-2": "__init__", + "recent-bigframes-api-3": "head", + "recent-bigframes-api-4": "__init__", + } + assert labels is not None + assert len(labels) == 7 + assert labels == expected_dict + + +def test_create_job_configs_labels_length_limit_met_and_labels_is_none(): + df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]}) + # Test running methods more than the labels' length limit + for i in range(66): + df.head() + api_methods = log_adapter._api_methods + + labels = io_bq.create_job_configs_labels( + job_configs_labels=None, api_methods=api_methods + ) + assert labels is not None + assert len(labels) == 64 + assert "head" in labels.values() + + +def test_create_job_configs_labels_length_limit_met(): + cur_labels = { + "bigframes-api": "read_pandas", + "source": "bigquery-dataframes-temp", + } + for i in range(60): + key = f"bigframes-api-test-{i}" + value = f"test{i}" + cur_labels[key] = value + # If cur_labels length is 62, we can only add one label from api_methods + df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]}) + # Test running two methods + df.head() + df.max() + api_methods = log_adapter._api_methods + + labels = io_bq.create_job_configs_labels( + job_configs_labels=cur_labels, api_methods=api_methods + ) + assert labels is not None + assert len(labels) == 64 + assert "max" in labels.values() + assert "head" not in labels.values() + assert "bigframes-api" in labels.keys() + assert "source" in labels.keys() def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): @@ -125,5 +231,5 @@ def test_create_temp_table_default_expiration(): ), ) def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str): - sql = bigframes.session._io.bigquery.bq_schema_to_sql(schema) + sql = io_bq.bq_schema_to_sql(schema) assert sql == expected diff --git a/tests/unit/test_compute_options.py b/tests/unit/test_compute_options.py index 499a0a5fefa..a613bca7b90 100644 --- a/tests/unit/test_compute_options.py +++ b/tests/unit/test_compute_options.py @@ -18,13 +18,9 @@ def test_maximum_bytes_option(): session = resources.create_bigquery_session() - num_query_calls = 0 with bf.option_context("compute.maximum_bytes_billed", 10000): - # clear initial method calls - session.bqclient.method_calls = [] + session.bqclient.query.reset_mock() session._start_query("query") - for call in session.bqclient.method_calls: - _, _, kwargs = call - num_query_calls += 1 - assert kwargs["job_config"].maximum_bytes_billed == 10000 - assert num_query_calls > 0 + call = session.bqclient.query.call_args + assert call.kwargs["job_config"].maximum_bytes_billed == 10000 + session.bqclient.query.assert_called_once() From e2b7df5530de53f4b6a7d721a8d226c94cabee34 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 15 Nov 2023 00:26:15 +0000 Subject: [PATCH 161/197] chore(main): release 0.14.0 (#183) :robot: I have created a release *beep* *boop* --- ## [0.14.0](https://togithub.com/googleapis/python-bigquery-dataframes/compare/v0.13.0...v0.14.0) (2023-11-14) ### Features * Add 'cross' join support ([#176](https://togithub.com/googleapis/python-bigquery-dataframes/issues/176)) ([765446a](https://togithub.com/googleapis/python-bigquery-dataframes/commit/765446a929abe1ac076c3037afa7892f64105356)) * Add 'index', 'pad', 'nearest' interpolate methods ([#162](https://togithub.com/googleapis/python-bigquery-dataframes/issues/162)) ([6a28403](https://togithub.com/googleapis/python-bigquery-dataframes/commit/6a2840349a23035bdfdabacd1e231b41bbb5ed7a)) * Add series.sample (identical to existing dataframe.sample) ([#187](https://togithub.com/googleapis/python-bigquery-dataframes/issues/187)) ([37914a4](https://togithub.com/googleapis/python-bigquery-dataframes/commit/37914a4077c681881491f5c36d1a9c9f4255e18f)) * Add unordered sql compilation ([#156](https://togithub.com/googleapis/python-bigquery-dataframes/issues/156)) ([58f420c](https://togithub.com/googleapis/python-bigquery-dataframes/commit/58f420c91d94ca085e9810f36513ffe772bfddcf)) * Log most recent API calls as `recent-bigframes-api-xx` labels on BigQuery jobs ([#145](https://togithub.com/googleapis/python-bigquery-dataframes/issues/145)) ([4ea33b7](https://togithub.com/googleapis/python-bigquery-dataframes/commit/4ea33b7433532ae3a386a6ffa9eb57360ea39526)) * Read_gbq creates order deterministically without table copy ([#191](https://togithub.com/googleapis/python-bigquery-dataframes/issues/191)) ([8ab81de](https://togithub.com/googleapis/python-bigquery-dataframes/commit/8ab81dee4d0eee499094f2dd576550f0c59d7551)) * Support `date_series.astype("string[pyarrow]")` to cast DATE to STRING ([#186](https://togithub.com/googleapis/python-bigquery-dataframes/issues/186)) ([aee0e8e](https://togithub.com/googleapis/python-bigquery-dataframes/commit/aee0e8e2518c59bd1e0b07940c3309871fde8899)) * Support `series.at[row_label] = scalar` ([#173](https://togithub.com/googleapis/python-bigquery-dataframes/issues/173)) ([0c8bd33](https://togithub.com/googleapis/python-bigquery-dataframes/commit/0c8bd33806bb99206b8b12dbdf7d7485c6ffb759)) * Temporary resources no longer use BigQuery Sessions ([#194](https://togithub.com/googleapis/python-bigquery-dataframes/issues/194)) ([4a02cac](https://togithub.com/googleapis/python-bigquery-dataframes/commit/4a02cac88c7d7b46bed1fa813a862fc2ef9ef084)) ### Bug Fixes * All sort operation are now stable ([#195](https://togithub.com/googleapis/python-bigquery-dataframes/issues/195)) ([3a2761f](https://togithub.com/googleapis/python-bigquery-dataframes/commit/3a2761f3c38d0de8b8eda47fffa15b8412aa84b0)) * Default to 7 days expiration for `read_csv`, `read_json`, `read_parquet` ([#193](https://togithub.com/googleapis/python-bigquery-dataframes/issues/193)) ([03606cd](https://togithub.com/googleapis/python-bigquery-dataframes/commit/03606cda30eb7645bfd4534460112dcca56b0ab0)) * Deprecate the `remote_service_type` in llm model ([#180](https://togithub.com/googleapis/python-bigquery-dataframes/issues/180)) ([a8a409a](https://togithub.com/googleapis/python-bigquery-dataframes/commit/a8a409ab0bd1f99dfb442df0703bf8786e0fe58e)) * For reset_index on unnamed multiindex, always use level_[n] label ([#182](https://togithub.com/googleapis/python-bigquery-dataframes/issues/182)) ([f95000d](https://togithub.com/googleapis/python-bigquery-dataframes/commit/f95000d3f88662be4d88c8b0152f1b838e99ec55)) * Match pandas behavior when assigning listlike to empty dfs ([#172](https://togithub.com/googleapis/python-bigquery-dataframes/issues/172)) ([c1d1f42](https://togithub.com/googleapis/python-bigquery-dataframes/commit/c1d1f42a21cc089877f79ebb46a39ddef6958e04)) * Use anonymous dataset instead of session dataset for temp tables ([#181](https://togithub.com/googleapis/python-bigquery-dataframes/issues/181)) ([800d44e](https://togithub.com/googleapis/python-bigquery-dataframes/commit/800d44eb5eb77da5d87b2e005f5a2ed53842e7b5)) * Use random table for `read_pandas` ([#192](https://togithub.com/googleapis/python-bigquery-dataframes/issues/192)) ([741c75e](https://togithub.com/googleapis/python-bigquery-dataframes/commit/741c75e5797e26a1487ff3da76a07953d9537f3f)) * Use random table when loading data for `read_csv`, `read_json`, `read_parquet` ([#175](https://togithub.com/googleapis/python-bigquery-dataframes/issues/175)) ([9d2e6dc](https://togithub.com/googleapis/python-bigquery-dataframes/commit/9d2e6dc1ae4e11e80da4aabe0daa3a6044137cc6)) ### Documentation * Add code samples for `read_gbq_function` using community UDFs ([#188](https://togithub.com/googleapis/python-bigquery-dataframes/issues/188)) ([7506eab](https://togithub.com/googleapis/python-bigquery-dataframes/commit/7506eabf2e58159507809e36abfe90c417dfe92f)) * Add docstring code samples for `Series.apply` and `DataFrame.map` ([#185](https://togithub.com/googleapis/python-bigquery-dataframes/issues/185)) ([c816d84](https://togithub.com/googleapis/python-bigquery-dataframes/commit/c816d843e6f3c5a944cd4395ed0e1e91cec49812)) * Add llm kmeans notebook as an included example ([#177](https://togithub.com/googleapis/python-bigquery-dataframes/issues/177)) ([d49ae42](https://togithub.com/googleapis/python-bigquery-dataframes/commit/d49ae42a379fafd601cc94227e7f8f14b3d5f8c3)) * Use `head()` to get top `n` results, not to preview results ([#190](https://togithub.com/googleapis/python-bigquery-dataframes/issues/190)) ([87f84c9](https://togithub.com/googleapis/python-bigquery-dataframes/commit/87f84c9e58e7d0ea521ac386c9f02791cdddd19f)) --- 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 | 35 +++++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc327b2e966..1f76b782728 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,41 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.14.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.13.0...v0.14.0) (2023-11-14) + + +### Features + +* Add 'cross' join support ([#176](https://github.com/googleapis/python-bigquery-dataframes/issues/176)) ([765446a](https://github.com/googleapis/python-bigquery-dataframes/commit/765446a929abe1ac076c3037afa7892f64105356)) +* Add 'index', 'pad', 'nearest' interpolate methods ([#162](https://github.com/googleapis/python-bigquery-dataframes/issues/162)) ([6a28403](https://github.com/googleapis/python-bigquery-dataframes/commit/6a2840349a23035bdfdabacd1e231b41bbb5ed7a)) +* Add series.sample (identical to existing dataframe.sample) ([#187](https://github.com/googleapis/python-bigquery-dataframes/issues/187)) ([37914a4](https://github.com/googleapis/python-bigquery-dataframes/commit/37914a4077c681881491f5c36d1a9c9f4255e18f)) +* Add unordered sql compilation ([#156](https://github.com/googleapis/python-bigquery-dataframes/issues/156)) ([58f420c](https://github.com/googleapis/python-bigquery-dataframes/commit/58f420c91d94ca085e9810f36513ffe772bfddcf)) +* Log most recent API calls as `recent-bigframes-api-xx` labels on BigQuery jobs ([#145](https://github.com/googleapis/python-bigquery-dataframes/issues/145)) ([4ea33b7](https://github.com/googleapis/python-bigquery-dataframes/commit/4ea33b7433532ae3a386a6ffa9eb57360ea39526)) +* Read_gbq creates order deterministically without table copy ([#191](https://github.com/googleapis/python-bigquery-dataframes/issues/191)) ([8ab81de](https://github.com/googleapis/python-bigquery-dataframes/commit/8ab81dee4d0eee499094f2dd576550f0c59d7551)) +* Support `date_series.astype("string[pyarrow]")` to cast DATE to STRING ([#186](https://github.com/googleapis/python-bigquery-dataframes/issues/186)) ([aee0e8e](https://github.com/googleapis/python-bigquery-dataframes/commit/aee0e8e2518c59bd1e0b07940c3309871fde8899)) +* Support `series.at[row_label] = scalar` ([#173](https://github.com/googleapis/python-bigquery-dataframes/issues/173)) ([0c8bd33](https://github.com/googleapis/python-bigquery-dataframes/commit/0c8bd33806bb99206b8b12dbdf7d7485c6ffb759)) +* Temporary resources no longer use BigQuery Sessions ([#194](https://github.com/googleapis/python-bigquery-dataframes/issues/194)) ([4a02cac](https://github.com/googleapis/python-bigquery-dataframes/commit/4a02cac88c7d7b46bed1fa813a862fc2ef9ef084)) + + +### Bug Fixes + +* All sort operation are now stable ([#195](https://github.com/googleapis/python-bigquery-dataframes/issues/195)) ([3a2761f](https://github.com/googleapis/python-bigquery-dataframes/commit/3a2761f3c38d0de8b8eda47fffa15b8412aa84b0)) +* Default to 7 days expiration for `read_csv`, `read_json`, `read_parquet` ([#193](https://github.com/googleapis/python-bigquery-dataframes/issues/193)) ([03606cd](https://github.com/googleapis/python-bigquery-dataframes/commit/03606cda30eb7645bfd4534460112dcca56b0ab0)) +* Deprecate the `remote_service_type` in llm model ([#180](https://github.com/googleapis/python-bigquery-dataframes/issues/180)) ([a8a409a](https://github.com/googleapis/python-bigquery-dataframes/commit/a8a409ab0bd1f99dfb442df0703bf8786e0fe58e)) +* For reset_index on unnamed multiindex, always use level_[n] label ([#182](https://github.com/googleapis/python-bigquery-dataframes/issues/182)) ([f95000d](https://github.com/googleapis/python-bigquery-dataframes/commit/f95000d3f88662be4d88c8b0152f1b838e99ec55)) +* Match pandas behavior when assigning listlike to empty dfs ([#172](https://github.com/googleapis/python-bigquery-dataframes/issues/172)) ([c1d1f42](https://github.com/googleapis/python-bigquery-dataframes/commit/c1d1f42a21cc089877f79ebb46a39ddef6958e04)) +* Use anonymous dataset instead of session dataset for temp tables ([#181](https://github.com/googleapis/python-bigquery-dataframes/issues/181)) ([800d44e](https://github.com/googleapis/python-bigquery-dataframes/commit/800d44eb5eb77da5d87b2e005f5a2ed53842e7b5)) +* Use random table for `read_pandas` ([#192](https://github.com/googleapis/python-bigquery-dataframes/issues/192)) ([741c75e](https://github.com/googleapis/python-bigquery-dataframes/commit/741c75e5797e26a1487ff3da76a07953d9537f3f)) +* Use random table when loading data for `read_csv`, `read_json`, `read_parquet` ([#175](https://github.com/googleapis/python-bigquery-dataframes/issues/175)) ([9d2e6dc](https://github.com/googleapis/python-bigquery-dataframes/commit/9d2e6dc1ae4e11e80da4aabe0daa3a6044137cc6)) + + +### Documentation + +* Add code samples for `read_gbq_function` using community UDFs ([#188](https://github.com/googleapis/python-bigquery-dataframes/issues/188)) ([7506eab](https://github.com/googleapis/python-bigquery-dataframes/commit/7506eabf2e58159507809e36abfe90c417dfe92f)) +* Add docstring code samples for `Series.apply` and `DataFrame.map` ([#185](https://github.com/googleapis/python-bigquery-dataframes/issues/185)) ([c816d84](https://github.com/googleapis/python-bigquery-dataframes/commit/c816d843e6f3c5a944cd4395ed0e1e91cec49812)) +* Add llm kmeans notebook as an included example ([#177](https://github.com/googleapis/python-bigquery-dataframes/issues/177)) ([d49ae42](https://github.com/googleapis/python-bigquery-dataframes/commit/d49ae42a379fafd601cc94227e7f8f14b3d5f8c3)) +* Use `head()` to get top `n` results, not to preview results ([#190](https://github.com/googleapis/python-bigquery-dataframes/issues/190)) ([87f84c9](https://github.com/googleapis/python-bigquery-dataframes/commit/87f84c9e58e7d0ea521ac386c9f02791cdddd19f)) + ## [0.13.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.12.0...v0.13.0) (2023-11-07) diff --git a/bigframes/version.py b/bigframes/version.py index 0a5df274799..5a94f72649a 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.13.0" +__version__ = "0.14.0" From f957b278b39e0a472a3153e9e1906c2d5f2ac2e5 Mon Sep 17 00:00:00 2001 From: Henry Solberg Date: Tue, 14 Nov 2023 18:52:15 -0800 Subject: [PATCH 162/197] docs: add an example notebook about line graphs (#197) 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_covid_line_graphs.ipynb | 598 ++++++++++++++++++ noxfile.py | 1 + 2 files changed, 599 insertions(+) create mode 100644 notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb diff --git a/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb b/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb new file mode 100644 index 00000000000..8b18cc89674 --- /dev/null +++ b/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb @@ -0,0 +1,598 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9GIt_orUtNvA" + }, + "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": "h7AT6h2ItNvD" + }, + "source": [ + "## Use BigQuery DataFrames to visualize COVID-19 data\n", + "\n", + "\n", + "\n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + "
" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "n-MFJQxLtNvE" + }, + "source": [ + "## Overview\n", + "\n", + "The goal of this notebook is to demonstrate creating line graphs from a ~20 million-row BigQuery dataset using BigQuery DataFrames. We will first create a plain line graph using matplotlip, then we will downsample and download our data to create a graph with a line of best fit using seaborn.\n", + "\n", + "If you're like me, during 2020 (and/or later years) you often found yourself looking at charts like [these](https://health.google.com/covid-19/open-data/explorer/statistics) visualizing COVID-19 cases over time. For our first graph, we're going to recreate one of those charts by filtering, summing, and then graphing COVID-19 data from the United States. BigQuery DataFrame's default integration with matplotlib will get us a satisfying result for this first graph.\n", + "\n", + "For our second graph, though, we want to use a scatterplot with a line of best fit, something that matplotlib will not do for us automatically. So, we'll demonstrate how to downsample our data and use seaborn to make our plot. Our second graph will be of symptom-related search trends against new cases of COVID-19, so we'll see if searches for things like \"cough\" and \"fever\" are more common in the places and times where more new cases of COVID-19 occur." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "ffqBzbNztNvF" + }, + "source": [ + "### Dataset\n", + "\n", + "This notebook uses the [BigQuery COVID-19 Open Data](https://pantheon.corp.google.com/marketplace/product/bigquery-public-datasets/covid19-open-data). In this dataset, each row represents a new observation of the COVID-19 situation in a particular time and place. We will use the \"new_confirmed\" column, which contains the number of new COVID-19 cases at each observation, along with the \"search_trends_cough\", \"search_trends_fever\", and \"search_trends_bruise\" columns, which are [Google Trends](https://trends.google.com/trends/) data for searches related to cough, fever, and bruises. In the first section of the notebook, we will also use the \"country_code\" and \"date\" columns to compile one data point per day for a particular country." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Nf__tMR-tNvF" + }, + "source": [ + "### Costs\n", + "\n", + "This tutorial uses billable components of Google Cloud:\n", + "\n", + "* BigQuery (compute)\n", + "\n", + "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models),\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": "7_rsbkCktNvG" + }, + "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": "XZKC6iMFxmMG" + }, + "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": "4aooKMmnxrWF" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "pv5A8Tm-yC1U" + }, + "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": "bk03Rt_HyGx-" + }, + "outputs": [], + "source": [ + "REGION = \"US\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "B9RWxD1btNvK" + }, + "source": [ + "Now we are ready to use BigQuery DataFrames!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "wJ0gXezj2w1t" + }, + "source": [ + "## Visualization #1: Cases over time in the US" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "xckgWno6ouHY" + }, + "source": [ + "### Set up project and filter data" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "-uiY0hh4tNvK" + }, + "source": [ + "First, let's do project setup. We use options to tell BigQuery DataFrames what project and what region to use for our cloud computing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "R7STCS8xB5d2" + }, + "outputs": [], + "source": [ + "import bigframes.pandas as bf\n", + "\n", + "bf.options.bigquery.project = PROJECT_ID\n", + "bf.options.bigquery.location = REGION" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "v6FGschEowht" + }, + "source": [ + "Next, we read the data from a publicly available BigQuery dataset. This will take ~1 minute." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "zDSwoBo1CU3G" + }, + "outputs": [], + "source": [ + "all_data = bf.read_gbq(\"bigquery-public-data.covid19_open_data.covid19_open_data\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9qV2y3iHp13y" + }, + "source": [ + "Using pandas syntax, we will select from our all_data input dataframe only those rows where the country_code is US. This is called row filtering." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "UjMT_qhjf8Fu" + }, + "outputs": [], + "source": [ + "usa_data = all_data[all_data[\"country_code\"] == \"US\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "IYCUayWkwq8c" + }, + "source": [ + "We're only concerned with the date and the total number of confirmed cases for now, so select just those two columns as well." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "IaoUf57ZwrJ8" + }, + "outputs": [], + "source": [ + "usa_data = usa_data[[\"date\", \"new_confirmed\"]]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "94oqNRnDvGkr" + }, + "source": [ + "### Sum data" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "TNCQWZW83U0b" + }, + "source": [ + "`usa_data.groupby(\"date\")` will give us a groupby object that lets us perform operations on groups of rows with the same date. We call sum on that object to get the sum for each day. This process might be familiar to pandas users." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tYDoaKgJChiq" + }, + "outputs": [], + "source": [ + "# numeric_only = True because we don't want to sum dates\n", + "new_cases_usa = usa_data.groupby(\"date\").sum(numeric_only = True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3jcwFPgK5BLh" + }, + "source": [ + "### Line graph" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8GvJAgnH5Nzi" + }, + "source": [ + "BigQuery DataFrames implements some of the interface required by matplotlib. This means we can pass our DataFrame right into `pyplot.plt` and using the default settings, matplotlib will draw a simple line graph for us." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "gFbCgfFC2gHw" + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "# matplotlin will draw a line graph by default\n", + "plt.plot(new_cases_usa)\n", + "# Rotate the labels on the x axis so that they don't overlap\n", + "plt.xticks(rotation=45)\n", + "# label the y axis for clarity\n", + "plt.ylabel(\"New Cases\")\n", + "\n", + "# Show the plot\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sM5-HFDx70RG" + }, + "source": [ + "## Visualization #2: Symptom-related searches compared to new cases" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "se1b6Vf4XB9_" + }, + "source": [ + "### Filter data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Wl2o-NYMoygb" + }, + "source": [ + "We're curious if searches for symptoms like \"cough\" and \"fever\" went up in the same times and places that new COVID-19 cases occured, compared to non-symptoms like \"bruise.\" Let's plot searches vs. new cases to see if it looks like there's a correlation." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "olfnCzyg8jYi" + }, + "source": [ + "First, we select the new cases column and the search trends we're interested in." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "LqqHzjty8jk0" + }, + "outputs": [], + "source": [ + "symptom_data = all_data[[\"new_confirmed\", \"search_trends_cough\", \"search_trends_fever\", \"search_trends_bruise\"]]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b3DlJX-k9SPk" + }, + "source": [ + "Not all rows have data for all of these columns, so let's select only the rows that do." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "g4MeM8Oe9Q6X" + }, + "outputs": [], + "source": [ + "symptom_data = symptom_data.dropna()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "IlXt__om9QYI" + }, + "source": [ + "We want to use a line of best fit to make the correlation stand out. Matplotlib does not include a feature for lines of best fit, but seaborn, which is built on matplotlib, does.\n", + "\n", + "BigQuery DataFrames does not currently integrate with seaborn by default. So we will demonstrate how to downsample and download a DataFrame, and use seaborn on the downloaded data." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "MmfgKMaEXNbL" + }, + "source": [ + "### Downsample and download" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "wIuG1JRTPAk9" + }, + "source": [ + "BigQuery DataFrames options let us set up the sampling functionality we need. Calls to `to_pandas()` usually download all the data available in our BigQuery table and store it locally as a pandas DataFrame. `pd.options.sampling.enable_downsampling = True` will make future calls to `to_pandas` use downsampling to download only part of the data, and `pd.options.sampling.max_download_size` allows us to set the amount of data to download." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "x95ZgBkyDMP4" + }, + "outputs": [], + "source": [ + "bf.options.sampling.enable_downsampling = True # enable downsampling\n", + "bf.options.sampling.max_download_size = 5 # download only 5 mb of data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "C6sCXkrQPJC_" + }, + "source": [ + "Download the data and note the message letting us know that downsampling is being used." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "V0OK02D7PJSL" + }, + "outputs": [], + "source": [ + "local_symptom_data = symptom_data.to_pandas(sampling_method=\"uniform\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "T9Hub_EAXWvY" + }, + "source": [ + "### Graph with lines of best fit" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "hoQ9TPgUPJnN" + }, + "source": [ + "We will now use seaborn to make the plots with the lines of best fit for cough, fever, and bruise. Note that since we're working with a local pandas dataframe, you could use any other Python library or technique you're familiar with, but we'll stick to seaborn for this notebook.\n", + "\n", + "Seaborn will take a few seconds to calculate the lines. Since cough and fever are symptoms of COVID-19, but bruising isn't, we expect the slope of the line of best fit to be positive in the first two graphs, but not the third, indicating that there is a correlation between new COVID-19 cases and cough- and fever-related searches." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EG7qM3R18bOb" + }, + "outputs": [], + "source": [ + "import seaborn as sns\n", + "\n", + "# first, convert to a data type that is suitable for seaborn\n", + "local_symptom_data[\"new_confirmed\"] = \\\n", + " local_symptom_data[\"new_confirmed\"].astype(float)\n", + "local_symptom_data[\"search_trends_cough\"] = \\\n", + " local_symptom_data[\"search_trends_cough\"].astype(float)\n", + "\n", + "# draw the graph. This might take ~30 seconds.\n", + "sns.regplot(x=\"new_confirmed\", y=\"search_trends_cough\", data=local_symptom_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5nVy61rEGaM4" + }, + "outputs": [], + "source": [ + "# similarly, for fever\n", + "\n", + "local_symptom_data[\"search_trends_fever\"] = \\\n", + " local_symptom_data[\"search_trends_fever\"].astype(float)\n", + "sns.regplot(x=\"new_confirmed\", y=\"search_trends_fever\", data=local_symptom_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "-S1A9E3WGaYH" + }, + "outputs": [], + "source": [ + "# similarly, for bruise\n", + "local_symptom_data[\"search_trends_bruise\"] = \\\n", + " local_symptom_data[\"search_trends_bruise\"].astype(float)\n", + "sns.regplot(\n", + " x=\"new_confirmed\",\n", + " y=\"search_trends_bruise\",\n", + " data=local_symptom_data\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Hd2A8707Uhz2" + }, + "source": [ + "We see that the slope of the line is positive in the graphs for cough and fever, but flat for bruise. That means that in places with increasing new cases of COVID-19, we saw increasing searches for cough and fever, but we didn't see increasing searches for unrelated symptoms like bruises. Interesting!" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Recap" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We used matplotlib to draw a line graph of COVID-19 cases over time in the USA. Then, we used downsampling to download only a portion of the available data, and used seaborn locally to plot lines of best fit to observe corellation between COVID-19 cases and searches for related vs. unrelated symptoms.\n", + "\n", + "Thank you for using BigQuery DataFrames!" + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "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.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/noxfile.py b/noxfile.py index 3dd23ba04f6..da61232fc79 100644 --- a/noxfile.py +++ b/noxfile.py @@ -615,6 +615,7 @@ def notebook(session): "notebooks/vertex_sdk/sdk2_bigframes_pytorch.ipynb", "notebooks/vertex_sdk/sdk2_bigframes_sklearn.ipynb", "notebooks/vertex_sdk/sdk2_bigframes_tensorflow.ipynb", + "notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb", # The experimental notebooks imagine features that don't yet # exist or only exist as temporary prototypes. "notebooks/experimental/longer_ml_demo.ipynb", From 8324f133547ec35da5eefc0a8b02fe0f3887d81d Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Wed, 15 Nov 2023 18:58:14 -0800 Subject: [PATCH 163/197] fix: correctly handle null values when initializing fingerprint ordering (#210) 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/session/__init__.py | 9 ++++++--- tests/system/small/test_dataframe.py | 8 ++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 069bd5d2602..928123ce749 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1120,8 +1120,9 @@ def _create_total_ordering( ordering_hash_part = guid.generate_guid("bigframes_ordering_") ordering_rand_part = guid.generate_guid("bigframes_ordering_") + # All inputs into hash must be non-null or resulting hash will be null str_values = list( - map(lambda col: _convert_to_string(table[col]), table.columns) + map(lambda col: _convert_to_nonnull_string(table[col]), table.columns) ) full_row_str = ( str_values[0].concat(*str_values[1:]) @@ -1419,7 +1420,7 @@ def _can_cluster_bq(field: bigquery.SchemaField): ) -def _convert_to_string(column: ibis_types.Column) -> ibis_types.StringColumn: +def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringValue: col_type = column.type() if ( col_type.is_numeric() @@ -1436,4 +1437,6 @@ def _convert_to_string(column: ibis_types.Column) -> ibis_types.StringColumn: # TO_JSON_STRING works with all data types, but isn't the most efficient # Needed for JSON, STRUCT and ARRAY datatypes result = vendored_ibis_ops.ToJsonString(column).to_expr() # type: ignore - return typing.cast(ibis_types.StringColumn, result) + # Escape backslashes and use backslash as delineator + escaped = typing.cast(ibis_types.StringColumn, result.fillna("")).replace("\\", "\\\\") # type: ignore + return typing.cast(ibis_types.StringColumn, ibis.literal("\\")).concat(escaped) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index e522878229a..a0cf25807cf 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2703,6 +2703,14 @@ def test_sample(scalars_dfs, frac, n, random_state): assert bf_result.shape[1] == scalars_df.shape[1] +def test_sample_determinism(penguins_df_default_index): + df = penguins_df_default_index.sample(n=100, random_state=12345).head(15) + bf_result = df.to_pandas() + bf_result2 = df.to_pandas() + + pandas.testing.assert_frame_equal(bf_result, bf_result2) + + def test_sample_raises_value_error(scalars_dfs): scalars_df, _ = scalars_dfs with pytest.raises( From 5ab5059f7db5d0f2be735dca76bc8e5163287c4d Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 15 Nov 2023 20:25:53 -0800 Subject: [PATCH 164/197] chore(main): release 0.14.1 (#207) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 12 ++++++++++++ bigframes/version.py | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f76b782728..091967513aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,18 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.14.1](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.14.0...v0.14.1) (2023-11-16) + + +### Bug Fixes + +* Correctly handle null values when initializing fingerprint ordering ([#210](https://github.com/googleapis/python-bigquery-dataframes/issues/210)) ([8324f13](https://github.com/googleapis/python-bigquery-dataframes/commit/8324f133547ec35da5eefc0a8b02fe0f3887d81d)) + + +### Documentation + +* Add an example notebook about line graphs ([#197](https://github.com/googleapis/python-bigquery-dataframes/issues/197)) ([f957b27](https://github.com/googleapis/python-bigquery-dataframes/commit/f957b278b39e0a472a3153e9e1906c2d5f2ac2e5)) + ## [0.14.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.13.0...v0.14.0) (2023-11-14) diff --git a/bigframes/version.py b/bigframes/version.py index 5a94f72649a..46e57e5b888 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.14.0" +__version__ = "0.14.1" From e8532b1d999d26ea1ebdd30efb8f2c0a93a6a28d Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 16 Nov 2023 10:08:27 -0800 Subject: [PATCH 165/197] fix: polish the llm+kmeans notebook (#208) --- .../bq_dataframes_llm_code_generation.ipynb | 2 +- .../bq_dataframes_llm_kmeans.ipynb | 1181 +++++++++++++++-- 2 files changed, 1057 insertions(+), 126 deletions(-) diff --git a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index 0f113b84c6d..0a41447a53e 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -34,7 +34,7 @@ "\n", "\n", " \n", diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb index 46c49552889..ae038136393 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -31,7 +31,7 @@ "
\n", - " \n", + " \n", " \"Colab Run in Colab\n", " \n", "
\n", "\n", " \n", @@ -118,14 +118,10 @@ "\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,run.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com,cloudresourcemanager.googleapis.com) to enable the following APIs:\n", + "3. [Click here](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com,bigqueryconnection.googleapis.com,aiplatform.googleapis.com) to enable the following APIs:\n", "\n", " * BigQuery API\n", " * BigQuery Connection 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)." @@ -143,9 +139,17 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Updated property [core/project].\n" + ] + } + ], "source": [ "# set your project ID below\n", "PROJECT_ID = \"\" # @param {type:\"string\"}\n", @@ -166,7 +170,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -232,87 +236,6 @@ "# auth.authenticate_user()" ] }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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." - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Connect to Vertex AI\n", - "\n", - "In order to use PaLM2TextGenerator, we will need to set up a [cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from google.cloud import bigquery_connection_v1 as bq_connection\n", - "\n", - "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)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "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:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "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'" - ] - }, { "attachments": {}, "cell_type": "markdown", @@ -336,12 +259,12 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Project Setup" + "BigQuery DataFrames setup" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "metadata": { "id": "R7STCS8xB5d2" }, @@ -353,6 +276,14 @@ "bf.options.bigquery.location = REGION" ] }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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." + ] + }, { "attachments": {}, "cell_type": "markdown", @@ -365,7 +296,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "metadata": { "id": "zDSwoBo1CU3G" }, @@ -376,11 +307,101 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "metadata": { "id": "tYDoaKgJChiq" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 9f096761-e3b5-4d58-a9f7-485ced67afca is DONE. 2.3 GB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job ee8fecb1-2e30-407d-9e2e-9e76061da9e7 is DONE. 2.3 GB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "
\n", - " \n", + " \n", " \"Colab Run in Colab\n", " \n", "
\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
consumer_complaint_narrative
0I signed a contract as a condition of employme...
1First, I want to disclose that XXXX and XXXX b...
2Frequent calls from Focused Receivables Manage...
3I recently contacted Enhanced Recovery Company...
4This began when I subscribed to XXXX XXXX inte...
\n", + "

5 rows × 1 columns

\n", + "[5 rows x 1 columns in total]" + ], + "text/plain": [ + " consumer_complaint_narrative\n", + "0 I signed a contract as a condition of employme...\n", + "1 First, I want to disclose that XXXX and XXXX b...\n", + "2 Frequent calls from Focused Receivables Manage...\n", + "3 I recently contacted Enhanced Recovery Company...\n", + "4 This began when I subscribed to XXXX XXXX inte...\n", + "\n", + "[5 rows x 1 columns]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "issues_df = input_df[[\"consumer_complaint_narrative\"]].dropna()\n", "issues_df.head(n=5) # View the first five complaints" @@ -391,12 +412,12 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Download 10000 complaints to use with PaLM2TextEmbeddingGenerator" + "Downsample DataFrame to 10,000 records for model training." ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "metadata": { "id": "OltYSUEcsSOW" }, @@ -418,11 +439,24 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "metadata": { "id": "li38q8FzDDMu" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 52d2e961-7896-497c-8b03-ab7374737679 is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], "source": [ "from bigframes.ml.llm import PaLM2TextEmbeddingGenerator\n", "\n", @@ -431,11 +465,125 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 29, "metadata": { "id": "cOuSOQ5FDewD" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job d093d51a-8eda-442f-80cd-568cb76e00b3 is DONE. 10.6 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 6419df65-3e96-41a7-a7b5-3d058e18763a is DONE. 80.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 917f09ea-c468-4363-a856-b1091e5f775f is DONE. 80.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 5c9679e7-192c-40b5-a14b-edc0fa113eaa is DONE. 61.5 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "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", + "
text_embedding
422[-0.012013785541057587, 0.003669967409223318, ...
616[-0.014948881231248379, -0.04672442376613617, ...
833[-0.01951478235423565, -0.027120858430862427, ...
1370[-0.03140445053577423, -0.048797041177749634, ...
1430[-0.02244548313319683, -0.03336532413959503, 0...
\n", + "

5 rows × 1 columns

\n", + "
[5 rows x 1 columns in total]" + ], + "text/plain": [ + " text_embedding\n", + "422 [-0.012013785541057587, 0.003669967409223318, ...\n", + "616 [-0.014948881231248379, -0.04672442376613617, ...\n", + "833 [-0.01951478235423565, -0.027120858430862427, ...\n", + "1370 [-0.03140445053577423, -0.048797041177749634, ...\n", + "1430 [-0.02244548313319683, -0.03336532413959503, 0...\n", + "\n", + "[5 rows x 1 columns]" + ] + }, + "execution_count": 29, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# Will take ~3 minutes to compute the embeddings\n", "predicted_embeddings = model.predict(downsampled_issues_df)\n", @@ -445,14 +593,263 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 30, "metadata": { "id": "4H_etYfsEOFP" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job ce9cb0f9-4b0d-40a1-81f3-d6e60dd6c684 is DONE. 160.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job aa692a30-5706-46ad-8029-faf2fac66234 is DONE. 72.2 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "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", + "
consumer_complaint_narrativetext_embedding
2580664Hello, my name is XXXX XXXX, and I am writing ...[0.0003211698785889894, -0.01816680282354355, ...
1806973This is XXXX XXXX and I am submitting this com...[-0.009485247544944286, -0.025846892967820168,...
2055053XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX...[-0.010950954630970955, -0.0249345600605011, 0...
2515231When I reinvestigated my credit report, I real...[-0.009660656563937664, -0.05793113633990288, ...
2633049Checking my credit report XX/XX/2018 with all ...[-0.0022159104701131582, -0.03330004960298538,...
3117273I contacted TransUnion and spoke a credit rep ...[-0.015955328941345215, -0.006488671060651541,...
698814XXXX XXXX XXXX. makes daily calls to me cell c...[0.005397460889071226, -0.01276913657784462, 0...
267826Can we please reopen Case : XXXX? \n", + "\n", + "Wells Farg...[0.004065403249114752, -0.0005381882656365633,...
54019My rights under 15 USC 1681 have been violated...[0.013823015615344048, -0.02010691538453102, 0...
141050To whom it may concern : My personal informati...[0.008104532025754452, -0.01856449618935585, 0...
2962076I have had a CashApp account since last year, ...[-0.0003019514260813594, -0.03750108182430267,...
2481105that some of the information was erroneous. Th...[-0.014868081547319889, -0.0443895161151886, -...
431562I have disputed the referenced accounts to the...[-0.0020524838473647833, -0.04830990731716156,...
1953029On, XX/XX/22, I attempted to complete a transa...[-0.01599179394543171, -0.0074900356121361256,...
2395979Subject : XXXX XXXX XXXX compensation, refund,...[-0.0035950862802565098, -0.014652969315648079...
455524I paid off my mortgage on XX/XX/2019. The comp...[-0.01100730150938034, -0.03495829552412033, 0...
2155924This kind of account is placed as a charged of...[-0.028635455295443535, -0.028604287654161453,...
1069497This is one of many issues I have had with Wel...[0.008871790021657944, -0.028502725064754486, ...
3181689I have disputed this account with MONTEREY FIN...[-0.004721717908978462, -0.03673810139298439, ...
274268Lender is not updating my loan status in the V...[-0.009221495129168034, -0.0289347805082798, 0...
1671305XXXX is a peer to peer lending conmpany that u...[-0.02911308966577053, -0.01850792020559311, -...
886026( DISPUTE CODE - XXXX ) My personal informatio...[-0.007220877334475517, -0.016615957021713257,...
1044431I filed a complaint against PNC this year and ...[0.002848619595170021, -0.035117778927087784, ...
1938481I applied for a modification and was approved....[-0.03114932030439377, -0.0421406552195549, 0....
1987834Ive been Disputting my XXXX XXXX I opened this...[-0.009406660683453083, -0.020967338234186172,...
\n", + "

25 rows × 2 columns

\n", + "
[10000 rows x 2 columns in total]" + ], + "text/plain": [ + " consumer_complaint_narrative \\\n", + "2580664 Hello, my name is XXXX XXXX, and I am writing ... \n", + "1806973 This is XXXX XXXX and I am submitting this com... \n", + "2055053 XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX... \n", + "2515231 When I reinvestigated my credit report, I real... \n", + "2633049 Checking my credit report XX/XX/2018 with all ... \n", + "3117273 I contacted TransUnion and spoke a credit rep ... \n", + "698814 XXXX XXXX XXXX. makes daily calls to me cell c... \n", + "267826 Can we please reopen Case : XXXX? \n", + "\n", + "Wells Farg... \n", + "54019 My rights under 15 USC 1681 have been violated... \n", + "141050 To whom it may concern : My personal informati... \n", + "2962076 I have had a CashApp account since last year, ... \n", + "2481105 that some of the information was erroneous. Th... \n", + "431562 I have disputed the referenced accounts to the... \n", + "1953029 On, XX/XX/22, I attempted to complete a transa... \n", + "2395979 Subject : XXXX XXXX XXXX compensation, refund,... \n", + "455524 I paid off my mortgage on XX/XX/2019. The comp... \n", + "2155924 This kind of account is placed as a charged of... \n", + "1069497 This is one of many issues I have had with Wel... \n", + "3181689 I have disputed this account with MONTEREY FIN... \n", + "274268 Lender is not updating my loan status in the V... \n", + "1671305 XXXX is a peer to peer lending conmpany that u... \n", + "886026 ( DISPUTE CODE - XXXX ) My personal informatio... \n", + "1044431 I filed a complaint against PNC this year and ... \n", + "1938481 I applied for a modification and was approved.... \n", + "1987834 Ive been Disputting my XXXX XXXX I opened this... \n", + "\n", + " text_embedding \n", + "2580664 [0.0003211698785889894, -0.01816680282354355, ... \n", + "1806973 [-0.009485247544944286, -0.025846892967820168,... \n", + "2055053 [-0.010950954630970955, -0.0249345600605011, 0... \n", + "2515231 [-0.009660656563937664, -0.05793113633990288, ... \n", + "2633049 [-0.0022159104701131582, -0.03330004960298538,... \n", + "3117273 [-0.015955328941345215, -0.006488671060651541,... \n", + "698814 [0.005397460889071226, -0.01276913657784462, 0... \n", + "267826 [0.004065403249114752, -0.0005381882656365633,... \n", + "54019 [0.013823015615344048, -0.02010691538453102, 0... \n", + "141050 [0.008104532025754452, -0.01856449618935585, 0... \n", + "2962076 [-0.0003019514260813594, -0.03750108182430267,... \n", + "2481105 [-0.014868081547319889, -0.0443895161151886, -... \n", + "431562 [-0.0020524838473647833, -0.04830990731716156,... \n", + "1953029 [-0.01599179394543171, -0.0074900356121361256,... \n", + "2395979 [-0.0035950862802565098, -0.014652969315648079... \n", + "455524 [-0.01100730150938034, -0.03495829552412033, 0... \n", + "2155924 [-0.028635455295443535, -0.028604287654161453,... \n", + "1069497 [0.008871790021657944, -0.028502725064754486, ... \n", + "3181689 [-0.004721717908978462, -0.03673810139298439, ... \n", + "274268 [-0.009221495129168034, -0.0289347805082798, 0... \n", + "1671305 [-0.02911308966577053, -0.01850792020559311, -... \n", + "886026 [-0.007220877334475517, -0.016615957021713257,... \n", + "1044431 [0.002848619595170021, -0.035117778927087784, ... \n", + "1938481 [-0.03114932030439377, -0.0421406552195549, 0.... \n", + "1987834 [-0.009406660683453083, -0.020967338234186172,... \n", + "...\n", + "\n", + "[10000 rows x 2 columns]" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# Join the complaints with their embeddings in the same DataFrame\n", - "combined_df = downsampled_issues_df.join(predicted_embeddings)" + "combined_df = downsampled_issues_df.join(predicted_embeddings, how=\"left\")\n", + "combined_df" ] }, { @@ -470,12 +867,12 @@ "id": "OUZ3NNbzo1Tb" }, "source": [ - "## Step 2: KMeans clustering" + "## Step 2: Create k-means model and predict clusters" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 31, "metadata": { "id": "AhNTnEC5FRz2" }, @@ -496,14 +893,152 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 32, "metadata": { "id": "6poSxh-fGJF7" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 65eb317d-59f1-4d10-acd1-4b7f3778114c is DONE. 61.7 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 156e445e-cc01-4b30-84cc-ac1c98a69b81 is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 5befc212-f4a3-4e33-b1b2-01e809acdcbd is DONE. 61.9 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job bd271178-8b8d-45dc-ac57-7f0194d0daac is DONE. 80.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job bbfb9cca-622d-4bf5-9fc0-6d9a85287d41 is DONE. 80.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job a5f30b32-9fb0-42b4-b426-d8484f008bdb is DONE. 160.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "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", + "
CENTROID_ID
4222
6163
8335
13707
14303
\n", + "

5 rows × 1 columns

\n", + "
[5 rows x 1 columns in total]" + ], + "text/plain": [ + " CENTROID_ID\n", + "422 2\n", + "616 3\n", + "833 5\n", + "1370 7\n", + "1430 3\n", + "\n", + "[5 rows x 1 columns]" + ] + }, + "execution_count": 32, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# Use KMeans clustering to calculate our groups. Will take ~3 minutes.\n", - "cluster_model.fit(combined_df[[\"text_embedding\"]])\n", + "cluster_model.fit(combined_df[\"text_embedding\"])\n", "clustered_result = cluster_model.predict(combined_df[[\"text_embedding\"]])\n", "# Notice the CENTROID_ID column, which is the ID number of the group that\n", "# each complaint belongs to.\n", @@ -512,12 +1047,123 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 33, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 7a41196e-ea67-44ac-95a7-7dce620d6d21 is DONE. 320.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 8008b482-1a0d-461f-a215-4676d9d918dc is DONE. 72.4 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "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", + "
consumer_complaint_narrativetext_embeddingCENTROID_ID
2580664Hello, my name is XXXX XXXX, and I am writing ...[0.0003211698785889894, -0.01816680282354355, ...2
1806973This is XXXX XXXX and I am submitting this com...[-0.009485247544944286, -0.025846892967820168,...5
2055053XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX...[-0.010950954630970955, -0.0249345600605011, 0...3
2515231When I reinvestigated my credit report, I real...[-0.009660656563937664, -0.05793113633990288, ...5
2633049Checking my credit report XX/XX/2018 with all ...[-0.0022159104701131582, -0.03330004960298538,...3
\n", + "

5 rows × 3 columns

\n", + "
[5 rows x 3 columns in total]" + ], + "text/plain": [ + " consumer_complaint_narrative \\\n", + "2580664 Hello, my name is XXXX XXXX, and I am writing ... \n", + "1806973 This is XXXX XXXX and I am submitting this com... \n", + "2055053 XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX... \n", + "2515231 When I reinvestigated my credit report, I real... \n", + "2633049 Checking my credit report XX/XX/2018 with all ... \n", + "\n", + " text_embedding CENTROID_ID \n", + "2580664 [0.0003211698785889894, -0.01816680282354355, ... 2 \n", + "1806973 [-0.009485247544944286, -0.025846892967820168,... 5 \n", + "2055053 [-0.010950954630970955, -0.0249345600605011, 0... 3 \n", + "2515231 [-0.009660656563937664, -0.05793113633990288, ... 5 \n", + "2633049 [-0.0022159104701131582, -0.03330004960298538,... 3 \n", + "\n", + "[5 rows x 3 columns]" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# Join the group number to the complaints and their text embeddings\n", - "combined_clustered_result = combined_df.join(clustered_result)" + "combined_clustered_result = combined_df.join(clustered_result)\n", + "\n", + "combined_clustered_result.head(n=5)" ] }, { @@ -535,7 +1181,7 @@ "id": "21rNsFMHo8hO" }, "source": [ - "## Step 3: Summarize the complaints" + "## Step 3: Use PaLM2 LLM model to summarize complaint clusters" ] }, { @@ -548,11 +1194,36 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 34, "metadata": { "id": "2E7wXM_jGqo6" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 50c7c0dd-94a2-494e-a37f-6a838a518f6c is DONE. 11.0 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job d96c847f-c292-4804-bd05-fd643c41c7a5 is DONE. 11.0 MB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], "source": [ "# Using bigframes, with syntax identical to pandas,\n", "# filter out the first and second groups\n", @@ -569,11 +1240,100 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 36, "metadata": { "id": "ZNDiueI9IP5e" }, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "comment list 1:\n", + "1. XXXX is a peer to peer lending conmpany that uses borrowers crypto to collateralize loans from investors ( like myself ). I've been investing with them for almost XXXX years and currently have {$240000.00} tied up in lending products with XXXX. \n", + "As of XXXX days ago we received an email saying all business operations have been ceased and no withdrawals or deposits will be allowed. They said they'll update customers within 10 days, but no one can reach anyone at the company to find out any more details as they are not answering calls nor returning emails. It also appears the company has scrubbed its XXXX page and the XXXX pages of top executives. \n", + "\n", + "All collateral and client 's investment funds are supposedly held at or processed through XXXX XXXX XXXX ( registered SEC company ). XXXX XXXX keeps telling us to contact XXXX and won't give us any information, so we have no way to find out what's happening with our funds/collateral or if everything is gone. We have a XXXX channel up where people are gathering evidence, documentation, etc. This is probably the best place to start to get a broad view of what's happening. Details below. \n", + "\n", + "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX CONST LLC ( Business ID : XXXX ) FoXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX 'Cease of Operations ' email received by all investors XXXX XX/XX/2022 at XXXX : \" Dear XXXX Users, Given the collapses of several cryptocurrencies so far this year and the rapidly deteriorating market conditions that have been prompting heavy withdrawals across all XXXX lending and XXXX exchange platforms recently, we are sad to inform you that we are unable to continue to operate our business as usual. As such, we are limiting our business activities, including pausing user withdrawals as allowed under our Terms of XXXX. \n", + "No deposit or investment request will be processed at this time. \n", + "\n", + "Our team is working diligently towards our objective of maximizing value for all of our Users, and our top priority continues to be to protect your interests. As we explore all options available to us, we will provide updates to you as we go. \n", + "\n", + "We hope to communicate with you within the next XXXX business days on the next steps to address the situation. We appreciate your patience in this trying time. \n", + "\n", + "Sincerely yoursXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", + "2. Submitted XX/XX/XXXX\n", + "Typed XX/XX/XXXX:\n", + "\n", + "XX/XX/XXXX\n", + "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX, XXXX XXXX\n", + "PH:. XXXX\n", + "PH: XXXX\n", + "EM:\n", + "XXXX\n", + "XXXX\n", + "XXXX XXXX \n", + "XXXX XXXX\n", + "Date of Birth XX/XX/XXXX\n", + "SS#: XXXX\n", + "TO:\n", + "* Consumer Financial Protection Brueau\n", + "* Department of Veteran Affairs, Office of the Inspector General\n", + "My name is XXXX XXXX XXXX, I've received more than one email from Discover Card in my XXXX XXXX, past emails from Discover Card were unautherized deletions.\n", + "From: Discover Card XXXX\n", + "To: You XXXX\n", + "Date: XX/XX/XXXX, XXXX XXXX XXXX From: Discover Card XXXX>\n", + "To Recipient \n", + "Date Mon, XX/XX/XXXX XXXX XXXX\n", + "I dont and havent ever had a Discover Checking, Savings, Business Accounts nor Loans of any kind through any Bank called Discover. The 1st time I was contacted by Discover Card I resided alone from XX/XX/XXXX to XX/XX/XXXXat XXXX XXXX XXXX at XXXX XXXX XXXX XXXX XXXX in XXXX, XXXX years prior to me moving here to XXXX, XXXX in XX/XX/XXXX. When \n", + "\n", + "\n", + "Discover Card had 1st contacted me in XXXX, XXXX it was associated with my XXXX XXXX XXXX website related online Merchants Account. Not once have I ever applied for or had any Website Merchant Accounts here in XXXX; I only applied for online online Merchant Accounts associated with my XXXX related Accounts I purchased while residing in XXXX, XXXX. Some of my website related information was stolen both in XXXX, XXXX and here in XXXX along with my other property that hasn't been returned to me. I don't and haven't ever had any XXXX XXXX related Agreements,Contracts or Credit Cards offered to Veterans associated with ones businesses. Nor have I ever applied for or had a Business License or Business Permit in any City or State inspite of my diverse interest. Not once have I ever allowed another be it an Paralegal, Payee, Attorney, Employers, Landlords, Veteran Organizations including Vocational Rehabilitation Programs, XXXX( XXXX XXXX XXXX, XXXX XXXX, Entertainment Companies, Banks, Celebrity Personal Assistant Agencies or Celebs, Shelters, Charities, HUD, Housing Arthority, Department of Veteran Affairs, Military, Law Enforcement or anyone else nor their employess to sign any business related Agreements or Contracts on my behalf; not even my family members or friends. \n", + "None of my XXXX XXXX attempts were associated with my Employers, Department of Veteran Affairs,Vocational Rehabilitation Programs Military, Landlords, HUD( Housing Authority),Friends, Family nor did I ever sign related Agreements or Contracts with them. Not once had I ever provided anyone the passwords to be able to sign into my accounts rather were aware of my accounts or not. Yes, my desktop computer that was stolen along with my other property XX/XX/XXXX was registered with my Online Merchant Account. I had paid for my Merchant related Accounts through my same XXXX XXXX XXXX Account I purchased both of my XXXX XXXX XXXX related accounts through. That was 1st once during the Summer of XX/XX/XXXX and 2nd my related website months later, while I resided in XXXX XXXX and I worked for XXXX. I never offered nor did I ever sign any business Contracts or Agreements with XXXX nor my Landlord or their staff associted with any of my online websites or Merchant Accounts. My XXXX XXXX XXXX Compensation was deposited into both of my XXXX XXXX XXXX Accounts at that time. My account was changed during the Summer of XX/XX/XXXXbecause of theft of my Bank Card. None of my Checking,Savings, past Credit Cards or Business related were shared accounts in which others were allowed to \n", + "use to make purchases. I had written checks from my XXXX XXXX XXXX account to pay for my XXXX XXXX XXXX XXXX on the XXXX XXXX here in XXXX in XX/XX/XXXX before it's name changed to XXXX XXXX. Prior to me using my same account open a Checking account in person at XXXX XXXX before it's name was changed to XXXX XXXX. Where my XXXX XXXX XXXX XXXX has been deposited since that time. I had used my XXXX XXXX Checking to pay for my XXXX XXXX XXXX XXXX both before theft of my property XX/XX/XXXX and that was also prior to the theft of my property from my XXXX XXXX XXXX XXXX in XX/XX/XXXX.\n", + "I've stated this many times:\n", + "I paid for my 1st XXXX XXXX XXXX Membership while employed at XXXX using my XXXX XXXX XXXX account XXXX my XXXX XXXX XXXX XXXX was also deposited. That was changed to XXXX because I didn't receive my 1st XXXX XXXX XXXX Card the bank sent to XXXX XXXX residence on XXXX XXXX in XX/XX/XXXX while I was there. In which both my XXXX salary and XXXX XXXX XXXX XXXX were deposited into my account, no money from XXXX XXXX nor anyone else that was at that residence was given to nor were any of my children there. Nor did XXXX or any other person at that residence ever give me my missing Bank Card not even after I moved out and stayed a month at XXXX XXXX XXXX using my replacement card to pay for my Hotel room. Which is the same account I used to pay for XXXX XXXX Membership, XXXX XXXX XXXX, XXXX XXXX Membership fees, and various online Merchant Account activation related fees.\n", + "* XXXX XXXX XXXX.\n", + "XXXX XXXX XXXX XXXX. Membership\n", + "\n", + "# XXXX\n", + "* XXXX XXXX Membership\n", + "# XXXX\n", + "* Total Merchant Services XXXX and XXXX.\n", + "* XXXX XXXX XXXX XXXX XXXX\n", + "* XXXX XXXX changed my $XXXX a month fees to my XXXX XXXX XXXX account #XXXX.\n", + "XX/XX/XXXX - XX/XX/XXXX XXXX XXXX, XXXX.\n", + "\n", + "Rep: XXXX XXXX XXXX, Fl \n", + "XXXX\n", + "XXXX Website \n", + "XXXX\n", + "Software and website owner, I performed Internet advertising and marketing, to promote this software and website. I worked and XXXX from my home XXXX XXXX XXXX XXXX XXXX , XXXX. I purchased XXXX XXXX XXXX-Software Electronic Book CD and was given a website to promote the software on the internet. The XXXX was given a copy of my website owner certificate document submitted to me when I purchased the software marketing program as well copies of my other school transcripts in addition to XXXX XXXX XXXX for example. XXXX, represented the first initials of my children's names. I wasn't ever paid and I'm still owed the money. Nor did my marketing program have anything to do with any schools, college nor university programs nor did I ever offer or sign any agreement to include it such. Nor did my XXXX XXXX XXXX have anything to do with any other employers, Department of Family and Children, Military, Veteran Organizations or Food Stamp programs, Section 8 nor Indianapolis Housing Authority for example; only me.\n", + "Thank you,\n", + "XXXX XXXX\n", + "3. ACCORDING TO 15 U.S. CODE 6803-DISCLOSURE OF INSTITUTION PRIVACY POLICY, AND ACCORDING TO U.S. CODE 6802- OBLIGATIONS WITH RESPECT TO DISCLOSURES OF PERSONAL INFORMATION. ( b ) OPT OUT ( 1 ) IN GENERAL A FINANCIAL INSTITUTION MAY NOT DISCLOSE NONPUBLIC PERSONAL INFORMATION TO A NONAFFILIATED THIRD PARTY ( TRANSUNION, XXXX, AND XXXX. ) UNLESS- ( A ) SUCH FINANCIAL INSTITUTION CLEARLY AND CONSPICUOUSLY DISCLOSES TO THE CONSUMER, IN WRITING OR IN ELECTRONIC FORM OR OTHER FORM PERMITTED BY THE REGULATIONS PRESCRIBED UNDER SECTION 6804 OF THIS TITLE. ALSO ACCORDING TO THE \" XXXX ACT '', FINANCIAL INSTITUTIONS MUST TELL THEIR CUSTOMERS ABOUT THEIR INFORMATION-SHARING PRACTICES AND EXPLAIN TO CUSTOMERS THEIR RIGHT TO \" OPT OUT '' IF THEY DON'T WANT THEIR INFORMATION SHARED WITH CERTAIN THIRD PARTIES. UNDER THE FDCPA, A COLLECTOR MUST PROVIDE YOU WITH INFORMATION ABOUT THE DEBT IN ITS INITIAL COMMUNICATION OR WITHIN FIVE DAYS AFTER THE INITIAL COMMUNICATION. ALSO, THE FDCPA STATES, \" YOU CAN NOT ATTEMPT TO COLLECT AN DEBT WHILE A PERSON ( THE CONSUMER ) SUPRESS VALIDATION. TRANSUNION, XXXX, XXXX, AND THE ACCOUNTS LISTED BELOW HAVE CLEARLY VIOLATED MY RIGHTS : XXXX ACCOUNT # XXXX, XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX ACCOUNT # XXXX, XXXX XXXX XXXX XXXX ACCOUNT # XXXX, AND XXXX ACCOUNT # XXXX. FAILURE TO RESPOND SATISFACTORILY WITH DELETIONS OF ALL THE ABOVE ACCOUNTS WILL RESULT IN LEGAL ACTIONS BEING TAKEN AGAINST, TRANSUNION, XXXX, XXXX, WHICH I'LL BE SEEKING A {$1000.00} PER VIOLATION FOR DEFAMATION OF CHARACTER ( PER SE ) NEGLIGENT ENABLEMENT OF IDENTITY FRAUD. 15 USC 1681 VIOLATIONS FOR WILLFUL NONCOMPLIANCE-616 CIVIL LIABILITY FOR WILLFUL NONCOPLIANCE. THIS IS THE THIRD TIME I'VE SUBMITTED A COMPLAINT, AND THE REPONSE I GET IS \" YOU CAN NOT LOCATE MY CREDIT REPORT! '' THIS IS CLEARLY NEGLIGENCE.\n", + "4. I do not know how this works, but I need it done or somehow corrected. My name is XXXX XXXX, XXXX XXXX XXXX XXXX TN XXXXMy SS XXXX DOB XXXX. I had some issues with my income being affected by the COVID-19PANDEMICSHUTDOWN. I was under the 1 CARESAct, Pub. L. 116-136, section 4021, codified at FCRAsection 623 ( a ) ( 1 ) ( F ) ( i ) ( I ), 15 U.S.C.1681s- 2 ( a ) ( 1 ) ( F ) ( i ) ( I ). I am requesting some accommodations so I care to protect the integrity of my credit file. US DEPT OF ED / XXXX # XXXX, # XXXX accounts are reporting on XXXX, XXXX The was 30,60, 90 DAYS LATEsince requested assistance due to the pandemic. I found a few accounts that I have never done any business with these companies and the accounts do not belong on my report : XXXX XXXX # XXXX, XXXX XXXX XXXX XXXX # XXXX. \n", + "\n", + "I have some issues with the misspelling of my name, my correct spelling is XXXX XXXX. Please remove any other variation of my name they are not correct. The following addresses do not belong to me please delete them : XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXXSC, XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", + "5. I want to know if this is even legal?! How can they disclose information without knowing its a correct email?!\n", + "\n", + "comment list 2:\n", + "1. Hello, my name is XXXX XXXX, and I am writing to delete the following information in my file. The items I need deleted are listed in the report. I am a victim of identity theft and did not make the charge. I ask that the items be deleted to correct my credit report. I reported the theft of my identity to the Federal Trade Commission and I also have enclosed copies of the Federal Trade Commissions Identity Theft Affidavit. Please delete the items as soon as possible. The accounts are being reported currently open and the accounts need to be closed. \n", + "XXXX account number XXXX opened on XX/XX/2022 for the amount {$530.00} XXXX XXXX XXXX account number XXXX opened on XX/XX/2022 for the amount of {$140.00} The accounts are being reported currently open and need to be closed immediately. \n", + "Based on, 15 U.S. Code 1681c2 a consumer reporting agency shall block the reporting of any information in the file of a consumer that the consumer identifies as information that resulted from an alleged identity theft, not later than 4 business days after the date of receipt. This account should not be furnished on my consumer report. As a consumer I am demanding the deletion of the accounts listed IMMEDIATELY.\n", + "2. To whom it may concern : My personal information was breach in the internet as result accounts had been open in my name, I was advise to fill out an Id theft report to help me deal with this situation, I have listed each one of the accounts that do not belong to me. This is my second request to remove unverified items in my report, but XXXX keep rposting these account with out providing any type of original document as the FCRA provide, you need to provide me with original documents or remove these account immediately.\n", + "3. Ive been Disputting my XXXX XXXX I opened this account and someone got my information and used my card, I contacted XXXX over and over, they removed the negative reporting from my XXXX report but still reporting it negative on my XXXX and Expean this is very unfair to me because Im a victim of identity theft\n", + "4. Today, XX/XX/2021, I received three items in the mail, one envelope containing an unsolicited debit card from Navy Federal credit Union and the other two, with a letter each describing The Important Rights on two accounts should these accounts become delinquent under New York law. \n", + "\n", + "First of all, I never applied for these accounts with Navy Federal, not have I authorized anyone to do so on my behalf. I immediately contacted Navy Federal via phone and was told I was most likely a victim of identity theft and that I should monitor my credit and use a credit monitoring service. I was also asked for my email and mailing information in order to receive a letter from them regarding this issue. \n", + "\n", + "My main concern is having someone using my identity to illegally open bank accounts and commit fraud, destroying my credit and finances in the process. This bank is in another state from where I reside. I have not lived in Virginia nor do I intend to do so in the foreseeable future.\n", + "5. My personal information ( including my SSN, Drivers License Info, Addresses, and more ) was stolen from a hacking, and Equifax did n't tell the public about the hack until more than a month after the hacking. During this time, three Equifax executives were caught inside trading. It really shows how Equifax cares about other people!\n", + "\n" + ] + } + ], "source": [ "# Build plain-text prompts to send to PaLM 2. Use only 5 complaints from each group.\n", "prompt1 = 'comment list 1:\\n'\n", @@ -592,11 +1352,100 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 37, "metadata": { "id": "BfHGJLirzSvH" }, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Please highlight the most obvious difference betweenthe two lists of comments:\n", + "comment list 1:\n", + "1. XXXX is a peer to peer lending conmpany that uses borrowers crypto to collateralize loans from investors ( like myself ). I've been investing with them for almost XXXX years and currently have {$240000.00} tied up in lending products with XXXX. \n", + "As of XXXX days ago we received an email saying all business operations have been ceased and no withdrawals or deposits will be allowed. They said they'll update customers within 10 days, but no one can reach anyone at the company to find out any more details as they are not answering calls nor returning emails. It also appears the company has scrubbed its XXXX page and the XXXX pages of top executives. \n", + "\n", + "All collateral and client 's investment funds are supposedly held at or processed through XXXX XXXX XXXX ( registered SEC company ). XXXX XXXX keeps telling us to contact XXXX and won't give us any information, so we have no way to find out what's happening with our funds/collateral or if everything is gone. We have a XXXX channel up where people are gathering evidence, documentation, etc. This is probably the best place to start to get a broad view of what's happening. Details below. \n", + "\n", + "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX CONST LLC ( Business ID : XXXX ) FoXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX 'Cease of Operations ' email received by all investors XXXX XX/XX/2022 at XXXX : \" Dear XXXX Users, Given the collapses of several cryptocurrencies so far this year and the rapidly deteriorating market conditions that have been prompting heavy withdrawals across all XXXX lending and XXXX exchange platforms recently, we are sad to inform you that we are unable to continue to operate our business as usual. As such, we are limiting our business activities, including pausing user withdrawals as allowed under our Terms of XXXX. \n", + "No deposit or investment request will be processed at this time. \n", + "\n", + "Our team is working diligently towards our objective of maximizing value for all of our Users, and our top priority continues to be to protect your interests. As we explore all options available to us, we will provide updates to you as we go. \n", + "\n", + "We hope to communicate with you within the next XXXX business days on the next steps to address the situation. We appreciate your patience in this trying time. \n", + "\n", + "Sincerely yoursXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", + "2. Submitted XX/XX/XXXX\n", + "Typed XX/XX/XXXX:\n", + "\n", + "XX/XX/XXXX\n", + "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX, XXXX XXXX\n", + "PH:. XXXX\n", + "PH: XXXX\n", + "EM:\n", + "XXXX\n", + "XXXX\n", + "XXXX XXXX \n", + "XXXX XXXX\n", + "Date of Birth XX/XX/XXXX\n", + "SS#: XXXX\n", + "TO:\n", + "* Consumer Financial Protection Brueau\n", + "* Department of Veteran Affairs, Office of the Inspector General\n", + "My name is XXXX XXXX XXXX, I've received more than one email from Discover Card in my XXXX XXXX, past emails from Discover Card were unautherized deletions.\n", + "From: Discover Card XXXX\n", + "To: You XXXX\n", + "Date: XX/XX/XXXX, XXXX XXXX XXXX From: Discover Card XXXX>\n", + "To Recipient \n", + "Date Mon, XX/XX/XXXX XXXX XXXX\n", + "I dont and havent ever had a Discover Checking, Savings, Business Accounts nor Loans of any kind through any Bank called Discover. The 1st time I was contacted by Discover Card I resided alone from XX/XX/XXXX to XX/XX/XXXXat XXXX XXXX XXXX at XXXX XXXX XXXX XXXX XXXX in XXXX, XXXX years prior to me moving here to XXXX, XXXX in XX/XX/XXXX. When \n", + "\n", + "\n", + "Discover Card had 1st contacted me in XXXX, XXXX it was associated with my XXXX XXXX XXXX website related online Merchants Account. Not once have I ever applied for or had any Website Merchant Accounts here in XXXX; I only applied for online online Merchant Accounts associated with my XXXX related Accounts I purchased while residing in XXXX, XXXX. Some of my website related information was stolen both in XXXX, XXXX and here in XXXX along with my other property that hasn't been returned to me. I don't and haven't ever had any XXXX XXXX related Agreements,Contracts or Credit Cards offered to Veterans associated with ones businesses. Nor have I ever applied for or had a Business License or Business Permit in any City or State inspite of my diverse interest. Not once have I ever allowed another be it an Paralegal, Payee, Attorney, Employers, Landlords, Veteran Organizations including Vocational Rehabilitation Programs, XXXX( XXXX XXXX XXXX, XXXX XXXX, Entertainment Companies, Banks, Celebrity Personal Assistant Agencies or Celebs, Shelters, Charities, HUD, Housing Arthority, Department of Veteran Affairs, Military, Law Enforcement or anyone else nor their employess to sign any business related Agreements or Contracts on my behalf; not even my family members or friends. \n", + "None of my XXXX XXXX attempts were associated with my Employers, Department of Veteran Affairs,Vocational Rehabilitation Programs Military, Landlords, HUD( Housing Authority),Friends, Family nor did I ever sign related Agreements or Contracts with them. Not once had I ever provided anyone the passwords to be able to sign into my accounts rather were aware of my accounts or not. Yes, my desktop computer that was stolen along with my other property XX/XX/XXXX was registered with my Online Merchant Account. I had paid for my Merchant related Accounts through my same XXXX XXXX XXXX Account I purchased both of my XXXX XXXX XXXX related accounts through. That was 1st once during the Summer of XX/XX/XXXX and 2nd my related website months later, while I resided in XXXX XXXX and I worked for XXXX. I never offered nor did I ever sign any business Contracts or Agreements with XXXX nor my Landlord or their staff associted with any of my online websites or Merchant Accounts. My XXXX XXXX XXXX Compensation was deposited into both of my XXXX XXXX XXXX Accounts at that time. My account was changed during the Summer of XX/XX/XXXXbecause of theft of my Bank Card. None of my Checking,Savings, past Credit Cards or Business related were shared accounts in which others were allowed to \n", + "use to make purchases. I had written checks from my XXXX XXXX XXXX account to pay for my XXXX XXXX XXXX XXXX on the XXXX XXXX here in XXXX in XX/XX/XXXX before it's name changed to XXXX XXXX. Prior to me using my same account open a Checking account in person at XXXX XXXX before it's name was changed to XXXX XXXX. Where my XXXX XXXX XXXX XXXX has been deposited since that time. I had used my XXXX XXXX Checking to pay for my XXXX XXXX XXXX XXXX both before theft of my property XX/XX/XXXX and that was also prior to the theft of my property from my XXXX XXXX XXXX XXXX in XX/XX/XXXX.\n", + "I've stated this many times:\n", + "I paid for my 1st XXXX XXXX XXXX Membership while employed at XXXX using my XXXX XXXX XXXX account XXXX my XXXX XXXX XXXX XXXX was also deposited. That was changed to XXXX because I didn't receive my 1st XXXX XXXX XXXX Card the bank sent to XXXX XXXX residence on XXXX XXXX in XX/XX/XXXX while I was there. In which both my XXXX salary and XXXX XXXX XXXX XXXX were deposited into my account, no money from XXXX XXXX nor anyone else that was at that residence was given to nor were any of my children there. Nor did XXXX or any other person at that residence ever give me my missing Bank Card not even after I moved out and stayed a month at XXXX XXXX XXXX using my replacement card to pay for my Hotel room. Which is the same account I used to pay for XXXX XXXX Membership, XXXX XXXX XXXX, XXXX XXXX Membership fees, and various online Merchant Account activation related fees.\n", + "* XXXX XXXX XXXX.\n", + "XXXX XXXX XXXX XXXX. Membership\n", + "\n", + "# XXXX\n", + "* XXXX XXXX Membership\n", + "# XXXX\n", + "* Total Merchant Services XXXX and XXXX.\n", + "* XXXX XXXX XXXX XXXX XXXX\n", + "* XXXX XXXX changed my $XXXX a month fees to my XXXX XXXX XXXX account #XXXX.\n", + "XX/XX/XXXX - XX/XX/XXXX XXXX XXXX, XXXX.\n", + "\n", + "Rep: XXXX XXXX XXXX, Fl \n", + "XXXX\n", + "XXXX Website \n", + "XXXX\n", + "Software and website owner, I performed Internet advertising and marketing, to promote this software and website. I worked and XXXX from my home XXXX XXXX XXXX XXXX XXXX , XXXX. I purchased XXXX XXXX XXXX-Software Electronic Book CD and was given a website to promote the software on the internet. The XXXX was given a copy of my website owner certificate document submitted to me when I purchased the software marketing program as well copies of my other school transcripts in addition to XXXX XXXX XXXX for example. XXXX, represented the first initials of my children's names. I wasn't ever paid and I'm still owed the money. Nor did my marketing program have anything to do with any schools, college nor university programs nor did I ever offer or sign any agreement to include it such. Nor did my XXXX XXXX XXXX have anything to do with any other employers, Department of Family and Children, Military, Veteran Organizations or Food Stamp programs, Section 8 nor Indianapolis Housing Authority for example; only me.\n", + "Thank you,\n", + "XXXX XXXX\n", + "3. ACCORDING TO 15 U.S. CODE 6803-DISCLOSURE OF INSTITUTION PRIVACY POLICY, AND ACCORDING TO U.S. CODE 6802- OBLIGATIONS WITH RESPECT TO DISCLOSURES OF PERSONAL INFORMATION. ( b ) OPT OUT ( 1 ) IN GENERAL A FINANCIAL INSTITUTION MAY NOT DISCLOSE NONPUBLIC PERSONAL INFORMATION TO A NONAFFILIATED THIRD PARTY ( TRANSUNION, XXXX, AND XXXX. ) UNLESS- ( A ) SUCH FINANCIAL INSTITUTION CLEARLY AND CONSPICUOUSLY DISCLOSES TO THE CONSUMER, IN WRITING OR IN ELECTRONIC FORM OR OTHER FORM PERMITTED BY THE REGULATIONS PRESCRIBED UNDER SECTION 6804 OF THIS TITLE. ALSO ACCORDING TO THE \" XXXX ACT '', FINANCIAL INSTITUTIONS MUST TELL THEIR CUSTOMERS ABOUT THEIR INFORMATION-SHARING PRACTICES AND EXPLAIN TO CUSTOMERS THEIR RIGHT TO \" OPT OUT '' IF THEY DON'T WANT THEIR INFORMATION SHARED WITH CERTAIN THIRD PARTIES. UNDER THE FDCPA, A COLLECTOR MUST PROVIDE YOU WITH INFORMATION ABOUT THE DEBT IN ITS INITIAL COMMUNICATION OR WITHIN FIVE DAYS AFTER THE INITIAL COMMUNICATION. ALSO, THE FDCPA STATES, \" YOU CAN NOT ATTEMPT TO COLLECT AN DEBT WHILE A PERSON ( THE CONSUMER ) SUPRESS VALIDATION. TRANSUNION, XXXX, XXXX, AND THE ACCOUNTS LISTED BELOW HAVE CLEARLY VIOLATED MY RIGHTS : XXXX ACCOUNT # XXXX, XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX ACCOUNT # XXXX, XXXX XXXX XXXX XXXX ACCOUNT # XXXX, AND XXXX ACCOUNT # XXXX. FAILURE TO RESPOND SATISFACTORILY WITH DELETIONS OF ALL THE ABOVE ACCOUNTS WILL RESULT IN LEGAL ACTIONS BEING TAKEN AGAINST, TRANSUNION, XXXX, XXXX, WHICH I'LL BE SEEKING A {$1000.00} PER VIOLATION FOR DEFAMATION OF CHARACTER ( PER SE ) NEGLIGENT ENABLEMENT OF IDENTITY FRAUD. 15 USC 1681 VIOLATIONS FOR WILLFUL NONCOMPLIANCE-616 CIVIL LIABILITY FOR WILLFUL NONCOPLIANCE. THIS IS THE THIRD TIME I'VE SUBMITTED A COMPLAINT, AND THE REPONSE I GET IS \" YOU CAN NOT LOCATE MY CREDIT REPORT! '' THIS IS CLEARLY NEGLIGENCE.\n", + "4. I do not know how this works, but I need it done or somehow corrected. My name is XXXX XXXX, XXXX XXXX XXXX XXXX TN XXXXMy SS XXXX DOB XXXX. I had some issues with my income being affected by the COVID-19PANDEMICSHUTDOWN. I was under the 1 CARESAct, Pub. L. 116-136, section 4021, codified at FCRAsection 623 ( a ) ( 1 ) ( F ) ( i ) ( I ), 15 U.S.C.1681s- 2 ( a ) ( 1 ) ( F ) ( i ) ( I ). I am requesting some accommodations so I care to protect the integrity of my credit file. US DEPT OF ED / XXXX # XXXX, # XXXX accounts are reporting on XXXX, XXXX The was 30,60, 90 DAYS LATEsince requested assistance due to the pandemic. I found a few accounts that I have never done any business with these companies and the accounts do not belong on my report : XXXX XXXX # XXXX, XXXX XXXX XXXX XXXX # XXXX. \n", + "\n", + "I have some issues with the misspelling of my name, my correct spelling is XXXX XXXX. Please remove any other variation of my name they are not correct. The following addresses do not belong to me please delete them : XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXXSC, XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", + "5. I want to know if this is even legal?! How can they disclose information without knowing its a correct email?!\n", + "comment list 2:\n", + "1. Hello, my name is XXXX XXXX, and I am writing to delete the following information in my file. The items I need deleted are listed in the report. I am a victim of identity theft and did not make the charge. I ask that the items be deleted to correct my credit report. I reported the theft of my identity to the Federal Trade Commission and I also have enclosed copies of the Federal Trade Commissions Identity Theft Affidavit. Please delete the items as soon as possible. The accounts are being reported currently open and the accounts need to be closed. \n", + "XXXX account number XXXX opened on XX/XX/2022 for the amount {$530.00} XXXX XXXX XXXX account number XXXX opened on XX/XX/2022 for the amount of {$140.00} The accounts are being reported currently open and need to be closed immediately. \n", + "Based on, 15 U.S. Code 1681c2 a consumer reporting agency shall block the reporting of any information in the file of a consumer that the consumer identifies as information that resulted from an alleged identity theft, not later than 4 business days after the date of receipt. This account should not be furnished on my consumer report. As a consumer I am demanding the deletion of the accounts listed IMMEDIATELY.\n", + "2. To whom it may concern : My personal information was breach in the internet as result accounts had been open in my name, I was advise to fill out an Id theft report to help me deal with this situation, I have listed each one of the accounts that do not belong to me. This is my second request to remove unverified items in my report, but XXXX keep rposting these account with out providing any type of original document as the FCRA provide, you need to provide me with original documents or remove these account immediately.\n", + "3. Ive been Disputting my XXXX XXXX I opened this account and someone got my information and used my card, I contacted XXXX over and over, they removed the negative reporting from my XXXX report but still reporting it negative on my XXXX and Expean this is very unfair to me because Im a victim of identity theft\n", + "4. Today, XX/XX/2021, I received three items in the mail, one envelope containing an unsolicited debit card from Navy Federal credit Union and the other two, with a letter each describing The Important Rights on two accounts should these accounts become delinquent under New York law. \n", + "\n", + "First of all, I never applied for these accounts with Navy Federal, not have I authorized anyone to do so on my behalf. I immediately contacted Navy Federal via phone and was told I was most likely a victim of identity theft and that I should monitor my credit and use a credit monitoring service. I was also asked for my email and mailing information in order to receive a letter from them regarding this issue. \n", + "\n", + "My main concern is having someone using my identity to illegally open bank accounts and commit fraud, destroying my credit and finances in the process. This bank is in another state from where I reside. I have not lived in Virginia nor do I intend to do so in the foreseeable future.\n", + "5. My personal information ( including my SSN, Drivers License Info, Addresses, and more ) was stolen from a hacking, and Equifax did n't tell the public about the hack until more than a month after the hacking. During this time, three Equifax executives were caught inside trading. It really shows how Equifax cares about other people!\n", + "\n" + ] + } + ], "source": [ "# The plain English request we will make of PaLM 2\n", "prompt = (\n", @@ -616,22 +1465,42 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 38, "metadata": { "id": "mL5P0_3X04dE" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 66e3af22-91cb-400a-92c3-69e7cd12ee01 is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], "source": [ "from bigframes.ml.llm import PaLM2TextGenerator\n", "\n", +<<<<<<< HEAD + "q_a_model = PaLM2TextGenerator()" +======= + "# Create a BigQuery Cloud resource connection\n", + "CONN_NAME = \"bqdf-llm\"\n", "session = bf.get_global_session()\n", + "\n", "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", "q_a_model = PaLM2TextGenerator(session=session, connection_name=connection)" +>>>>>>> origin/lmm-kmeans-notebook ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 39, "metadata": { "id": "ICWHsqAW1FNk" }, @@ -643,11 +1512,58 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 40, "metadata": { "id": "gB7e1LXU1pst" }, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job 653add17-29be-408c-8882-064217f8556e is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 8fd16954-853a-45fd-80bc-65b1242429e2 is DONE. 8 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job d9929bcb-26ce-4844-b68e-f4a980b90ede is DONE. 171 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "' The first comment list is about people complaining about companies or services, while the second comment list is about people reporting identity theft or fraud.'" + ] + }, + "execution_count": 40, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# Send the request for PaLM 2 to generate a response to our prompt\n", "major_difference = q_a_model.predict(df)\n", @@ -662,6 +1578,21 @@ "source": [ "We now see PaLM2TextGenerator's characterization of the different comment groups. Thanks for using BigQuery DataFrames!" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Summary and next steps\n", + "\n", +<<<<<<< HEAD + "You've used the ML and LLM capabilities of BigQuery DataFrames to help analyze and understand a large dataset of unstructured feedback.\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", +>>>>>>> origin/lmm-kmeans-notebook + "\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": { @@ -682,7 +1613,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.10.13" } }, "nbformat": 4, From 416171a70d91d4a6b71622ba72685147ab7d6186 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Thu, 16 Nov 2023 11:04:18 -0800 Subject: [PATCH 166/197] feat!: model.predict returns all the columns (#204) 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/ml/cluster.py | 4 +- bigframes/ml/decomposition.py | 9 +- bigframes/ml/ensemble.py | 49 +- bigframes/ml/forecasting.py | 9 +- bigframes/ml/imported.py | 22 +- bigframes/ml/linear_model.py | 24 +- bigframes/ml/llm.py | 17 +- .../getting_started/ml_fundamentals.ipynb | 3586 ++++++++++------- .../sklearn_linear_regression.ipynb | 1192 +++--- tests/system/large/ml/test_cluster.py | 4 +- tests/system/large/ml/test_ensemble.py | 2 +- tests/system/large/ml/test_pipeline.py | 4 +- tests/system/small/ml/test_cluster.py | 4 +- tests/system/small/ml/test_ensemble.py | 20 +- tests/system/small/ml/test_forecasting.py | 4 +- tests/system/small/ml/test_imported.py | 8 +- tests/system/small/ml/test_linear_model.py | 8 +- tests/system/small/ml/test_llm.py | 22 +- .../sklearn/cluster/_kmeans.py | 16 +- .../sklearn/linear_model/_base.py | 6 +- .../bigframes_vendored/xgboost/sklearn.py | 2 +- 21 files changed, 2737 insertions(+), 2275 deletions(-) diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index 772b90f666e..c9f52ba0b6b 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import cast, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union from google.cloud import bigquery @@ -92,7 +92,7 @@ def predict( (X,) = utils.convert_to_dataframe(X) - return cast(bpd.DataFrame, self._bqml_model.predict(X)[["CENTROID_ID"]]) + return self._bqml_model.predict(X) def to_gbq(self, model_name: str, replace: bool = False) -> KMeans: """Save the model to BigQuery. diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 8e6be6d28c3..7cda7a69933 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import cast, List, Optional, Union +from typing import List, Optional, Union from google.cloud import bigquery @@ -106,12 +106,7 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: (X,) = utils.convert_to_dataframe(X) - return cast( - bpd.DataFrame, - self._bqml_model.predict(X)[ - ["principal_component_" + str(i + 1) for i in range(self.n_components)] - ], - ) + return self._bqml_model.predict(X) def to_gbq(self, model_name: str, replace: bool = False) -> PCA: """Save the model to BigQuery. diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index 19ca8608ffb..fcb3fe53438 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import cast, Dict, List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from google.cloud import bigquery @@ -168,16 +168,7 @@ def predict( raise RuntimeError("A model must be fitted before predict") (X,) = utils.convert_to_dataframe(X) - df = self._bqml_model.predict(X) - return cast( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, @@ -328,19 +319,9 @@ def _fit( 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( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, @@ -486,19 +467,9 @@ def predict( ) -> 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( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, @@ -661,19 +632,9 @@ def predict( ) -> 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( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index 8e309d5e736..cf23854fa01 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import cast, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union from google.cloud import bigquery @@ -24,8 +24,6 @@ from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd -_PREDICT_OUTPUT_COLUMNS = ["forecast_timestamp", "forecast_value"] - class ARIMAPlus(base.SupervisedTrainablePredictor): """Time Series ARIMA Plus model.""" @@ -100,10 +98,7 @@ def predict(self, X=None) -> bpd.DataFrame: if not self._bqml_model: raise RuntimeError("A model must be fitted before predict") - return cast( - bpd.DataFrame, - self._bqml_model.forecast()[_PREDICT_OUTPUT_COLUMNS], - ) + return self._bqml_model.forecast() def score( self, diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py index fb8aa98befd..f6afc9aa382 100644 --- a/bigframes/ml/imported.py +++ b/bigframes/ml/imported.py @@ -78,16 +78,7 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: (X,) = utils.convert_to_dataframe(X) - df = self._bqml_model.predict(X) - return cast( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def to_gbq(self, model_name: str, replace: bool = False) -> TensorFlowModel: """Save the model to BigQuery. @@ -161,16 +152,7 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: (X,) = utils.convert_to_dataframe(X) - df = self._bqml_model.predict(X) - return cast( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def to_gbq(self, model_name: str, replace: bool = False) -> ONNXModel: """Save the model to BigQuery. diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index f11879500b9..433d9fbc380 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, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from google.cloud import bigquery @@ -145,16 +145,7 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: (X,) = utils.convert_to_dataframe(X) - df = self._bqml_model.predict(X) - return cast( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, @@ -267,16 +258,7 @@ def predict( (X,) = utils.convert_to_dataframe(X) - df = self._bqml_model.predict(X) - return cast( - bpd.DataFrame, - df[ - [ - cast(str, field.name) - for field in self._bqml_model.model.label_columns - ] - ], - ) + return self._bqml_model.predict(X) def score( self, diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 3cfc28e61f8..93e2ba825f4 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -149,7 +149,8 @@ def predict( Returns: - bigframes.dataframe.DataFrame: Output DataFrame with only 1 column as the output text results.""" + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. + """ # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models if temperature < 0.0 or temperature > 1.0: @@ -181,11 +182,7 @@ def predict( "top_p": top_p, "flatten_json_output": True, } - df = self._bqml_model.generate_text(X, options) - return cast( - bpd.DataFrame, - df[[_TEXT_GENERATE_RESULT_COLUMN]], - ) + return self._bqml_model.generate_text(X, options) class PaLM2TextEmbeddingGenerator(base.Predictor): @@ -269,7 +266,7 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: 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: - bigframes.dataframe.DataFrame: Output DataFrame with only 1 column as the output embedding results + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. """ # Params reference: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models @@ -287,8 +284,4 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: options = { "flatten_json_output": True, } - df = self._bqml_model.generate_text_embedding(X, options) - return cast( - bpd.DataFrame, - df[[_EMBED_TEXT_RESULT_COLUMN]], - ) + return self._bqml_model.generate_text_embedding(X, options) diff --git a/notebooks/getting_started/ml_fundamentals.ipynb b/notebooks/getting_started/ml_fundamentals.ipynb index 2f566dd7049..165bd90f314 100644 --- a/notebooks/getting_started/ml_fundamentals.ipynb +++ b/notebooks/getting_started/ml_fundamentals.ipynb @@ -14,46 +14,16 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 1, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "0c8a8bc0b4d64448aef68d6a98fae666", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job 28e903c6-e874-4b99-8f53-0755e0b0c188 is RUNNING. " ] }, "metadata": {}, @@ -61,13 +31,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "9680fd748e0546b4a010fda0155c5027", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e8aba858-7660-4274-8d90-8d2b0382f8f6 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 7950d6a7-3747-4454-bba2-9660e830647f is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -117,250 +85,250 @@ " \n", " 0\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 36.6\n", - " 18.4\n", - " 184.0\n", - " 3475.0\n", - " FEMALE\n", + " Biscoe\n", + " 40.1\n", + " 18.9\n", + " 188.0\n", + " 4300.0\n", + " MALE\n", " \n", " \n", " 1\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 39.8\n", - " 19.1\n", - " 184.0\n", - " 4650.0\n", + " Torgersen\n", + " 39.1\n", + " 18.7\n", + " 181.0\n", + " 3750.0\n", " MALE\n", " \n", " \n", " 2\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.9\n", - " 18.9\n", - " 184.0\n", - " 3900.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 47.4\n", + " 14.6\n", + " 212.0\n", + " 4725.0\n", + " FEMALE\n", " \n", " \n", " 3\n", " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 46.5\n", - " 17.9\n", - " 192.0\n", - " 3500.0\n", + " 42.5\n", + " 16.7\n", + " 187.0\n", + " 3350.0\n", " FEMALE\n", " \n", " \n", " 4\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.3\n", - " 16.8\n", - " 192.0\n", - " 3000.0\n", - " FEMALE\n", + " Biscoe\n", + " 43.2\n", + " 19.0\n", + " 197.0\n", + " 4775.0\n", + " MALE\n", " \n", " \n", " 5\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 43.2\n", - " 18.5\n", - " 192.0\n", - " 4100.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.7\n", + " 15.3\n", + " 219.0\n", + " 5200.0\n", " MALE\n", " \n", " \n", " 6\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 46.9\n", - " 16.6\n", - " 192.0\n", - " 2700.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 41.3\n", + " 21.1\n", + " 195.0\n", + " 4400.0\n", + " MALE\n", " \n", " \n", " 7\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 50.5\n", - " 18.4\n", - " 200.0\n", - " 3400.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 45.2\n", + " 13.8\n", + " 215.0\n", + " 4750.0\n", " FEMALE\n", " \n", " \n", " 8\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 49.5\n", - " 19.0\n", - " 200.0\n", - " 3800.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.5\n", + " 13.5\n", + " 210.0\n", + " 4550.0\n", + " FEMALE\n", " \n", " \n", " 9\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.2\n", - " 20.1\n", - " 200.0\n", - " 3975.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 50.5\n", + " 15.2\n", + " 216.0\n", + " 5000.0\n", + " FEMALE\n", " \n", " \n", " 10\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.8\n", - " 18.9\n", - " 208.0\n", - " 4300.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 48.2\n", + " 15.6\n", + " 221.0\n", + " 5100.0\n", " MALE\n", " \n", " \n", " 11\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 39.0\n", - " 18.7\n", - " 185.0\n", - " 3650.0\n", - " MALE\n", + " 38.1\n", + " 18.6\n", + " 190.0\n", + " 3700.0\n", + " FEMALE\n", " \n", " \n", " 12\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.0\n", - " 16.9\n", - " 185.0\n", - " 3000.0\n", - " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 50.7\n", + " 15.0\n", + " 223.0\n", + " 5550.0\n", + " MALE\n", " \n", " \n", " 13\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 47.0\n", - " 17.3\n", - " 185.0\n", - " 3700.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 37.8\n", + " 20.0\n", + " 190.0\n", + " 4250.0\n", + " MALE\n", " \n", " \n", " 14\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 34.0\n", - " 17.1\n", - " 185.0\n", - " 3400.0\n", + " Biscoe\n", + " 35.0\n", + " 17.9\n", + " 190.0\n", + " 3450.0\n", " FEMALE\n", " \n", " \n", " 15\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.0\n", - " 16.5\n", - " 185.0\n", - " 3400.0\n", - " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 48.7\n", + " 15.7\n", + " 208.0\n", + " 5350.0\n", + " MALE\n", " \n", " \n", " 16\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 45.7\n", - " 17.3\n", - " 193.0\n", - " 3600.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Torgersen\n", + " 34.6\n", + " 21.1\n", + " 198.0\n", + " 4400.0\n", + " MALE\n", " \n", " \n", " 17\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 50.6\n", - " 19.4\n", - " 193.0\n", - " 3800.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.8\n", + " 15.4\n", + " 215.0\n", + " 5150.0\n", " MALE\n", " \n", " \n", " 18\n", - " Adelie Penguin (Pygoscelis adeliae)\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 39.7\n", - " 17.9\n", - " 193.0\n", - " 4250.0\n", + " 50.3\n", + " 20.0\n", + " 197.0\n", + " 3300.0\n", " MALE\n", " \n", " \n", " 19\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 37.8\n", + " 37.2\n", " 18.1\n", - " 193.0\n", - " 3750.0\n", + " 178.0\n", + " 3900.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", + " 51.0\n", + " 18.8\n", + " 203.0\n", + " 4100.0\n", + " MALE\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", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 40.5\n", + " 17.9\n", + " 187.0\n", + " 3200.0\n", + " FEMALE\n", " \n", " \n", " 22\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.2\n", - " 17.1\n", - " 193.0\n", - " 3400.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 45.5\n", + " 13.9\n", + " 210.0\n", + " 4200.0\n", " FEMALE\n", " \n", " \n", " 23\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 36.8\n", + " 42.2\n", " 18.5\n", - " 193.0\n", - " 3500.0\n", + " 180.0\n", + " 3550.0\n", " FEMALE\n", " \n", " \n", " 24\n", " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 49.6\n", - " 18.2\n", - " 193.0\n", + " 51.7\n", + " 20.3\n", + " 194.0\n", " 3775.0\n", " MALE\n", " \n", @@ -370,86 +338,86 @@ "[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", + " species island \\\n", + "penguin_id \n", + "0 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "1 Adelie Penguin (Pygoscelis adeliae) Torgersen \n", + "2 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "4 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "5 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "6 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "7 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "8 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "9 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "10 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "12 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "13 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "14 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "15 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "16 Adelie Penguin (Pygoscelis adeliae) Torgersen \n", + "17 Gentoo penguin (Pygoscelis papua) Biscoe \n", + "18 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream \n", + "21 Adelie Penguin (Pygoscelis adeliae) Biscoe \n", + "22 Gentoo penguin (Pygoscelis papua) Biscoe \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", + "0 40.1 18.9 188.0 4300.0 \n", + "1 39.1 18.7 181.0 3750.0 \n", + "2 47.4 14.6 212.0 4725.0 \n", + "3 42.5 16.7 187.0 3350.0 \n", + "4 43.2 19.0 197.0 4775.0 \n", + "5 46.7 15.3 219.0 5200.0 \n", + "6 41.3 21.1 195.0 4400.0 \n", + "7 45.2 13.8 215.0 4750.0 \n", + "8 46.5 13.5 210.0 4550.0 \n", + "9 50.5 15.2 216.0 5000.0 \n", + "10 48.2 15.6 221.0 5100.0 \n", + "11 38.1 18.6 190.0 3700.0 \n", + "12 50.7 15.0 223.0 5550.0 \n", + "13 37.8 20.0 190.0 4250.0 \n", + "14 35.0 17.9 190.0 3450.0 \n", + "15 48.7 15.7 208.0 5350.0 \n", + "16 34.6 21.1 198.0 4400.0 \n", + "17 46.8 15.4 215.0 5150.0 \n", + "18 50.3 20.0 197.0 3300.0 \n", + "19 37.2 18.1 178.0 3900.0 \n", + "20 51.0 18.8 203.0 4100.0 \n", + "21 40.5 17.9 187.0 3200.0 \n", + "22 45.5 13.9 210.0 4200.0 \n", + "23 42.2 18.5 180.0 3550.0 \n", + "24 51.7 20.3 194.0 3775.0 \n", "\n", " sex \n", "penguin_id \n", - "0 FEMALE \n", + "0 MALE \n", "1 MALE \n", - "2 MALE \n", + "2 FEMALE \n", "3 FEMALE \n", - "4 FEMALE \n", + "4 MALE \n", "5 MALE \n", - "6 FEMALE \n", + "6 MALE \n", "7 FEMALE \n", - "8 MALE \n", - "9 MALE \n", + "8 FEMALE \n", + "9 FEMALE \n", "10 MALE \n", - "11 MALE \n", - "12 FEMALE \n", - "13 FEMALE \n", + "11 FEMALE \n", + "12 MALE \n", + "13 MALE \n", "14 FEMALE \n", - "15 FEMALE \n", - "16 FEMALE \n", + "15 MALE \n", + "16 MALE \n", "17 MALE \n", "18 MALE \n", "19 MALE \n", - "20 FEMALE \n", - "21 MALE \n", + "20 MALE \n", + "21 FEMALE \n", "22 FEMALE \n", "23 FEMALE \n", "24 MALE \n", @@ -458,7 +426,7 @@ "[334 rows x 7 columns]" ] }, - "execution_count": 18, + "execution_count": 1, "metadata": {}, "output_type": "execute_result" } @@ -490,18 +458,16 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 2, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "171160f246eb43d1832aeefb055c0851", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job deda90a8-6ec7-419c-8067-e85777bd916f is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 1408053d-cb80-4870-af28-e94b90a20a6d is DONE. 28.9 kB processed. " ] }, "metadata": {}, @@ -509,13 +475,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "eaffac40f94745728e6bd618bebd2c53", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job efe8fa0a-d450-475a-99d5-36beeb985247 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 262885fe-973c-4338-a853-227f9db4835a is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -523,13 +487,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "68e7ecdc639f4d3ab482830bf6a9da04", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 5022c56d-e605-4cab-be1b-1ecf189588a1 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job fb1dc831-7f6f-42ce-96da-1292d73919b4 is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -537,13 +499,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "ebfe197fd88348129ebe2f7d288bf4b9", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 175bd293-d448-4510-b926-1d8cfb4eb5e7 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job e79add79-f1e4-4cf0-bb97-04d153222f19 is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -551,13 +511,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "2ae69ea7da5247e8a1f7cd0e049629cb", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job a3a2e68c-f5f3-4237-99ad-44974f29d090 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job cb5ee343-f86e-4795-b0ce-d58854e72e5c is RUNNING. " ] }, "metadata": {}, @@ -596,18 +554,16 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 3, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "5ed4206cd3ad4cd485315605bf033df2", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job db3365fb-67ca-44cc-a117-88a80dc63cca is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job e65af31c-feda-468d-89c9-dec033574640 is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -615,13 +571,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "ac72db21945542558fdd62093d9dc0c3", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job ab78f7ab-a115-448b-92d0-19c091a831ca is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 0455f252-2b94-457e-bad5-672b91d9b51f is RUNNING. " ] }, "metadata": {}, @@ -667,47 +621,47 @@ " \n", " \n", " \n", - " 156\n", - " Biscoe\n", - " 46.2\n", - " 14.5\n", - " 209.0\n", - " FEMALE\n", - " Gentoo penguin (Pygoscelis papua)\n", + " 249\n", + " Torgersen\n", + " 41.1\n", + " 18.6\n", + " 189.0\n", + " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 189\n", + " 36\n", " Biscoe\n", - " 35.3\n", - " 18.9\n", - " 187.0\n", + " 43.4\n", + " 14.4\n", + " 218.0\n", " FEMALE\n", - " Adelie Penguin (Pygoscelis adeliae)\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 279\n", + " 74\n", " Biscoe\n", - " 45.1\n", - " 14.5\n", - " 215.0\n", + " 42.8\n", + " 14.2\n", + " 209.0\n", " FEMALE\n", " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 245\n", - " Biscoe\n", - " 49.5\n", - " 16.2\n", - " 229.0\n", - " MALE\n", - " Gentoo penguin (Pygoscelis papua)\n", + " 235\n", + " Dream\n", + " 34.0\n", + " 17.1\n", + " 185.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 343\n", - " Torgersen\n", - " 37.3\n", - " 20.5\n", - " 199.0\n", + " 117\n", + " Dream\n", + " 37.8\n", + " 18.1\n", + " 193.0\n", " MALE\n", " Adelie Penguin (Pygoscelis adeliae)\n", " \n", @@ -719,24 +673,24 @@ "text/plain": [ " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", "penguin_id \n", - "156 Biscoe 46.2 14.5 209.0 \n", - "189 Biscoe 35.3 18.9 187.0 \n", - "279 Biscoe 45.1 14.5 215.0 \n", - "245 Biscoe 49.5 16.2 229.0 \n", - "343 Torgersen 37.3 20.5 199.0 \n", + "249 Torgersen 41.1 18.6 189.0 \n", + "36 Biscoe 43.4 14.4 218.0 \n", + "74 Biscoe 42.8 14.2 209.0 \n", + "235 Dream 34.0 17.1 185.0 \n", + "117 Dream 37.8 18.1 193.0 \n", "\n", " sex species \n", "penguin_id \n", - "156 FEMALE Gentoo penguin (Pygoscelis papua) \n", - "189 FEMALE Adelie Penguin (Pygoscelis adeliae) \n", - "279 FEMALE Gentoo penguin (Pygoscelis papua) \n", - "245 MALE Gentoo penguin (Pygoscelis papua) \n", - "343 MALE Adelie Penguin (Pygoscelis adeliae) \n", + "249 MALE Adelie Penguin (Pygoscelis adeliae) \n", + "36 FEMALE Gentoo penguin (Pygoscelis papua) \n", + "74 FEMALE Gentoo penguin (Pygoscelis papua) \n", + "235 FEMALE Adelie Penguin (Pygoscelis adeliae) \n", + "117 MALE Adelie Penguin (Pygoscelis adeliae) \n", "\n", "[5 rows x 6 columns]" ] }, - "execution_count": 20, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } @@ -749,18 +703,16 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 4, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "d6dd794f89724099950dcc927d63d0f5", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 22a72cad-11a6-4f8e-b16d-f92853b8112e is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job d5a173bd-a7dc-42fa-8468-b088d47ccfe0 is RUNNING. " ] }, "metadata": {}, @@ -768,13 +720,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "a8ab7ca12e0d43a6803483480e837c6e", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job bc952727-8806-4fe2-abf2-c3a8a2bd9b6d is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job c6b6518b-2689-4dc1-a5b0-2a9ab75301eb is RUNNING. " ] }, "metadata": {}, @@ -810,24 +760,24 @@ " \n", " \n", " \n", - " 156\n", - " 4800.0\n", + " 249\n", + " 3325.0\n", " \n", " \n", - " 189\n", - " 3800.0\n", + " 36\n", + " 4600.0\n", " \n", " \n", - " 279\n", - " 5000.0\n", + " 74\n", + " 4700.0\n", " \n", " \n", - " 245\n", - " 5800.0\n", + " 235\n", + " 3400.0\n", " \n", " \n", - " 343\n", - " 3775.0\n", + " 117\n", + " 3750.0\n", " \n", " \n", "\n", @@ -837,16 +787,16 @@ "text/plain": [ " body_mass_g\n", "penguin_id \n", - "156 4800.0\n", - "189 3800.0\n", - "279 5000.0\n", - "245 5800.0\n", - "343 3775.0\n", + "249 3325.0\n", + "36 4600.0\n", + "74 4700.0\n", + "235 3400.0\n", + "117 3750.0\n", "\n", "[5 rows x 1 columns]" ] }, - "execution_count": 21, + "execution_count": 4, "metadata": {}, "output_type": "execute_result" } @@ -880,18 +830,16 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 5, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "380c57dc3fe54fbd8ad2fb23f1e66e37", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job f239341e-785f-43e1-bfe0-683132d6f15f is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 03a0eb1c-747e-4c2a-b7b5-d3e4e5a78134 is RUNNING. " ] }, "metadata": {}, @@ -899,13 +847,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "3db47aadba854beca71960d846838dc4", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 2d5bbbb9-efc4-4f4e-a8dc-2c7b66b0e5e0 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 70608c84-dac8-4e77-8a9e-00d823b24f37 is RUNNING. " ] }, "metadata": {}, @@ -913,13 +859,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "1de81f2944a44cbda3f16fa8a1fae813", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 66120e1c-2471-4a0c-8b82-aeb189c8866a is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job d18fdc32-2152-45d3-8c62-bf9b1556ec47 is RUNNING. " ] }, "metadata": {}, @@ -927,13 +871,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "b06cae61a4534388a4e9ed26ce442cc2", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 62825fc4-5b77-43e5-a3e4-525ebfd1285b is DONE. 2.1 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 2a022682-535f-4dc0-80ba-1640306ad9ef is RUNNING. " ] }, "metadata": {}, @@ -941,13 +883,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "977c8eae2c9848e98c5478c41af82633", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 656d1d69-b4ff-4db6-9f2d-28dcf91e2fd7 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job c145b39d-7d02-4394-80f0-fc605b2ba256 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -955,13 +895,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "aefc3085fee04c438d0327d400b4b72a", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 466507c8-1474-4725-93e5-baf8ee292e39 is DONE. 8.5 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job fc156a2b-db95-44a3-9ad1-d95b9d290080 is RUNNING. " ] }, "metadata": {}, @@ -1002,153 +940,153 @@ " \n", " \n", " 0\n", - " -1.344188\n", - " 0.642519\n", - " -1.193942\n", - " \n", - " \n", - " 1\n", - " -0.750047\n", - " 1.005876\n", - " -1.193942\n", + " -0.750505\n", + " 0.84903\n", + " -0.937262\n", " \n", " \n", " 2\n", - " -0.545811\n", - " 0.90206\n", - " -1.193942\n", + " 0.622496\n", + " -1.322402\n", + " 0.804051\n", " \n", " \n", - " 4\n", - " -1.214219\n", - " -0.188011\n", - " -0.619171\n", + " 3\n", + " -0.299107\n", + " -0.261935\n", + " -1.009817\n", " \n", " \n", " 5\n", - " -0.118772\n", - " 0.694427\n", - " -0.619171\n", + " 0.490839\n", + " -0.968913\n", + " 1.311935\n", " \n", " \n", " 6\n", - " 0.568203\n", - " -0.291828\n", - " -0.619171\n", + " -0.524806\n", + " 1.959995\n", + " -0.429379\n", " \n", " \n", " 7\n", - " 1.236611\n", - " 0.642519\n", - " -0.044401\n", + " 0.208715\n", + " -1.726389\n", + " 1.021716\n", " \n", " \n", " 9\n", - " -0.675779\n", - " 1.524957\n", - " -0.044401\n", + " 1.205551\n", + " -1.019412\n", + " 1.09427\n", " \n", " \n", " 10\n", - " -0.564378\n", - " 0.90206\n", - " 0.530369\n", - " \n", - " \n", - " 11\n", - " -0.898582\n", - " 0.798243\n", - " -1.122096\n", + " 0.772962\n", + " -0.817418\n", + " 1.457044\n", " \n", " \n", " 12\n", - " -1.26992\n", - " -0.136103\n", - " -1.122096\n", - " \n", - " \n", - " 13\n", - " 0.58677\n", - " 0.071529\n", - " -1.122096\n", + " 1.243168\n", + " -1.120408\n", + " 1.602153\n", " \n", " \n", " 14\n", - " -1.826927\n", - " -0.032287\n", - " -1.122096\n", + " -1.709725\n", + " 0.344046\n", + " -0.792152\n", " \n", " \n", - " 15\n", - " -1.26992\n", - " -0.343736\n", - " -1.122096\n", - " \n", - " \n", - " 16\n", - " 0.3454\n", - " 0.071529\n", - " -0.547325\n", + " 17\n", + " 0.509647\n", + " -0.918415\n", + " 1.021716\n", " \n", " \n", " 18\n", - " -0.768614\n", - " 0.382978\n", - " -0.547325\n", + " 1.167935\n", + " 1.404513\n", + " -0.284269\n", " \n", " \n", " 19\n", - " -1.121385\n", - " 0.486795\n", - " -0.547325\n", + " -1.295944\n", + " 0.445043\n", + " -1.662809\n", " \n", " \n", " 20\n", - " 0.512502\n", - " 0.33107\n", - " -0.547325\n", + " 1.299593\n", + " 0.798532\n", + " 0.151059\n", " \n", " \n", " 21\n", - " 1.385146\n", - " 1.057784\n", - " -0.547325\n", + " -0.675272\n", + " 0.344046\n", + " -1.009817\n", " \n", " \n", " 22\n", - " -0.675779\n", - " -0.032287\n", - " -0.547325\n", + " 0.26514\n", + " -1.675891\n", + " 0.658942\n", " \n", " \n", " 24\n", - " 1.069509\n", - " 0.538703\n", - " -0.547325\n", + " 1.43125\n", + " 1.556008\n", + " -0.501934\n", + " \n", + " \n", + " 25\n", + " 0.302756\n", + " 0.041055\n", + " -0.574488\n", " \n", " \n", " 26\n", - " -0.43441\n", - " 0.694427\n", - " 0.027445\n", + " 0.302756\n", + " -1.675891\n", + " 0.949161\n", + " \n", + " \n", + " 27\n", + " 0.227523\n", + " -1.776888\n", + " 0.658942\n", " \n", " \n", " 28\n", - " 1.923586\n", - " 1.888314\n", - " 0.027445\n", + " 1.318401\n", + " -0.362932\n", + " 1.747263\n", + " \n", + " \n", + " 29\n", + " 2.202388\n", + " 1.303516\n", + " 0.441278\n", " \n", " \n", " 30\n", - " 1.292312\n", - " 0.694427\n", - " 0.027445\n", + " -0.919779\n", + " 1.959995\n", + " -0.356824\n", " \n", " \n", " 31\n", - " -1.994029\n", - " -0.551368\n", - " -1.62502\n", + " 1.036277\n", + " -0.615424\n", + " 1.747263\n", + " \n", + " \n", + " 32\n", + " -0.223874\n", + " 0.19255\n", + " -0.356824\n", " \n", " \n", "\n", @@ -1158,65 +1096,65 @@ "text/plain": [ " standard_scaled_culmen_length_mm standard_scaled_culmen_depth_mm \\\n", "penguin_id \n", - "0 -1.344188 0.642519 \n", - "1 -0.750047 1.005876 \n", - "2 -0.545811 0.90206 \n", - "4 -1.214219 -0.188011 \n", - "5 -0.118772 0.694427 \n", - "6 0.568203 -0.291828 \n", - "7 1.236611 0.642519 \n", - "9 -0.675779 1.524957 \n", - "10 -0.564378 0.90206 \n", - "11 -0.898582 0.798243 \n", - "12 -1.26992 -0.136103 \n", - "13 0.58677 0.071529 \n", - "14 -1.826927 -0.032287 \n", - "15 -1.26992 -0.343736 \n", - "16 0.3454 0.071529 \n", - "18 -0.768614 0.382978 \n", - "19 -1.121385 0.486795 \n", - "20 0.512502 0.33107 \n", - "21 1.385146 1.057784 \n", - "22 -0.675779 -0.032287 \n", - "24 1.069509 0.538703 \n", - "26 -0.43441 0.694427 \n", - "28 1.923586 1.888314 \n", - "30 1.292312 0.694427 \n", - "31 -1.994029 -0.551368 \n", + "0 -0.750505 0.84903 \n", + "2 0.622496 -1.322402 \n", + "3 -0.299107 -0.261935 \n", + "5 0.490839 -0.968913 \n", + "6 -0.524806 1.959995 \n", + "7 0.208715 -1.726389 \n", + "9 1.205551 -1.019412 \n", + "10 0.772962 -0.817418 \n", + "12 1.243168 -1.120408 \n", + "14 -1.709725 0.344046 \n", + "17 0.509647 -0.918415 \n", + "18 1.167935 1.404513 \n", + "19 -1.295944 0.445043 \n", + "20 1.299593 0.798532 \n", + "21 -0.675272 0.344046 \n", + "22 0.26514 -1.675891 \n", + "24 1.43125 1.556008 \n", + "25 0.302756 0.041055 \n", + "26 0.302756 -1.675891 \n", + "27 0.227523 -1.776888 \n", + "28 1.318401 -0.362932 \n", + "29 2.202388 1.303516 \n", + "30 -0.919779 1.959995 \n", + "31 1.036277 -0.615424 \n", + "32 -0.223874 0.19255 \n", "\n", " standard_scaled_flipper_length_mm \n", "penguin_id \n", - "0 -1.193942 \n", - "1 -1.193942 \n", - "2 -1.193942 \n", - "4 -0.619171 \n", - "5 -0.619171 \n", - "6 -0.619171 \n", - "7 -0.044401 \n", - "9 -0.044401 \n", - "10 0.530369 \n", - "11 -1.122096 \n", - "12 -1.122096 \n", - "13 -1.122096 \n", - "14 -1.122096 \n", - "15 -1.122096 \n", - "16 -0.547325 \n", - "18 -0.547325 \n", - "19 -0.547325 \n", - "20 -0.547325 \n", - "21 -0.547325 \n", - "22 -0.547325 \n", - "24 -0.547325 \n", - "26 0.027445 \n", - "28 0.027445 \n", - "30 0.027445 \n", - "31 -1.62502 \n", + "0 -0.937262 \n", + "2 0.804051 \n", + "3 -1.009817 \n", + "5 1.311935 \n", + "6 -0.429379 \n", + "7 1.021716 \n", + "9 1.09427 \n", + "10 1.457044 \n", + "12 1.602153 \n", + "14 -0.792152 \n", + "17 1.021716 \n", + "18 -0.284269 \n", + "19 -1.662809 \n", + "20 0.151059 \n", + "21 -1.009817 \n", + "22 0.658942 \n", + "24 -0.501934 \n", + "25 -0.574488 \n", + "26 0.949161 \n", + "27 0.658942 \n", + "28 1.747263 \n", + "29 0.441278 \n", + "30 -0.356824 \n", + "31 1.747263 \n", + "32 -0.356824 \n", "...\n", "\n", "[267 rows x 3 columns]" ] }, - "execution_count": 22, + "execution_count": 5, "metadata": {}, "output_type": "execute_result" } @@ -1237,32 +1175,16 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": 6, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "74f3c24c0a434e12bf6a56dc4809b501", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job c6268b07-0d3d-4fe0-971d-cc99fd98cd7e is RUNNING. Open Job" + ], "text/plain": [ - "HTML(value='Query job 31550d88-fc7b-4fcb-9975-9ed24bf2e009 is RUNNING. " ] }, "metadata": {}, @@ -1270,13 +1192,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "5a04e46a7d0248b1ae523f2ca6903ee8", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 1e17f5f7-2956-4bdd-baa9-c07591481341 is DONE. 536 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 5ec7c8b1-037c-466c-a51e-963f8274e76b is RUNNING. " ] }, "metadata": {}, @@ -1284,13 +1204,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "62563820bfb245be85bbc1bf3dfb993c", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e2fde7a6-67b4-45a4-91d4-1cb9eff66ae5 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 4e860716-bc41-4ef6-83ff-310d085ed7cc is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -1298,13 +1216,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "98aff3bfded44868bf120451c89df9f5", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e0683619-23c5-44fd-8930-9d3c9d02729a is DONE. 2.1 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 6b96a757-42fe-4b65-92fd-a3ae339fe769 is RUNNING. " ] }, "metadata": {}, @@ -1344,154 +1260,154 @@ " \n", " \n", " \n", - " 3\n", - " 0.493935\n", - " 0.382978\n", - " -0.619171\n", - " \n", - " \n", - " 8\n", - " 1.050942\n", - " 0.953968\n", - " -0.044401\n", + " 1\n", + " -0.938587\n", + " 0.748033\n", + " -1.445145\n", " \n", " \n", - " 17\n", - " 1.255178\n", - " 1.1616\n", - " -0.547325\n", + " 4\n", + " -0.16745\n", + " 0.899528\n", + " -0.284269\n", " \n", " \n", - " 23\n", - " -1.307054\n", - " 0.694427\n", - " -0.547325\n", + " 8\n", + " 0.453222\n", + " -1.877885\n", + " 0.658942\n", " \n", " \n", - " 25\n", - " 1.515114\n", - " 0.486795\n", - " 0.027445\n", + " 11\n", + " -1.12667\n", + " 0.697535\n", + " -0.792152\n", " \n", " \n", - " 27\n", - " 1.236611\n", - " 1.265417\n", - " 0.027445\n", + " 13\n", + " -1.183094\n", + " 1.404513\n", + " -0.792152\n", " \n", " \n", - " 29\n", - " 1.403713\n", - " 0.953968\n", - " 0.027445\n", + " 15\n", + " 0.867003\n", + " -0.766919\n", + " 0.513833\n", " \n", " \n", - " 34\n", - " 0.419668\n", - " 0.538703\n", - " -1.62502\n", + " 16\n", + " -1.784958\n", + " 1.959995\n", + " -0.211715\n", " \n", " \n", - " 35\n", - " -1.455589\n", - " 0.694427\n", - " -1.050249\n", + " 23\n", + " -0.355532\n", + " 0.647036\n", + " -1.5177\n", " \n", " \n", - " 39\n", - " 0.326833\n", - " 1.1616\n", - " -0.475479\n", + " 34\n", + " -0.600039\n", + " -1.776888\n", + " 0.949161\n", " \n", " \n", - " 51\n", - " -1.065684\n", - " 0.227254\n", - " -0.978403\n", + " 36\n", + " -0.129833\n", + " -1.423399\n", + " 1.23938\n", " \n", " \n", - " 52\n", - " -0.248741\n", - " 0.071529\n", - " -0.978403\n", + " 42\n", + " -1.615684\n", + " -0.514427\n", + " -0.429379\n", " \n", " \n", - " 60\n", - " 0.531069\n", - " 0.382978\n", - " -0.403633\n", + " 48\n", + " 0.415606\n", + " -0.716421\n", + " 1.021716\n", " \n", " \n", " 61\n", - " 0.401101\n", - " 0.90206\n", - " -0.403633\n", + " 0.396797\n", + " -1.170907\n", + " 1.457044\n", " \n", " \n", " 64\n", - " -1.455589\n", - " 0.33107\n", - " -0.403633\n", + " 0.434414\n", + " -1.120408\n", + " 1.09427\n", " \n", " \n", " 65\n", - " -0.564378\n", - " 0.642519\n", - " -0.403633\n", + " -1.220711\n", + " 1.051024\n", + " -1.445145\n", " \n", " \n", - " 67\n", - " 1.273745\n", - " 1.317325\n", - " 0.171138\n", + " 68\n", + " -1.484026\n", + " -0.009443\n", + " -1.009817\n", " \n", " \n", - " 83\n", - " 2.629128\n", - " 0.33107\n", - " -1.409481\n", + " 70\n", + " 1.638141\n", + " 1.404513\n", + " 0.296168\n", " \n", " \n", - " 85\n", - " -1.288487\n", - " 0.746335\n", - " -0.83471\n", + " 72\n", + " 0.829387\n", + " 0.142052\n", + " -0.719598\n", " \n", " \n", - " 93\n", - " -0.508677\n", - " 0.486795\n", - " 0.314831\n", + " 74\n", + " -0.242683\n", + " -1.524396\n", + " 0.586387\n", " \n", " \n", - " 104\n", - " 0.382534\n", - " -0.032287\n", - " -0.762864\n", + " 77\n", + " -1.277136\n", + " -0.211437\n", + " -0.647043\n", " \n", " \n", - " 105\n", - " -1.065684\n", - " 0.746335\n", - " -0.762864\n", + " 81\n", + " 0.208715\n", + " -1.221405\n", + " 0.804051\n", + " \n", + " \n", + " 91\n", + " 1.261976\n", + " 0.647036\n", + " 0.005949\n", " \n", " \n", - " 108\n", - " 1.162343\n", - " 0.382978\n", - " -0.762864\n", + " 96\n", + " 0.246331\n", + " -1.322402\n", + " 0.731497\n", " \n", " \n", - " 113\n", - " 1.496547\n", - " 1.213509\n", - " 0.386677\n", + " 105\n", + " -1.803766\n", + " 0.445043\n", + " -1.009817\n", " \n", " \n", - " 130\n", - " -0.341575\n", - " 1.213509\n", - " -0.044401\n", + " 111\n", + " -1.164286\n", + " 0.697535\n", + " -2.098138\n", " \n", " \n", "\n", @@ -1501,65 +1417,65 @@ "text/plain": [ " standard_scaled_culmen_length_mm standard_scaled_culmen_depth_mm \\\n", "penguin_id \n", - "3 0.493935 0.382978 \n", - "8 1.050942 0.953968 \n", - "17 1.255178 1.1616 \n", - "23 -1.307054 0.694427 \n", - "25 1.515114 0.486795 \n", - "27 1.236611 1.265417 \n", - "29 1.403713 0.953968 \n", - "34 0.419668 0.538703 \n", - "35 -1.455589 0.694427 \n", - "39 0.326833 1.1616 \n", - "51 -1.065684 0.227254 \n", - "52 -0.248741 0.071529 \n", - "60 0.531069 0.382978 \n", - "61 0.401101 0.90206 \n", - "64 -1.455589 0.33107 \n", - "65 -0.564378 0.642519 \n", - "67 1.273745 1.317325 \n", - "83 2.629128 0.33107 \n", - "85 -1.288487 0.746335 \n", - "93 -0.508677 0.486795 \n", - "104 0.382534 -0.032287 \n", - "105 -1.065684 0.746335 \n", - "108 1.162343 0.382978 \n", - "113 1.496547 1.213509 \n", - "130 -0.341575 1.213509 \n", + "1 -0.938587 0.748033 \n", + "4 -0.16745 0.899528 \n", + "8 0.453222 -1.877885 \n", + "11 -1.12667 0.697535 \n", + "13 -1.183094 1.404513 \n", + "15 0.867003 -0.766919 \n", + "16 -1.784958 1.959995 \n", + "23 -0.355532 0.647036 \n", + "34 -0.600039 -1.776888 \n", + "36 -0.129833 -1.423399 \n", + "42 -1.615684 -0.514427 \n", + "48 0.415606 -0.716421 \n", + "61 0.396797 -1.170907 \n", + "64 0.434414 -1.120408 \n", + "65 -1.220711 1.051024 \n", + "68 -1.484026 -0.009443 \n", + "70 1.638141 1.404513 \n", + "72 0.829387 0.142052 \n", + "74 -0.242683 -1.524396 \n", + "77 -1.277136 -0.211437 \n", + "81 0.208715 -1.221405 \n", + "91 1.261976 0.647036 \n", + "96 0.246331 -1.322402 \n", + "105 -1.803766 0.445043 \n", + "111 -1.164286 0.697535 \n", "\n", " standard_scaled_flipper_length_mm \n", "penguin_id \n", - "3 -0.619171 \n", - "8 -0.044401 \n", - "17 -0.547325 \n", - "23 -0.547325 \n", - "25 0.027445 \n", - "27 0.027445 \n", - "29 0.027445 \n", - "34 -1.62502 \n", - "35 -1.050249 \n", - "39 -0.475479 \n", - "51 -0.978403 \n", - "52 -0.978403 \n", - "60 -0.403633 \n", - "61 -0.403633 \n", - "64 -0.403633 \n", - "65 -0.403633 \n", - "67 0.171138 \n", - "83 -1.409481 \n", - "85 -0.83471 \n", - "93 0.314831 \n", - "104 -0.762864 \n", - "105 -0.762864 \n", - "108 -0.762864 \n", - "113 0.386677 \n", - "130 -0.044401 \n", + "1 -1.445145 \n", + "4 -0.284269 \n", + "8 0.658942 \n", + "11 -0.792152 \n", + "13 -0.792152 \n", + "15 0.513833 \n", + "16 -0.211715 \n", + "23 -1.5177 \n", + "34 0.949161 \n", + "36 1.23938 \n", + "42 -0.429379 \n", + "48 1.021716 \n", + "61 1.457044 \n", + "64 1.09427 \n", + "65 -1.445145 \n", + "68 -1.009817 \n", + "70 0.296168 \n", + "72 -0.719598 \n", + "74 0.586387 \n", + "77 -0.647043 \n", + "81 0.804051 \n", + "91 0.005949 \n", + "96 0.731497 \n", + "105 -1.009817 \n", + "111 -2.098138 \n", "...\n", "\n", "[67 rows x 3 columns]" ] }, - "execution_count": 23, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -1581,32 +1497,16 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": 7, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "d642a617d27f4e2493c80dbdd1686193", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job a8d8afa4-d91e-487e-8709-8727a73ab453 is RUNNING. Open Job" + ], "text/plain": [ - "HTML(value='Query job b9afd624-4345-4160-8809-05786563ce35 is RUNNING. " ] }, "metadata": {}, @@ -1614,13 +1514,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "09217776c2294e8b929a56e7a73fbfa8", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 41962e2e-4d14-4053-9297-3ce61699551a is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job c918fc7c-a956-4259-b5c5-09c2eac615cd is RUNNING. " ] }, "metadata": {}, @@ -1628,13 +1526,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "9c1581fc9fcb49739d1d81b73506b894", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 5d3c22c9-c972-4213-8557-726c9e0aca37 is DONE. 22.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 1d855341-282f-4d10-9ba9-3ce6683b729a is RUNNING. " ] }, "metadata": {}, @@ -1642,13 +1538,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "d7749eb7cf554697a60c90f3718ad582", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 9cb7b33f-ea05-4cf4-9f92-bb3aa4ea8d10 is DONE. 2.1 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job c257ff78-3e15-4296-82f5-ba6c2eb6a6ff is RUNNING. " ] }, "metadata": {}, @@ -1656,13 +1550,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "e900465918224249bccc781d992aadbb", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job fe1f35d6-d82c-4aab-a284-637b72554f5b is DONE. 29.2 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job a17eec0c-10d0-4943-95be-60fced57d5cb is RUNNING. " ] }, "metadata": {}, @@ -1670,13 +1562,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "b0272ee35c5745a491b7c5883b3fbb1b", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 37bc90ff-59cb-4b0c-8f9d-73bcda43524a is DONE. 536 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 1db53c8a-cf45-4c69-a443-6b7a49fc3a07 is DONE. 536 Bytes processed. " ] }, "metadata": {}, @@ -1684,13 +1574,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "00f9d4b55bb94997aaebdae298cefab3", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e23f4724-fdd8-45a9-8c87-defd8d471035 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job ae870ee3-e633-4556-94e6-6669fa0bfde2 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -1698,13 +1586,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "9cd8e791be5844669cba10dc53f862ae", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 257378db-0569-42d7-965a-7757154c710b is DONE. 21.4 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job baa7c70c-eabc-49e1-bae9-fdd4891cdb6e is RUNNING. " ] }, "metadata": {}, @@ -1751,226 +1637,226 @@ " \n", " \n", " 0\n", - " [{'index': 2, 'value': 1.0}]\n", - " -1.344188\n", - " 0.642519\n", - " -1.193942\n", - " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", - " \n", - " \n", - " 1\n", + " -0.750505\n", + " 0.84903\n", + " -0.937262\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.750047\n", - " 1.005876\n", - " -1.193942\n", - " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 2\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.545811\n", - " 0.90206\n", - " -1.193942\n", - " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", + " 0.622496\n", + " -1.322402\n", + " 0.804051\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 4\n", - " [{'index': 2, 'value': 1.0}]\n", - " -1.214219\n", - " -0.188011\n", - " -0.619171\n", + " 3\n", " [{'index': 2, 'value': 1.0}]\n", + " -0.299107\n", + " -0.261935\n", + " -1.009817\n", " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", " 5\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.490839\n", + " -0.968913\n", + " 1.311935\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.118772\n", - " 0.694427\n", - " -0.619171\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 6\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.524806\n", + " 1.959995\n", + " -0.429379\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.568203\n", - " -0.291828\n", - " -0.619171\n", - " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 7\n", - " [{'index': 2, 'value': 1.0}]\n", - " 1.236611\n", - " 0.642519\n", - " -0.044401\n", - " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.208715\n", + " -1.726389\n", + " 1.021716\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 9\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.675779\n", - " 1.524957\n", - " -0.044401\n", - " [{'index': 3, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", + " 1.205551\n", + " -1.019412\n", + " 1.09427\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 10\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.772962\n", + " -0.817418\n", + " 1.457044\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.564378\n", - " 0.90206\n", - " 0.530369\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 11\n", + " 12\n", + " [{'index': 1, 'value': 1.0}]\n", + " 1.243168\n", + " -1.120408\n", + " 1.602153\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.898582\n", - " 0.798243\n", - " -1.122096\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 12\n", - " [{'index': 2, 'value': 1.0}]\n", - " -1.26992\n", - " -0.136103\n", - " -1.122096\n", - " [{'index': 2, 'value': 1.0}]\n", + " 14\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.709725\n", + " 0.344046\n", + " -0.792152\n", + " [{'index': 1, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 13\n", - " [{'index': 2, 'value': 1.0}]\n", - " 0.58677\n", - " 0.071529\n", - " -1.122096\n", - " [{'index': 2, 'value': 1.0}]\n", + " 17\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.509647\n", + " -0.918415\n", + " 1.021716\n", " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 14\n", + " 18\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.167935\n", + " 1.404513\n", + " -0.284269\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.826927\n", - " -0.032287\n", - " -1.122096\n", " [{'index': 2, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 15\n", + " 19\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.26992\n", - " -0.343736\n", - " -1.122096\n", + " -1.295944\n", + " 0.445043\n", + " -1.662809\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 16\n", + " 20\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.3454\n", - " 0.071529\n", - " -0.547325\n", + " 1.299593\n", + " 0.798532\n", + " 0.151059\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 18\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.768614\n", - " 0.382978\n", - " -0.547325\n", - " [{'index': 3, 'value': 1.0}]\n", + " 21\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.675272\n", + " 0.344046\n", + " -1.009817\n", + " [{'index': 1, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 19\n", - " [{'index': 2, 'value': 1.0}]\n", - " -1.121385\n", - " 0.486795\n", - " -0.547325\n", - " [{'index': 3, 'value': 1.0}]\n", + " 22\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.26514\n", + " -1.675891\n", + " 0.658942\n", " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 20\n", + " 24\n", " [{'index': 2, 'value': 1.0}]\n", - " 0.512502\n", - " 0.33107\n", - " -0.547325\n", + " 1.43125\n", + " 1.556008\n", + " -0.501934\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 21\n", + " 25\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.385146\n", - " 1.057784\n", - " -0.547325\n", - " [{'index': 3, 'value': 1.0}]\n", + " 0.302756\n", + " 0.041055\n", + " -0.574488\n", + " [{'index': 1, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 22\n", - " [{'index': 2, 'value': 1.0}]\n", - " -0.675779\n", - " -0.032287\n", - " -0.547325\n", - " [{'index': 2, 'value': 1.0}]\n", + " 26\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.302756\n", + " -1.675891\n", + " 0.949161\n", " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 24\n", - " [{'index': 2, 'value': 1.0}]\n", - " 1.069509\n", - " 0.538703\n", - " -0.547325\n", + " 27\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.227523\n", + " -1.776888\n", + " 0.658942\n", + " [{'index': 1, 'value': 1.0}]\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 26\n", + " 28\n", + " [{'index': 1, 'value': 1.0}]\n", + " 1.318401\n", + " -0.362932\n", + " 1.747263\n", " [{'index': 2, 'value': 1.0}]\n", - " -0.43441\n", - " 0.694427\n", - " 0.027445\n", " [{'index': 3, 'value': 1.0}]\n", - " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 28\n", + " 29\n", + " [{'index': 2, 'value': 1.0}]\n", + " 2.202388\n", + " 1.303516\n", + " 0.441278\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.923586\n", - " 1.888314\n", - " 0.027445\n", - " [{'index': 3, 'value': 1.0}]\n", " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", " 30\n", " [{'index': 2, 'value': 1.0}]\n", - " 1.292312\n", - " 0.694427\n", - " 0.027445\n", - " [{'index': 3, 'value': 1.0}]\n", + " -0.919779\n", + " 1.959995\n", + " -0.356824\n", " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", " 31\n", + " [{'index': 1, 'value': 1.0}]\n", + " 1.036277\n", + " -0.615424\n", + " 1.747263\n", " [{'index': 2, 'value': 1.0}]\n", - " -1.994029\n", - " -0.551368\n", - " -1.62502\n", + " [{'index': 3, 'value': 1.0}]\n", + " \n", + " \n", + " 32\n", + " [{'index': 3, 'value': 1.0}]\n", + " -0.223874\n", + " 0.19255\n", + " -0.356824\n", " [{'index': 2, 'value': 1.0}]\n", " [{'index': 1, 'value': 1.0}]\n", " \n", @@ -1982,121 +1868,121 @@ "text/plain": [ " onehotencoded_island standard_scaled_culmen_length_mm \\\n", "penguin_id \n", - "0 [{'index': 2, 'value': 1.0}] -1.344188 \n", - "1 [{'index': 2, 'value': 1.0}] -0.750047 \n", - "2 [{'index': 2, 'value': 1.0}] -0.545811 \n", - "4 [{'index': 2, 'value': 1.0}] -1.214219 \n", - "5 [{'index': 2, 'value': 1.0}] -0.118772 \n", - "6 [{'index': 2, 'value': 1.0}] 0.568203 \n", - "7 [{'index': 2, 'value': 1.0}] 1.236611 \n", - "9 [{'index': 2, 'value': 1.0}] -0.675779 \n", - "10 [{'index': 2, 'value': 1.0}] -0.564378 \n", - "11 [{'index': 2, 'value': 1.0}] -0.898582 \n", - "12 [{'index': 2, 'value': 1.0}] -1.26992 \n", - "13 [{'index': 2, 'value': 1.0}] 0.58677 \n", - "14 [{'index': 2, 'value': 1.0}] -1.826927 \n", - "15 [{'index': 2, 'value': 1.0}] -1.26992 \n", - "16 [{'index': 2, 'value': 1.0}] 0.3454 \n", - "18 [{'index': 2, 'value': 1.0}] -0.768614 \n", - "19 [{'index': 2, 'value': 1.0}] -1.121385 \n", - "20 [{'index': 2, 'value': 1.0}] 0.512502 \n", - "21 [{'index': 2, 'value': 1.0}] 1.385146 \n", - "22 [{'index': 2, 'value': 1.0}] -0.675779 \n", - "24 [{'index': 2, 'value': 1.0}] 1.069509 \n", - "26 [{'index': 2, 'value': 1.0}] -0.43441 \n", - "28 [{'index': 2, 'value': 1.0}] 1.923586 \n", - "30 [{'index': 2, 'value': 1.0}] 1.292312 \n", - "31 [{'index': 2, 'value': 1.0}] -1.994029 \n", + "0 [{'index': 1, 'value': 1.0}] -0.750505 \n", + "2 [{'index': 1, 'value': 1.0}] 0.622496 \n", + "3 [{'index': 2, 'value': 1.0}] -0.299107 \n", + "5 [{'index': 1, 'value': 1.0}] 0.490839 \n", + "6 [{'index': 1, 'value': 1.0}] -0.524806 \n", + "7 [{'index': 1, 'value': 1.0}] 0.208715 \n", + "9 [{'index': 1, 'value': 1.0}] 1.205551 \n", + "10 [{'index': 1, 'value': 1.0}] 0.772962 \n", + "12 [{'index': 1, 'value': 1.0}] 1.243168 \n", + "14 [{'index': 1, 'value': 1.0}] -1.709725 \n", + "17 [{'index': 1, 'value': 1.0}] 0.509647 \n", + "18 [{'index': 2, 'value': 1.0}] 1.167935 \n", + "19 [{'index': 2, 'value': 1.0}] -1.295944 \n", + "20 [{'index': 2, 'value': 1.0}] 1.299593 \n", + "21 [{'index': 1, 'value': 1.0}] -0.675272 \n", + "22 [{'index': 1, 'value': 1.0}] 0.26514 \n", + "24 [{'index': 2, 'value': 1.0}] 1.43125 \n", + "25 [{'index': 2, 'value': 1.0}] 0.302756 \n", + "26 [{'index': 1, 'value': 1.0}] 0.302756 \n", + "27 [{'index': 1, 'value': 1.0}] 0.227523 \n", + "28 [{'index': 1, 'value': 1.0}] 1.318401 \n", + "29 [{'index': 2, 'value': 1.0}] 2.202388 \n", + "30 [{'index': 2, 'value': 1.0}] -0.919779 \n", + "31 [{'index': 1, 'value': 1.0}] 1.036277 \n", + "32 [{'index': 3, 'value': 1.0}] -0.223874 \n", "\n", " standard_scaled_culmen_depth_mm \\\n", "penguin_id \n", - "0 0.642519 \n", - "1 1.005876 \n", - "2 0.90206 \n", - "4 -0.188011 \n", - "5 0.694427 \n", - "6 -0.291828 \n", - "7 0.642519 \n", - "9 1.524957 \n", - "10 0.90206 \n", - "11 0.798243 \n", - "12 -0.136103 \n", - "13 0.071529 \n", - "14 -0.032287 \n", - "15 -0.343736 \n", - "16 0.071529 \n", - "18 0.382978 \n", - "19 0.486795 \n", - "20 0.33107 \n", - "21 1.057784 \n", - "22 -0.032287 \n", - "24 0.538703 \n", - "26 0.694427 \n", - "28 1.888314 \n", - "30 0.694427 \n", - "31 -0.551368 \n", + "0 0.84903 \n", + "2 -1.322402 \n", + "3 -0.261935 \n", + "5 -0.968913 \n", + "6 1.959995 \n", + "7 -1.726389 \n", + "9 -1.019412 \n", + "10 -0.817418 \n", + "12 -1.120408 \n", + "14 0.344046 \n", + "17 -0.918415 \n", + "18 1.404513 \n", + "19 0.445043 \n", + "20 0.798532 \n", + "21 0.344046 \n", + "22 -1.675891 \n", + "24 1.556008 \n", + "25 0.041055 \n", + "26 -1.675891 \n", + "27 -1.776888 \n", + "28 -0.362932 \n", + "29 1.303516 \n", + "30 1.959995 \n", + "31 -0.615424 \n", + "32 0.19255 \n", "\n", " standard_scaled_flipper_length_mm onehotencoded_sex \\\n", "penguin_id \n", - "0 -1.193942 [{'index': 2, 'value': 1.0}] \n", - "1 -1.193942 [{'index': 3, 'value': 1.0}] \n", - "2 -1.193942 [{'index': 3, 'value': 1.0}] \n", - "4 -0.619171 [{'index': 2, 'value': 1.0}] \n", - "5 -0.619171 [{'index': 3, 'value': 1.0}] \n", - "6 -0.619171 [{'index': 2, 'value': 1.0}] \n", - "7 -0.044401 [{'index': 2, 'value': 1.0}] \n", - "9 -0.044401 [{'index': 3, 'value': 1.0}] \n", - "10 0.530369 [{'index': 3, 'value': 1.0}] \n", - "11 -1.122096 [{'index': 3, 'value': 1.0}] \n", - "12 -1.122096 [{'index': 2, 'value': 1.0}] \n", - "13 -1.122096 [{'index': 2, 'value': 1.0}] \n", - "14 -1.122096 [{'index': 2, 'value': 1.0}] \n", - "15 -1.122096 [{'index': 2, 'value': 1.0}] \n", - "16 -0.547325 [{'index': 2, 'value': 1.0}] \n", - "18 -0.547325 [{'index': 3, 'value': 1.0}] \n", - "19 -0.547325 [{'index': 3, 'value': 1.0}] \n", - "20 -0.547325 [{'index': 2, 'value': 1.0}] \n", - "21 -0.547325 [{'index': 3, 'value': 1.0}] \n", - "22 -0.547325 [{'index': 2, 'value': 1.0}] \n", - "24 -0.547325 [{'index': 3, 'value': 1.0}] \n", - "26 0.027445 [{'index': 3, 'value': 1.0}] \n", - "28 0.027445 [{'index': 3, 'value': 1.0}] \n", - "30 0.027445 [{'index': 3, 'value': 1.0}] \n", - "31 -1.62502 [{'index': 2, 'value': 1.0}] \n", + "0 -0.937262 [{'index': 2, 'value': 1.0}] \n", + "2 0.804051 [{'index': 1, 'value': 1.0}] \n", + "3 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "5 1.311935 [{'index': 2, 'value': 1.0}] \n", + "6 -0.429379 [{'index': 2, 'value': 1.0}] \n", + "7 1.021716 [{'index': 1, 'value': 1.0}] \n", + "9 1.09427 [{'index': 1, 'value': 1.0}] \n", + "10 1.457044 [{'index': 2, 'value': 1.0}] \n", + "12 1.602153 [{'index': 2, 'value': 1.0}] \n", + "14 -0.792152 [{'index': 1, 'value': 1.0}] \n", + "17 1.021716 [{'index': 2, 'value': 1.0}] \n", + "18 -0.284269 [{'index': 2, 'value': 1.0}] \n", + "19 -1.662809 [{'index': 2, 'value': 1.0}] \n", + "20 0.151059 [{'index': 2, 'value': 1.0}] \n", + "21 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "22 0.658942 [{'index': 1, 'value': 1.0}] \n", + "24 -0.501934 [{'index': 2, 'value': 1.0}] \n", + "25 -0.574488 [{'index': 1, 'value': 1.0}] \n", + "26 0.949161 [{'index': 1, 'value': 1.0}] \n", + "27 0.658942 [{'index': 1, 'value': 1.0}] \n", + "28 1.747263 [{'index': 2, 'value': 1.0}] \n", + "29 0.441278 [{'index': 2, 'value': 1.0}] \n", + "30 -0.356824 [{'index': 2, 'value': 1.0}] \n", + "31 1.747263 [{'index': 2, 'value': 1.0}] \n", + "32 -0.356824 [{'index': 2, 'value': 1.0}] \n", "\n", " onehotencoded_species \n", "penguin_id \n", "0 [{'index': 1, 'value': 1.0}] \n", - "1 [{'index': 1, 'value': 1.0}] \n", - "2 [{'index': 1, 'value': 1.0}] \n", - "4 [{'index': 1, 'value': 1.0}] \n", - "5 [{'index': 1, 'value': 1.0}] \n", - "6 [{'index': 2, 'value': 1.0}] \n", - "7 [{'index': 2, 'value': 1.0}] \n", - "9 [{'index': 1, 'value': 1.0}] \n", - "10 [{'index': 1, 'value': 1.0}] \n", - "11 [{'index': 1, 'value': 1.0}] \n", - "12 [{'index': 1, 'value': 1.0}] \n", - "13 [{'index': 2, 'value': 1.0}] \n", + "2 [{'index': 3, 'value': 1.0}] \n", + "3 [{'index': 2, 'value': 1.0}] \n", + "5 [{'index': 3, 'value': 1.0}] \n", + "6 [{'index': 1, 'value': 1.0}] \n", + "7 [{'index': 3, 'value': 1.0}] \n", + "9 [{'index': 3, 'value': 1.0}] \n", + "10 [{'index': 3, 'value': 1.0}] \n", + "12 [{'index': 3, 'value': 1.0}] \n", "14 [{'index': 1, 'value': 1.0}] \n", - "15 [{'index': 1, 'value': 1.0}] \n", - "16 [{'index': 2, 'value': 1.0}] \n", - "18 [{'index': 1, 'value': 1.0}] \n", + "17 [{'index': 3, 'value': 1.0}] \n", + "18 [{'index': 2, 'value': 1.0}] \n", "19 [{'index': 1, 'value': 1.0}] \n", "20 [{'index': 2, 'value': 1.0}] \n", - "21 [{'index': 2, 'value': 1.0}] \n", - "22 [{'index': 1, 'value': 1.0}] \n", + "21 [{'index': 1, 'value': 1.0}] \n", + "22 [{'index': 3, 'value': 1.0}] \n", "24 [{'index': 2, 'value': 1.0}] \n", - "26 [{'index': 1, 'value': 1.0}] \n", - "28 [{'index': 2, 'value': 1.0}] \n", - "30 [{'index': 2, 'value': 1.0}] \n", - "31 [{'index': 1, 'value': 1.0}] \n", + "25 [{'index': 2, 'value': 1.0}] \n", + "26 [{'index': 3, 'value': 1.0}] \n", + "27 [{'index': 3, 'value': 1.0}] \n", + "28 [{'index': 3, 'value': 1.0}] \n", + "29 [{'index': 2, 'value': 1.0}] \n", + "30 [{'index': 1, 'value': 1.0}] \n", + "31 [{'index': 3, 'value': 1.0}] \n", + "32 [{'index': 1, 'value': 1.0}] \n", "...\n", "\n", "[267 rows x 6 columns]" ] }, - "execution_count": 24, + "execution_count": 7, "metadata": {}, "output_type": "execute_result" } @@ -2138,18 +2024,28 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": 8, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "5db4c5c80ba4417db151aa561dab5ee7", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 7d9c9f8b-6b4c-451f-ae3d-06fb7090d148 is DONE. 21.4 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job be87ccfa-72ab-4858-9d4a-b2f5f8b2a5e6 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job ceced0cc-13a7-4b14-b42c-4d5f69e7e49a is RUNNING. " ] }, "metadata": {}, @@ -2157,13 +2053,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "e6b05d83de0e496d9e47392762046fc5", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 2d651fac-11bf-42da-8c18-bd33207379ca is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job a708b8df-6040-49b1-a6da-d2c0d162f247 is RUNNING. " ] }, "metadata": {}, @@ -2171,13 +2065,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "24d616c24a844abfbfd77ebd9f28486a", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 58836ccc-242b-4574-bc48-4c269e74dbf1 is DONE. 5.7 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job e9b9cbb5-f6a4-4d85-ba78-1edae77dce94 is RUNNING. " ] }, "metadata": {}, @@ -2185,13 +2077,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "ce49b66c6fa0460aa3ee28746765b6ac", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 1bf531f0-0fde-489b-ab36-6040a2a12377 is DONE. 536 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 6c0a41a7-a732-413a-b074-ba82f175eab8 is RUNNING. " ] }, "metadata": {}, @@ -2199,13 +2089,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "4a6010d73ca04ea9a133de99aa90da3c", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 4245f4e6-4d5b-404f-81d7-50f0553e2456 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 2d08b79d-9c36-4db7-824a-332fdd02e9fc is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -2213,13 +2101,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "ce9cfdca964a4062a52ebaae9d13ae59", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job ed951699-c005-450e-a8b6-0916ec234e7f is DONE. 5.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 7fa0bf53-1022-45ee-b3ac-78fa5c155585 is RUNNING. " ] }, "metadata": {}, @@ -2247,152 +2133,397 @@ " \n", " \n", " predicted_body_mass_g\n", + " onehotencoded_island\n", + " standard_scaled_culmen_length_mm\n", + " standard_scaled_culmen_depth_mm\n", + " standard_scaled_flipper_length_mm\n", + " onehotencoded_sex\n", + " onehotencoded_species\n", " \n", " \n", " penguin_id\n", " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", - " 3\n", - " 3394.118128\n", + " 1\n", + " 3781.402407\n", + " [{'index': 3, 'value': 1.0}]\n", + " -0.938587\n", + " 0.748033\n", + " -1.445145\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 8\n", - " 4048.685642\n", + " 4\n", + " 4124.107944\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.16745\n", + " 0.899528\n", + " -0.284269\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 17\n", - " 3976.454093\n", + " 8\n", + " 4670.344196\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.453222\n", + " -1.877885\n", + " 0.658942\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 23\n", - " 3541.582194\n", + " 11\n", + " 3529.417214\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.12667\n", + " 0.697535\n", + " -0.792152\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 25\n", - " 4032.844186\n", + " 13\n", + " 4014.101714\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.183094\n", + " 1.404513\n", + " -0.792152\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 27\n", - " 4118.351772\n", + " 15\n", + " 5212.41288\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.867003\n", + " -0.766919\n", + " 0.513833\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 29\n", - " 4087.767826\n", + " 16\n", + " 4163.595615\n", + " [{'index': 3, 'value': 1.0}]\n", + " -1.784958\n", + " 1.959995\n", + " -0.211715\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 34\n", - " 3183.755249\n", + " 23\n", + " 3392.453069\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.355532\n", + " 0.647036\n", + " -1.5177\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 35\n", - " 3418.802274\n", - " \n", - " \n", - " 39\n", - " 3519.186468\n", + " 34\n", + " 4698.305397\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.600039\n", + " -1.776888\n", + " 0.949161\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 51\n", - " 3398.135365\n", + " 36\n", + " 4828.226949\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.129833\n", + " -1.423399\n", + " 1.23938\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 52\n", - " 3223.615957\n", + " 42\n", + " 3430.58866\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.615684\n", + " -0.514427\n", + " -0.429379\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 60\n", - " 3445.014718\n", + " 48\n", + " 5314.260221\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.415606\n", + " -0.716421\n", + " 1.021716\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 61\n", - " 3505.638864\n", + " 5363.205372\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.396797\n", + " -1.170907\n", + " 1.457044\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 64\n", - " 3515.905786\n", + " 4855.908314\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.434414\n", + " -1.120408\n", + " 1.09427\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 65\n", - " 4028.363185\n", + " 3413.100524\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.220711\n", + " 1.051024\n", + " -1.445145\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 67\n", - " 4159.993943\n", + " 68\n", + " 3340.219002\n", + " [{'index': 3, 'value': 1.0}]\n", + " -1.484026\n", + " -0.009443\n", + " -1.009817\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 83\n", - " 3348.16883\n", + " 70\n", + " 4228.73157\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.638141\n", + " 1.404513\n", + " 0.296168\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 85\n", - " 3485.050273\n", + " 72\n", + " 3811.538478\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.829387\n", + " 0.142052\n", + " -0.719598\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 93\n", - " 4172.874548\n", + " 74\n", + " 4659.770763\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.242683\n", + " -1.524396\n", + " 0.586387\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 104\n", - " 3299.302424\n", + " 77\n", + " 3453.388804\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.277136\n", + " -0.211437\n", + " -0.647043\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 105\n", - " 3515.687917\n", + " 81\n", + " 4766.245033\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.208715\n", + " -1.221405\n", + " 0.804051\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", + " \n", + " \n", + " 91\n", + " 4057.807281\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.261976\n", + " 0.647036\n", + " 0.005949\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 108\n", - " 3405.224618\n", + " 96\n", + " 4739.827445\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.246331\n", + " -1.322402\n", + " 0.731497\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 113\n", - " 4209.140425\n", + " 105\n", + " 3394.891976\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.803766\n", + " 0.445043\n", + " -1.009817\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 130\n", - " 4197.905737\n", + " 111\n", + " 3201.493683\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.164286\n", + " 0.697535\n", + " -2.098138\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\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", - "3 3394.118128\n", - "8 4048.685642\n", - "17 3976.454093\n", - "23 3541.582194\n", - "25 4032.844186\n", - "27 4118.351772\n", - "29 4087.767826\n", - "34 3183.755249\n", - "35 3418.802274\n", - "39 3519.186468\n", - "51 3398.135365\n", - "52 3223.615957\n", - "60 3445.014718\n", - "61 3505.638864\n", - "64 3515.905786\n", - "65 4028.363185\n", - "67 4159.993943\n", - "83 3348.16883\n", - "85 3485.050273\n", - "93 4172.874548\n", - "104 3299.302424\n", - "105 3515.687917\n", - "108 3405.224618\n", - "113 4209.140425\n", - "130 4197.905737\n", - "...\n", + "

25 rows × 7 columns

\n", + "[67 rows x 7 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g onehotencoded_island \\\n", + "penguin_id \n", + "1 3781.402407 [{'index': 3, 'value': 1.0}] \n", + "4 4124.107944 [{'index': 1, 'value': 1.0}] \n", + "8 4670.344196 [{'index': 1, 'value': 1.0}] \n", + "11 3529.417214 [{'index': 2, 'value': 1.0}] \n", + "13 4014.101714 [{'index': 1, 'value': 1.0}] \n", + "15 5212.41288 [{'index': 1, 'value': 1.0}] \n", + "16 4163.595615 [{'index': 3, 'value': 1.0}] \n", + "23 3392.453069 [{'index': 2, 'value': 1.0}] \n", + "34 4698.305397 [{'index': 1, 'value': 1.0}] \n", + "36 4828.226949 [{'index': 1, 'value': 1.0}] \n", + "42 3430.58866 [{'index': 1, 'value': 1.0}] \n", + "48 5314.260221 [{'index': 1, 'value': 1.0}] \n", + "61 5363.205372 [{'index': 1, 'value': 1.0}] \n", + "64 4855.908314 [{'index': 1, 'value': 1.0}] \n", + "65 3413.100524 [{'index': 2, 'value': 1.0}] \n", + "68 3340.219002 [{'index': 3, 'value': 1.0}] \n", + "70 4228.73157 [{'index': 2, 'value': 1.0}] \n", + "72 3811.538478 [{'index': 2, 'value': 1.0}] \n", + "74 4659.770763 [{'index': 1, 'value': 1.0}] \n", + "77 3453.388804 [{'index': 2, 'value': 1.0}] \n", + "81 4766.245033 [{'index': 1, 'value': 1.0}] \n", + "91 4057.807281 [{'index': 2, 'value': 1.0}] \n", + "96 4739.827445 [{'index': 1, 'value': 1.0}] \n", + "105 3394.891976 [{'index': 1, 'value': 1.0}] \n", + "111 3201.493683 [{'index': 1, 'value': 1.0}] \n", "\n", - "[67 rows x 1 columns]" + " standard_scaled_culmen_length_mm standard_scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "1 -0.938587 0.748033 \n", + "4 -0.16745 0.899528 \n", + "8 0.453222 -1.877885 \n", + "11 -1.12667 0.697535 \n", + "13 -1.183094 1.404513 \n", + "15 0.867003 -0.766919 \n", + "16 -1.784958 1.959995 \n", + "23 -0.355532 0.647036 \n", + "34 -0.600039 -1.776888 \n", + "36 -0.129833 -1.423399 \n", + "42 -1.615684 -0.514427 \n", + "48 0.415606 -0.716421 \n", + "61 0.396797 -1.170907 \n", + "64 0.434414 -1.120408 \n", + "65 -1.220711 1.051024 \n", + "68 -1.484026 -0.009443 \n", + "70 1.638141 1.404513 \n", + "72 0.829387 0.142052 \n", + "74 -0.242683 -1.524396 \n", + "77 -1.277136 -0.211437 \n", + "81 0.208715 -1.221405 \n", + "91 1.261976 0.647036 \n", + "96 0.246331 -1.322402 \n", + "105 -1.803766 0.445043 \n", + "111 -1.164286 0.697535 \n", + "\n", + " standard_scaled_flipper_length_mm onehotencoded_sex \\\n", + "penguin_id \n", + "1 -1.445145 [{'index': 2, 'value': 1.0}] \n", + "4 -0.284269 [{'index': 2, 'value': 1.0}] \n", + "8 0.658942 [{'index': 1, 'value': 1.0}] \n", + "11 -0.792152 [{'index': 1, 'value': 1.0}] \n", + "13 -0.792152 [{'index': 2, 'value': 1.0}] \n", + "15 0.513833 [{'index': 2, 'value': 1.0}] \n", + "16 -0.211715 [{'index': 2, 'value': 1.0}] \n", + "23 -1.5177 [{'index': 1, 'value': 1.0}] \n", + "34 0.949161 [{'index': 1, 'value': 1.0}] \n", + "36 1.23938 [{'index': 1, 'value': 1.0}] \n", + "42 -0.429379 [{'index': 1, 'value': 1.0}] \n", + "48 1.021716 [{'index': 2, 'value': 1.0}] \n", + "61 1.457044 [{'index': 2, 'value': 1.0}] \n", + "64 1.09427 [{'index': 1, 'value': 1.0}] \n", + "65 -1.445145 [{'index': 1, 'value': 1.0}] \n", + "68 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "70 0.296168 [{'index': 2, 'value': 1.0}] \n", + "72 -0.719598 [{'index': 2, 'value': 1.0}] \n", + "74 0.586387 [{'index': 1, 'value': 1.0}] \n", + "77 -0.647043 [{'index': 1, 'value': 1.0}] \n", + "81 0.804051 [{'index': 1, 'value': 1.0}] \n", + "91 0.005949 [{'index': 2, 'value': 1.0}] \n", + "96 0.731497 [{'index': 1, 'value': 1.0}] \n", + "105 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "111 -2.098138 [{'index': 1, 'value': 1.0}] \n", + "\n", + " onehotencoded_species \n", + "penguin_id \n", + "1 [{'index': 1, 'value': 1.0}] \n", + "4 [{'index': 1, 'value': 1.0}] \n", + "8 [{'index': 3, 'value': 1.0}] \n", + "11 [{'index': 1, 'value': 1.0}] \n", + "13 [{'index': 1, 'value': 1.0}] \n", + "15 [{'index': 3, 'value': 1.0}] \n", + "16 [{'index': 1, 'value': 1.0}] \n", + "23 [{'index': 1, 'value': 1.0}] \n", + "34 [{'index': 3, 'value': 1.0}] \n", + "36 [{'index': 3, 'value': 1.0}] \n", + "42 [{'index': 1, 'value': 1.0}] \n", + "48 [{'index': 3, 'value': 1.0}] \n", + "61 [{'index': 3, 'value': 1.0}] \n", + "64 [{'index': 3, 'value': 1.0}] \n", + "65 [{'index': 1, 'value': 1.0}] \n", + "68 [{'index': 1, 'value': 1.0}] \n", + "70 [{'index': 2, 'value': 1.0}] \n", + "72 [{'index': 2, 'value': 1.0}] \n", + "74 [{'index': 3, 'value': 1.0}] \n", + "77 [{'index': 1, 'value': 1.0}] \n", + "81 [{'index': 3, 'value': 1.0}] \n", + "91 [{'index': 2, 'value': 1.0}] \n", + "96 [{'index': 3, 'value': 1.0}] \n", + "105 [{'index': 1, 'value': 1.0}] \n", + "111 [{'index': 1, 'value': 1.0}] \n", + "\n", + "[67 rows x 7 columns]" ] }, - "execution_count": 25, + "execution_count": 8, "metadata": {}, "output_type": "execute_result" } @@ -2423,18 +2554,16 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": 9, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "d7a16e04253a42b7a5ce247d8f63b656", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 027042f1-9a18-43d8-a378-ab9410e395b1 is DONE. 23.5 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 6f19614c-82c0-4f8b-b74b-9d91a894efdd is RUNNING. " ] }, "metadata": {}, @@ -2442,13 +2571,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "4a99ac15431e433595de1040872a4558", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 6c8484a0-a504-4e50-93d6-3d247c9ff558 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 51899e2d-f6ef-4e62-98b6-c11550f74f4b is RUNNING. " ] }, "metadata": {}, @@ -2456,13 +2583,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "90909b620e084f59b0f9da266257593f", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e81ca2de-df2e-41ec-af86-14f8dcec1b44 is DONE. 6.2 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 44d3fddc-74bc-4de0-a458-2c73b38f74fb is RUNNING. " ] }, "metadata": {}, @@ -2470,13 +2595,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "2a9c2c05041a4fb691809bab5310bb05", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 3e6d413c-f8c4-4390-95eb-3a1f5bc59aed is DONE. 536 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 33584475-f02b-4c98-9a51-e29996f4f950 is RUNNING. " ] }, "metadata": {}, @@ -2484,13 +2607,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "6b0677c228d54b409c66e5dfa98d7e00", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e448220d-0c50-45b7-bcbe-d1159b3d18ce is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job df25ba49-280e-424d-a357-dde71a9b35dd is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -2498,13 +2619,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "379ae6497fb34f969d21b2cd664e8bfa", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e167a234-828d-4f05-8654-63cf97e50ba3 is DONE. 10.2 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 6f92a04e-af7e-41d6-9303-6366c1751294 is RUNNING. " ] }, "metadata": {}, @@ -2532,152 +2651,452 @@ " \n", " \n", " CENTROID_ID\n", + " NEAREST_CENTROIDS_DISTANCE\n", + " onehotencoded_island\n", + " standard_scaled_culmen_length_mm\n", + " standard_scaled_culmen_depth_mm\n", + " standard_scaled_flipper_length_mm\n", + " onehotencoded_sex\n", + " onehotencoded_species\n", " \n", " \n", " penguin_id\n", " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", - " 3\n", - " 3\n", - " \n", - " \n", - " 8\n", + " 1\n", " 3\n", + " [{'CENTROID_ID': 3, 'DISTANCE': 1.236380597035...\n", + " [{'index': 3, 'value': 1.0}]\n", + " -0.938587\n", + " 0.748033\n", + " -1.445145\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 17\n", + " 4\n", " 3\n", + " [{'CENTROID_ID': 3, 'DISTANCE': 1.039497631856...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.16745\n", + " 0.899528\n", + " -0.284269\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 23\n", + " 8\n", " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 1.171040485975...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.453222\n", + " -1.877885\n", + " 0.658942\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 25\n", - " 3\n", + " 11\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 0.969102754012...\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.12667\n", + " 0.697535\n", + " -0.792152\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 27\n", + " 13\n", " 3\n", + " [{'CENTROID_ID': 3, 'DISTANCE': 1.113138945949...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.183094\n", + " 1.404513\n", + " -0.792152\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 29\n", - " 3\n", + " 15\n", + " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 1.070996026772...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.867003\n", + " -0.766919\n", + " 0.513833\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 34\n", + " 16\n", " 3\n", + " [{'CENTROID_ID': 3, 'DISTANCE': 1.780136190720...\n", + " [{'index': 3, 'value': 1.0}]\n", + " -1.784958\n", + " 1.959995\n", + " -0.211715\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 35\n", - " 1\n", + " 23\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 1.382540667483...\n", + " [{'index': 2, 'value': 1.0}]\n", + " -0.355532\n", + " 0.647036\n", + " -1.5177\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 39\n", - " 3\n", + " 34\n", + " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 1.598627908302...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.600039\n", + " -1.776888\n", + " 0.949161\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 51\n", + " 36\n", " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 1.095162305190...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.129833\n", + " -1.423399\n", + " 1.23938\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 52\n", - " 3\n", + " 42\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 1.275841743930...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.615684\n", + " -0.514427\n", + " -0.429379\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 60\n", - " 3\n", + " 48\n", + " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 0.882209023196...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.415606\n", + " -0.716421\n", + " 1.021716\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 61\n", - " 3\n", + " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 0.816202832282...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.396797\n", + " -1.170907\n", + " 1.457044\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 64\n", " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 0.735435721625...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.434414\n", + " -1.120408\n", + " 1.09427\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", " 65\n", - " 1\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 1.292559869148...\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.220711\n", + " 1.051024\n", + " -1.445145\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 67\n", - " 3\n", + " 68\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 0.876430138449...\n", + " [{'index': 3, 'value': 1.0}]\n", + " -1.484026\n", + " -0.009443\n", + " -1.009817\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 83\n", - " 3\n", + " 70\n", + " 4\n", + " [{'CENTROID_ID': 4, 'DISTANCE': 1.314229913955...\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.638141\n", + " 1.404513\n", + " 0.296168\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 85\n", - " 1\n", + " 72\n", + " 4\n", + " [{'CENTROID_ID': 4, 'DISTANCE': 0.938569518009...\n", + " [{'index': 2, 'value': 1.0}]\n", + " 0.829387\n", + " 0.142052\n", + " -0.719598\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 93\n", + " 74\n", " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 1.350320088546...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -0.242683\n", + " -1.524396\n", + " 0.586387\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 104\n", - " 3\n", + " 77\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 0.904806634663...\n", + " [{'index': 2, 'value': 1.0}]\n", + " -1.277136\n", + " -0.211437\n", + " -0.647043\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", - " 105\n", + " 81\n", " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 0.919082578073...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.208715\n", + " -1.221405\n", + " 0.804051\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 108\n", - " 3\n", + " 91\n", + " 4\n", + " [{'CENTROID_ID': 4, 'DISTANCE': 0.760360038086...\n", + " [{'index': 2, 'value': 1.0}]\n", + " 1.261976\n", + " 0.647036\n", + " 0.005949\n", + " [{'index': 2, 'value': 1.0}]\n", + " [{'index': 2, 'value': 1.0}]\n", " \n", " \n", - " 113\n", - " 3\n", + " 96\n", + " 1\n", + " [{'CENTROID_ID': 1, 'DISTANCE': 0.950188657227...\n", + " [{'index': 1, 'value': 1.0}]\n", + " 0.246331\n", + " -1.322402\n", + " 0.731497\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 3, 'value': 1.0}]\n", " \n", " \n", - " 130\n", - " 1\n", + " 105\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 1.101316467029...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.803766\n", + " 0.445043\n", + " -1.009817\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", + " \n", + " \n", + " 111\n", + " 2\n", + " [{'CENTROID_ID': 2, 'DISTANCE': 1.549061068385...\n", + " [{'index': 1, 'value': 1.0}]\n", + " -1.164286\n", + " 0.697535\n", + " -2.098138\n", + " [{'index': 1, 'value': 1.0}]\n", + " [{'index': 1, 'value': 1.0}]\n", " \n", " \n", "\n", - "

25 rows × 1 columns

\n", - "[67 rows x 1 columns in total]" + "

25 rows × 8 columns

\n", + "[67 rows x 8 columns in total]" ], "text/plain": [ - " CENTROID_ID\n", - "penguin_id \n", - "3 3\n", - "8 3\n", - "17 3\n", - "23 1\n", - "25 3\n", - "27 3\n", - "29 3\n", - "34 3\n", - "35 1\n", - "39 3\n", - "51 1\n", - "52 3\n", - "60 3\n", - "61 3\n", - "64 1\n", - "65 1\n", - "67 3\n", - "83 3\n", - "85 1\n", - "93 1\n", - "104 3\n", - "105 1\n", - "108 3\n", - "113 3\n", - "130 1\n", - "...\n", + " CENTROID_ID NEAREST_CENTROIDS_DISTANCE \\\n", + "penguin_id \n", + "1 3 [{'CENTROID_ID': 3, 'DISTANCE': 1.236380597035... \n", + "4 3 [{'CENTROID_ID': 3, 'DISTANCE': 1.039497631856... \n", + "8 1 [{'CENTROID_ID': 1, 'DISTANCE': 1.171040485975... \n", + "11 2 [{'CENTROID_ID': 2, 'DISTANCE': 0.969102754012... \n", + "13 3 [{'CENTROID_ID': 3, 'DISTANCE': 1.113138945949... \n", + "15 1 [{'CENTROID_ID': 1, 'DISTANCE': 1.070996026772... \n", + "16 3 [{'CENTROID_ID': 3, 'DISTANCE': 1.780136190720... \n", + "23 2 [{'CENTROID_ID': 2, 'DISTANCE': 1.382540667483... \n", + "34 1 [{'CENTROID_ID': 1, 'DISTANCE': 1.598627908302... \n", + "36 1 [{'CENTROID_ID': 1, 'DISTANCE': 1.095162305190... \n", + "42 2 [{'CENTROID_ID': 2, 'DISTANCE': 1.275841743930... \n", + "48 1 [{'CENTROID_ID': 1, 'DISTANCE': 0.882209023196... \n", + "61 1 [{'CENTROID_ID': 1, 'DISTANCE': 0.816202832282... \n", + "64 1 [{'CENTROID_ID': 1, 'DISTANCE': 0.735435721625... \n", + "65 2 [{'CENTROID_ID': 2, 'DISTANCE': 1.292559869148... \n", + "68 2 [{'CENTROID_ID': 2, 'DISTANCE': 0.876430138449... \n", + "70 4 [{'CENTROID_ID': 4, 'DISTANCE': 1.314229913955... \n", + "72 4 [{'CENTROID_ID': 4, 'DISTANCE': 0.938569518009... \n", + "74 1 [{'CENTROID_ID': 1, 'DISTANCE': 1.350320088546... \n", + "77 2 [{'CENTROID_ID': 2, 'DISTANCE': 0.904806634663... \n", + "81 1 [{'CENTROID_ID': 1, 'DISTANCE': 0.919082578073... \n", + "91 4 [{'CENTROID_ID': 4, 'DISTANCE': 0.760360038086... \n", + "96 1 [{'CENTROID_ID': 1, 'DISTANCE': 0.950188657227... \n", + "105 2 [{'CENTROID_ID': 2, 'DISTANCE': 1.101316467029... \n", + "111 2 [{'CENTROID_ID': 2, 'DISTANCE': 1.549061068385... \n", "\n", - "[67 rows x 1 columns]" + " onehotencoded_island standard_scaled_culmen_length_mm \\\n", + "penguin_id \n", + "1 [{'index': 3, 'value': 1.0}] -0.938587 \n", + "4 [{'index': 1, 'value': 1.0}] -0.16745 \n", + "8 [{'index': 1, 'value': 1.0}] 0.453222 \n", + "11 [{'index': 2, 'value': 1.0}] -1.12667 \n", + "13 [{'index': 1, 'value': 1.0}] -1.183094 \n", + "15 [{'index': 1, 'value': 1.0}] 0.867003 \n", + "16 [{'index': 3, 'value': 1.0}] -1.784958 \n", + "23 [{'index': 2, 'value': 1.0}] -0.355532 \n", + "34 [{'index': 1, 'value': 1.0}] -0.600039 \n", + "36 [{'index': 1, 'value': 1.0}] -0.129833 \n", + "42 [{'index': 1, 'value': 1.0}] -1.615684 \n", + "48 [{'index': 1, 'value': 1.0}] 0.415606 \n", + "61 [{'index': 1, 'value': 1.0}] 0.396797 \n", + "64 [{'index': 1, 'value': 1.0}] 0.434414 \n", + "65 [{'index': 2, 'value': 1.0}] -1.220711 \n", + "68 [{'index': 3, 'value': 1.0}] -1.484026 \n", + "70 [{'index': 2, 'value': 1.0}] 1.638141 \n", + "72 [{'index': 2, 'value': 1.0}] 0.829387 \n", + "74 [{'index': 1, 'value': 1.0}] -0.242683 \n", + "77 [{'index': 2, 'value': 1.0}] -1.277136 \n", + "81 [{'index': 1, 'value': 1.0}] 0.208715 \n", + "91 [{'index': 2, 'value': 1.0}] 1.261976 \n", + "96 [{'index': 1, 'value': 1.0}] 0.246331 \n", + "105 [{'index': 1, 'value': 1.0}] -1.803766 \n", + "111 [{'index': 1, 'value': 1.0}] -1.164286 \n", + "\n", + " standard_scaled_culmen_depth_mm \\\n", + "penguin_id \n", + "1 0.748033 \n", + "4 0.899528 \n", + "8 -1.877885 \n", + "11 0.697535 \n", + "13 1.404513 \n", + "15 -0.766919 \n", + "16 1.959995 \n", + "23 0.647036 \n", + "34 -1.776888 \n", + "36 -1.423399 \n", + "42 -0.514427 \n", + "48 -0.716421 \n", + "61 -1.170907 \n", + "64 -1.120408 \n", + "65 1.051024 \n", + "68 -0.009443 \n", + "70 1.404513 \n", + "72 0.142052 \n", + "74 -1.524396 \n", + "77 -0.211437 \n", + "81 -1.221405 \n", + "91 0.647036 \n", + "96 -1.322402 \n", + "105 0.445043 \n", + "111 0.697535 \n", + "\n", + " standard_scaled_flipper_length_mm onehotencoded_sex \\\n", + "penguin_id \n", + "1 -1.445145 [{'index': 2, 'value': 1.0}] \n", + "4 -0.284269 [{'index': 2, 'value': 1.0}] \n", + "8 0.658942 [{'index': 1, 'value': 1.0}] \n", + "11 -0.792152 [{'index': 1, 'value': 1.0}] \n", + "13 -0.792152 [{'index': 2, 'value': 1.0}] \n", + "15 0.513833 [{'index': 2, 'value': 1.0}] \n", + "16 -0.211715 [{'index': 2, 'value': 1.0}] \n", + "23 -1.5177 [{'index': 1, 'value': 1.0}] \n", + "34 0.949161 [{'index': 1, 'value': 1.0}] \n", + "36 1.23938 [{'index': 1, 'value': 1.0}] \n", + "42 -0.429379 [{'index': 1, 'value': 1.0}] \n", + "48 1.021716 [{'index': 2, 'value': 1.0}] \n", + "61 1.457044 [{'index': 2, 'value': 1.0}] \n", + "64 1.09427 [{'index': 1, 'value': 1.0}] \n", + "65 -1.445145 [{'index': 1, 'value': 1.0}] \n", + "68 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "70 0.296168 [{'index': 2, 'value': 1.0}] \n", + "72 -0.719598 [{'index': 2, 'value': 1.0}] \n", + "74 0.586387 [{'index': 1, 'value': 1.0}] \n", + "77 -0.647043 [{'index': 1, 'value': 1.0}] \n", + "81 0.804051 [{'index': 1, 'value': 1.0}] \n", + "91 0.005949 [{'index': 2, 'value': 1.0}] \n", + "96 0.731497 [{'index': 1, 'value': 1.0}] \n", + "105 -1.009817 [{'index': 1, 'value': 1.0}] \n", + "111 -2.098138 [{'index': 1, 'value': 1.0}] \n", + "\n", + " onehotencoded_species \n", + "penguin_id \n", + "1 [{'index': 1, 'value': 1.0}] \n", + "4 [{'index': 1, 'value': 1.0}] \n", + "8 [{'index': 3, 'value': 1.0}] \n", + "11 [{'index': 1, 'value': 1.0}] \n", + "13 [{'index': 1, 'value': 1.0}] \n", + "15 [{'index': 3, 'value': 1.0}] \n", + "16 [{'index': 1, 'value': 1.0}] \n", + "23 [{'index': 1, 'value': 1.0}] \n", + "34 [{'index': 3, 'value': 1.0}] \n", + "36 [{'index': 3, 'value': 1.0}] \n", + "42 [{'index': 1, 'value': 1.0}] \n", + "48 [{'index': 3, 'value': 1.0}] \n", + "61 [{'index': 3, 'value': 1.0}] \n", + "64 [{'index': 3, 'value': 1.0}] \n", + "65 [{'index': 1, 'value': 1.0}] \n", + "68 [{'index': 1, 'value': 1.0}] \n", + "70 [{'index': 2, 'value': 1.0}] \n", + "72 [{'index': 2, 'value': 1.0}] \n", + "74 [{'index': 3, 'value': 1.0}] \n", + "77 [{'index': 1, 'value': 1.0}] \n", + "81 [{'index': 3, 'value': 1.0}] \n", + "91 [{'index': 2, 'value': 1.0}] \n", + "96 [{'index': 3, 'value': 1.0}] \n", + "105 [{'index': 1, 'value': 1.0}] \n", + "111 [{'index': 1, 'value': 1.0}] \n", + "\n", + "[67 rows x 8 columns]" ] }, - "execution_count": 26, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -2704,7 +3123,7 @@ }, { "cell_type": "code", - "execution_count": 27, + "execution_count": 10, "metadata": {}, "outputs": [ { @@ -2721,7 +3140,7 @@ " ('linreg', LinearRegression())])" ] }, - "execution_count": 27, + "execution_count": 10, "metadata": {}, "output_type": "execute_result" } @@ -2748,18 +3167,16 @@ }, { "cell_type": "code", - "execution_count": 28, + "execution_count": 11, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "887bf58cebf14bdba95db828390fd33d", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job b11be0d8-e6f1-41cb-8cb2-25a38e7ef311 is DONE. 24.7 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job ed42cbb3-3d25-47ca-96c5-71a84e426a8c is RUNNING. " ] }, "metadata": {}, @@ -2767,13 +3184,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "24357055792a4eaaa60997fea0f76921", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job f32ea25c-be39-4726-a8f5-604ae83849a6 is DONE. 8.5 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 3fc74930-03b9-4a49-8ed3-c3edc4dd6e51 is RUNNING. " ] }, "metadata": {}, @@ -2781,13 +3196,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "bba878d6d3e345f1a29aea50f7101e8f", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 86e29b78-76f5-4937-8bde-407b99af04a2 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 38a4ce3b-5c2a-4d44-b826-f24529d6500b is RUNNING. " ] }, "metadata": {}, @@ -2795,13 +3208,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "4bc2c53aeb7d4a8280f9fbbe373f4b55", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job ca819734-0d41-4d9e-b743-09edae8c7fee is DONE. 29.6 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job ecad776d-77c8-4d94-8186-d5571b512b62 is RUNNING. " ] }, "metadata": {}, @@ -2809,13 +3220,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "f4f695cb0a224102b6e26adeb1827981", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 49bb5bed-cc84-47e0-9a90-08ab01e00548 is DONE. 536 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job c9bfc58f-ce2c-47a9-bbc7-b10d9de9b5a6 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -2823,13 +3232,23 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "cb1df595006d485288a1060299970e5e", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 1e40a085-2289-47dd-afd8-820413186b9f is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 8fd8036e-3753-433d-975b-c7b42406f648 is RUNNING. " + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 60319296-a480-4f51-b7ad-190ac6de963a is DONE. 6.2 kB processed. Open Job" + ], + "text/plain": [ + "" ] }, "metadata": {}, @@ -2857,152 +3276,369 @@ " \n", " \n", " predicted_body_mass_g\n", + " island\n", + " culmen_length_mm\n", + " culmen_depth_mm\n", + " flipper_length_mm\n", + " sex\n", + " species\n", " \n", " \n", " penguin_id\n", " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", - " 3\n", - " 3394.116212\n", - " \n", - " \n", - " 8\n", - " 4048.683645\n", + " 1\n", + " 3781.396682\n", + " Torgersen\n", + " 39.1\n", + " 18.7\n", + " 181.0\n", + " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 17\n", - " 3976.452358\n", + " 4\n", + " 4124.102574\n", + " Biscoe\n", + " 43.2\n", + " 19.0\n", + " 197.0\n", + " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 23\n", - " 3541.580346\n", + " 8\n", + " 4670.338389\n", + " Biscoe\n", + " 46.5\n", + " 13.5\n", + " 210.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 25\n", - " 4032.842027\n", + " 11\n", + " 3529.411644\n", + " Dream\n", + " 38.1\n", + " 18.6\n", + " 190.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 27\n", - " 4118.34983\n", + " 13\n", + " 4014.09632\n", + " Biscoe\n", + " 37.8\n", + " 20.0\n", + " 190.0\n", + " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 29\n", - " 4087.765797\n", + " 15\n", + " 5212.407319\n", + " Biscoe\n", + " 48.7\n", + " 15.7\n", + " 208.0\n", + " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 34\n", - " 3183.75379\n", + " 16\n", + " 4163.590502\n", + " Torgersen\n", + " 34.6\n", + " 21.1\n", + " 198.0\n", + " MALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 35\n", - " 3418.800633\n", + " 23\n", + " 3392.44731\n", + " Dream\n", + " 42.2\n", + " 18.5\n", + " 180.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 39\n", - " 3519.18471\n", + " 34\n", + " 4698.299674\n", + " Biscoe\n", + " 40.9\n", + " 13.7\n", + " 214.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 51\n", - " 3398.133564\n", + " 36\n", + " 4828.221398\n", + " Biscoe\n", + " 43.4\n", + " 14.4\n", + " 218.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 52\n", - " 3223.614107\n", + " 42\n", + " 3430.582874\n", + " Biscoe\n", + " 35.5\n", + " 16.2\n", + " 195.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 60\n", - " 3445.012713\n", + " 48\n", + " 5314.254798\n", + " Biscoe\n", + " 46.3\n", + " 15.8\n", + " 215.0\n", + " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", " 61\n", - " 3505.637004\n", + " 5363.19995\n", + " Biscoe\n", + " 46.2\n", + " 14.9\n", + " 221.0\n", + " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", " 64\n", - " 3515.903779\n", + " 4855.90281\n", + " Biscoe\n", + " 46.4\n", + " 15.0\n", + " 216.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", " 65\n", - " 4028.361259\n", + " 3413.094869\n", + " Dream\n", + " 37.6\n", + " 19.3\n", + " 181.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 67\n", - " 4159.991956\n", + " 68\n", + " 3340.213193\n", + " Torgersen\n", + " 36.2\n", + " 17.2\n", + " 187.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 83\n", - " 3348.167212\n", + " 70\n", + " 4228.726508\n", + " Dream\n", + " 52.8\n", + " 20.0\n", + " 205.0\n", + " MALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", " \n", " \n", - " 85\n", - " 3485.048557\n", + " 72\n", + " 3811.532821\n", + " Dream\n", + " 48.5\n", + " 17.5\n", + " 191.0\n", + " MALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", " \n", " \n", - " 93\n", - " 4172.872284\n", + " 74\n", + " 4659.765013\n", + " Biscoe\n", + " 42.8\n", + " 14.2\n", + " 209.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 104\n", - " 3299.300454\n", + " 77\n", + " 3453.383042\n", + " Dream\n", + " 37.3\n", + " 16.8\n", + " 192.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 105\n", - " 3515.68617\n", + " 81\n", + " 4766.239424\n", + " Biscoe\n", + " 45.2\n", + " 14.8\n", + " 212.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 108\n", - " 3405.222757\n", + " 91\n", + " 4057.801947\n", + " Dream\n", + " 50.8\n", + " 18.5\n", + " 201.0\n", + " MALE\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", + " \n", + " \n", + " 96\n", + " 4739.821792\n", + " Biscoe\n", + " 45.4\n", + " 14.6\n", + " 211.0\n", + " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", " \n", " \n", - " 113\n", - " 4209.13832\n", + " 105\n", + " 3394.886275\n", + " Biscoe\n", + " 34.5\n", + " 18.1\n", + " 187.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", " \n", " \n", - " 130\n", - " 4197.90382\n", + " 111\n", + " 3201.48777\n", + " Biscoe\n", + " 37.9\n", + " 18.6\n", + " 172.0\n", + " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\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", - "3 3394.116212\n", - "8 4048.683645\n", - "17 3976.452358\n", - "23 3541.580346\n", - "25 4032.842027\n", - "27 4118.34983\n", - "29 4087.765797\n", - "34 3183.75379\n", - "35 3418.800633\n", - "39 3519.18471\n", - "51 3398.133564\n", - "52 3223.614107\n", - "60 3445.012713\n", - "61 3505.637004\n", - "64 3515.903779\n", - "65 4028.361259\n", - "67 4159.991956\n", - "83 3348.167212\n", - "85 3485.048557\n", - "93 4172.872284\n", - "104 3299.300454\n", - "105 3515.68617\n", - "108 3405.222757\n", - "113 4209.13832\n", - "130 4197.90382\n", - "...\n", + "

25 rows × 7 columns

\n", + "[67 rows x 7 columns in total]" + ], + "text/plain": [ + " predicted_body_mass_g island culmen_length_mm \\\n", + "penguin_id \n", + "1 3781.396682 Torgersen 39.1 \n", + "4 4124.102574 Biscoe 43.2 \n", + "8 4670.338389 Biscoe 46.5 \n", + "11 3529.411644 Dream 38.1 \n", + "13 4014.09632 Biscoe 37.8 \n", + "15 5212.407319 Biscoe 48.7 \n", + "16 4163.590502 Torgersen 34.6 \n", + "23 3392.44731 Dream 42.2 \n", + "34 4698.299674 Biscoe 40.9 \n", + "36 4828.221398 Biscoe 43.4 \n", + "42 3430.582874 Biscoe 35.5 \n", + "48 5314.254798 Biscoe 46.3 \n", + "61 5363.19995 Biscoe 46.2 \n", + "64 4855.90281 Biscoe 46.4 \n", + "65 3413.094869 Dream 37.6 \n", + "68 3340.213193 Torgersen 36.2 \n", + "70 4228.726508 Dream 52.8 \n", + "72 3811.532821 Dream 48.5 \n", + "74 4659.765013 Biscoe 42.8 \n", + "77 3453.383042 Dream 37.3 \n", + "81 4766.239424 Biscoe 45.2 \n", + "91 4057.801947 Dream 50.8 \n", + "96 4739.821792 Biscoe 45.4 \n", + "105 3394.886275 Biscoe 34.5 \n", + "111 3201.48777 Biscoe 37.9 \n", "\n", - "[67 rows x 1 columns]" + " culmen_depth_mm flipper_length_mm sex \\\n", + "penguin_id \n", + "1 18.7 181.0 MALE \n", + "4 19.0 197.0 MALE \n", + "8 13.5 210.0 FEMALE \n", + "11 18.6 190.0 FEMALE \n", + "13 20.0 190.0 MALE \n", + "15 15.7 208.0 MALE \n", + "16 21.1 198.0 MALE \n", + "23 18.5 180.0 FEMALE \n", + "34 13.7 214.0 FEMALE \n", + "36 14.4 218.0 FEMALE \n", + "42 16.2 195.0 FEMALE \n", + "48 15.8 215.0 MALE \n", + "61 14.9 221.0 MALE \n", + "64 15.0 216.0 FEMALE \n", + "65 19.3 181.0 FEMALE \n", + "68 17.2 187.0 FEMALE \n", + "70 20.0 205.0 MALE \n", + "72 17.5 191.0 MALE \n", + "74 14.2 209.0 FEMALE \n", + "77 16.8 192.0 FEMALE \n", + "81 14.8 212.0 FEMALE \n", + "91 18.5 201.0 MALE \n", + "96 14.6 211.0 FEMALE \n", + "105 18.1 187.0 FEMALE \n", + "111 18.6 172.0 FEMALE \n", + "\n", + " species \n", + "penguin_id \n", + "1 Adelie Penguin (Pygoscelis adeliae) \n", + "4 Adelie Penguin (Pygoscelis adeliae) \n", + "8 Gentoo penguin (Pygoscelis papua) \n", + "11 Adelie Penguin (Pygoscelis adeliae) \n", + "13 Adelie Penguin (Pygoscelis adeliae) \n", + "15 Gentoo penguin (Pygoscelis papua) \n", + "16 Adelie Penguin (Pygoscelis adeliae) \n", + "23 Adelie Penguin (Pygoscelis adeliae) \n", + "34 Gentoo penguin (Pygoscelis papua) \n", + "36 Gentoo penguin (Pygoscelis papua) \n", + "42 Adelie Penguin (Pygoscelis adeliae) \n", + "48 Gentoo penguin (Pygoscelis papua) \n", + "61 Gentoo penguin (Pygoscelis papua) \n", + "64 Gentoo penguin (Pygoscelis papua) \n", + "65 Adelie Penguin (Pygoscelis adeliae) \n", + "68 Adelie Penguin (Pygoscelis adeliae) \n", + "70 Chinstrap penguin (Pygoscelis antarctica) \n", + "72 Chinstrap penguin (Pygoscelis antarctica) \n", + "74 Gentoo penguin (Pygoscelis papua) \n", + "77 Adelie Penguin (Pygoscelis adeliae) \n", + "81 Gentoo penguin (Pygoscelis papua) \n", + "91 Chinstrap penguin (Pygoscelis antarctica) \n", + "96 Gentoo penguin (Pygoscelis papua) \n", + "105 Adelie Penguin (Pygoscelis adeliae) \n", + "111 Adelie Penguin (Pygoscelis adeliae) \n", + "\n", + "[67 rows x 7 columns]" ] }, - "execution_count": 28, + "execution_count": 11, "metadata": {}, "output_type": "execute_result" } @@ -3034,60 +3670,16 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": 12, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "2d32081be31f44abb8de67e2209d76cd", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job 2a043039-670f-4eb8-9cf0-765ee6ed7de6 is RUNNING. Open Job" + ], "text/plain": [ - "HTML(value='Query job bc8b2042-1e13-441c-9531-300ed5badb7a is RUNNING. " ] }, "metadata": {}, @@ -3095,13 +3687,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "4588ae10de634460bf4026ddd9076351", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 7f1f565b-0f73-4a4e-b33f-8484fa260838 is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 5e867182-dd7a-4aff-87a8-f7596e900fd5 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -3109,13 +3699,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "8209cf8286a545ebb7b6ef9d002a43a1", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job d4b9d4a6-d75e-46e1-b092-ab58e8aef890 is DONE. 48 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job d4cdb016-8f1e-4960-8ed7-4524ccc5a8a8 is RUNNING. " ] }, "metadata": {}, @@ -3153,12 +3741,12 @@ " \n", " \n", " 0\n", - " 229.48269\n", - " 82962.794947\n", - " 0.004248\n", - " 206.728384\n", - " 0.88633\n", - " 0.892953\n", + " 216.444357\n", + " 72639.698707\n", + " 0.00463\n", + " 170.588356\n", + " 0.896396\n", + " 0.900547\n", " \n", " \n", "\n", @@ -3167,15 +3755,15 @@ ], "text/plain": [ " mean_absolute_error mean_squared_error mean_squared_log_error \\\n", - "0 229.48269 82962.794947 0.004248 \n", + "0 216.444357 72639.698707 0.00463 \n", "\n", " median_absolute_error r2_score explained_variance \n", - "0 206.728384 0.88633 0.892953 \n", + "0 170.588356 0.896396 0.900547 \n", "\n", "[1 rows x 6 columns]" ] }, - "execution_count": 29, + "execution_count": 12, "metadata": {}, "output_type": "execute_result" } @@ -3195,18 +3783,16 @@ }, { "cell_type": "code", - "execution_count": 30, + "execution_count": 14, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "f32692d89f00406499f4ea5aa55268fb", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 73448ee8-698b-435f-b11e-6fe2de3bcd8d is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job e57383ef-f043-458b-96c6-893e7c5b0c00 is RUNNING. " ] }, "metadata": {}, @@ -3214,13 +3800,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "72e5f23a99de4a818c8493b8b4f3854d", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job e002f59d-a03c-4ec9-a85a-93adbfd7bd17 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 1a9db485-477b-43e2-94eb-dea7dc21d45d is RUNNING. " ] }, "metadata": {}, @@ -3228,13 +3812,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "9d5333a91b504dd9be51c997715530ab", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 4ab1febc-fb55-473a-b295-69e4329cc5f0 is DONE. 30.0 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 4570a563-b8e0-4308-b8cb-c4731491d4f7 is RUNNING. " ] }, "metadata": {}, @@ -3243,10 +3825,10 @@ { "data": { "text/plain": [ - "0.8863300923278365" + "0.8963962044533755" ] }, - "execution_count": 30, + "execution_count": 14, "metadata": {}, "output_type": "execute_result" } @@ -3254,7 +3836,7 @@ "source": [ "from bigframes.ml.metrics import r2_score\n", "\n", - "r2_score(y_test, predicted_y_test)" + "r2_score(y_test, predicted_y_test[\"predicted_body_mass_g\"])" ] }, { @@ -3274,57 +3856,9 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 15, "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "fbc4a70f31d4465b974a7f7c9cc97731", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Copy job c2413be4-6972-4e36-8234-5063628b6d71 is RUNNING. Open Job" + ], "text/plain": [ - "HTML(value='Query job 31a5b656-000e-4238-9fd9-c6e644ca298f is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -67,13 +37,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "0f25faa156584cc59dda9b0e60f72534", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 12e0f983-695e-4903-8ff1-2f353d7e8cba is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job d8eed0ca-7ce9-4ed8-a592-e16af9f9db8d is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -113,250 +81,250 @@ " \n", " 0\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 36.6\n", - " 18.4\n", - " 184.0\n", - " 3475.0\n", - " FEMALE\n", + " Biscoe\n", + " 40.1\n", + " 18.9\n", + " 188.0\n", + " 4300.0\n", + " MALE\n", " \n", " \n", " 1\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 39.8\n", - " 19.1\n", - " 184.0\n", - " 4650.0\n", + " Torgersen\n", + " 39.1\n", + " 18.7\n", + " 181.0\n", + " 3750.0\n", " MALE\n", " \n", " \n", " 2\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.9\n", - " 18.9\n", - " 184.0\n", - " 3900.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 47.4\n", + " 14.6\n", + " 212.0\n", + " 4725.0\n", + " FEMALE\n", " \n", " \n", " 3\n", " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 46.5\n", - " 17.9\n", - " 192.0\n", - " 3500.0\n", + " 42.5\n", + " 16.7\n", + " 187.0\n", + " 3350.0\n", " FEMALE\n", " \n", " \n", " 4\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.3\n", - " 16.8\n", - " 192.0\n", - " 3000.0\n", - " FEMALE\n", + " Biscoe\n", + " 43.2\n", + " 19.0\n", + " 197.0\n", + " 4775.0\n", + " MALE\n", " \n", " \n", " 5\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 43.2\n", - " 18.5\n", - " 192.0\n", - " 4100.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.7\n", + " 15.3\n", + " 219.0\n", + " 5200.0\n", " MALE\n", " \n", " \n", " 6\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 46.9\n", - " 16.6\n", - " 192.0\n", - " 2700.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 41.3\n", + " 21.1\n", + " 195.0\n", + " 4400.0\n", + " MALE\n", " \n", " \n", " 7\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 50.5\n", - " 18.4\n", - " 200.0\n", - " 3400.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 45.2\n", + " 13.8\n", + " 215.0\n", + " 4750.0\n", " FEMALE\n", " \n", " \n", " 8\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 49.5\n", - " 19.0\n", - " 200.0\n", - " 3800.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.5\n", + " 13.5\n", + " 210.0\n", + " 4550.0\n", + " FEMALE\n", " \n", " \n", " 9\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.2\n", - " 20.1\n", - " 200.0\n", - " 3975.0\n", - " MALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 50.5\n", + " 15.2\n", + " 216.0\n", + " 5000.0\n", + " FEMALE\n", " \n", " \n", " 10\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.8\n", - " 18.9\n", - " 208.0\n", - " 4300.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 48.2\n", + " 15.6\n", + " 221.0\n", + " 5100.0\n", " MALE\n", " \n", " \n", " 11\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 39.0\n", - " 18.7\n", - " 185.0\n", - " 3650.0\n", - " MALE\n", + " 38.1\n", + " 18.6\n", + " 190.0\n", + " 3700.0\n", + " FEMALE\n", " \n", " \n", " 12\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.0\n", - " 16.9\n", - " 185.0\n", - " 3000.0\n", - " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 50.7\n", + " 15.0\n", + " 223.0\n", + " 5550.0\n", + " MALE\n", " \n", " \n", " 13\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 47.0\n", - " 17.3\n", - " 185.0\n", - " 3700.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 37.8\n", + " 20.0\n", + " 190.0\n", + " 4250.0\n", + " MALE\n", " \n", " \n", " 14\n", " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 34.0\n", - " 17.1\n", - " 185.0\n", - " 3400.0\n", + " Biscoe\n", + " 35.0\n", + " 17.9\n", + " 190.0\n", + " 3450.0\n", " FEMALE\n", " \n", " \n", " 15\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 37.0\n", - " 16.5\n", - " 185.0\n", - " 3400.0\n", - " FEMALE\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 48.7\n", + " 15.7\n", + " 208.0\n", + " 5350.0\n", + " MALE\n", " \n", " \n", " 16\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 45.7\n", - " 17.3\n", - " 193.0\n", - " 3600.0\n", - " FEMALE\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Torgersen\n", + " 34.6\n", + " 21.1\n", + " 198.0\n", + " 4400.0\n", + " MALE\n", " \n", " \n", " 17\n", - " Chinstrap penguin (Pygoscelis antarctica)\n", - " Dream\n", - " 50.6\n", - " 19.4\n", - " 193.0\n", - " 3800.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 46.8\n", + " 15.4\n", + " 215.0\n", + " 5150.0\n", " MALE\n", " \n", " \n", " 18\n", - " Adelie Penguin (Pygoscelis adeliae)\n", + " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 39.7\n", - " 17.9\n", - " 193.0\n", - " 4250.0\n", + " 50.3\n", + " 20.0\n", + " 197.0\n", + " 3300.0\n", " MALE\n", " \n", " \n", " 19\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 37.8\n", + " 37.2\n", " 18.1\n", - " 193.0\n", - " 3750.0\n", + " 178.0\n", + " 3900.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", + " 51.0\n", + " 18.8\n", + " 203.0\n", + " 4100.0\n", + " MALE\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", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Biscoe\n", + " 40.5\n", + " 17.9\n", + " 187.0\n", + " 3200.0\n", + " FEMALE\n", " \n", " \n", " 22\n", - " Adelie Penguin (Pygoscelis adeliae)\n", - " Dream\n", - " 40.2\n", - " 17.1\n", - " 193.0\n", - " 3400.0\n", + " Gentoo penguin (Pygoscelis papua)\n", + " Biscoe\n", + " 45.5\n", + " 13.9\n", + " 210.0\n", + " 4200.0\n", " FEMALE\n", " \n", " \n", " 23\n", " Adelie Penguin (Pygoscelis adeliae)\n", " Dream\n", - " 36.8\n", + " 42.2\n", " 18.5\n", - " 193.0\n", - " 3500.0\n", + " 180.0\n", + " 3550.0\n", " FEMALE\n", " \n", " \n", " 24\n", " Chinstrap penguin (Pygoscelis antarctica)\n", " Dream\n", - " 49.6\n", - " 18.2\n", - " 193.0\n", + " 51.7\n", + " 20.3\n", + " 194.0\n", " 3775.0\n", " MALE\n", " \n", @@ -366,65 +334,65 @@ "[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", + " species island culmen_length_mm \\\n", + "0 Adelie Penguin (Pygoscelis adeliae) Biscoe 40.1 \n", + "1 Adelie Penguin (Pygoscelis adeliae) Torgersen 39.1 \n", + "2 Gentoo penguin (Pygoscelis papua) Biscoe 47.4 \n", + "3 Chinstrap penguin (Pygoscelis antarctica) Dream 42.5 \n", + "4 Adelie Penguin (Pygoscelis adeliae) Biscoe 43.2 \n", + "5 Gentoo penguin (Pygoscelis papua) Biscoe 46.7 \n", + "6 Adelie Penguin (Pygoscelis adeliae) Biscoe 41.3 \n", + "7 Gentoo penguin (Pygoscelis papua) Biscoe 45.2 \n", + "8 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 \n", + "9 Gentoo penguin (Pygoscelis papua) Biscoe 50.5 \n", + "10 Gentoo penguin (Pygoscelis papua) Biscoe 48.2 \n", + "11 Adelie Penguin (Pygoscelis adeliae) Dream 38.1 \n", + "12 Gentoo penguin (Pygoscelis papua) Biscoe 50.7 \n", + "13 Adelie Penguin (Pygoscelis adeliae) Biscoe 37.8 \n", + "14 Adelie Penguin (Pygoscelis adeliae) Biscoe 35.0 \n", + "15 Gentoo penguin (Pygoscelis papua) Biscoe 48.7 \n", + "16 Adelie Penguin (Pygoscelis adeliae) Torgersen 34.6 \n", + "17 Gentoo penguin (Pygoscelis papua) Biscoe 46.8 \n", + "18 Chinstrap penguin (Pygoscelis antarctica) Dream 50.3 \n", + "19 Adelie Penguin (Pygoscelis adeliae) Dream 37.2 \n", + "20 Chinstrap penguin (Pygoscelis antarctica) Dream 51.0 \n", + "21 Adelie Penguin (Pygoscelis adeliae) Biscoe 40.5 \n", + "22 Gentoo penguin (Pygoscelis papua) Biscoe 45.5 \n", + "23 Adelie Penguin (Pygoscelis adeliae) Dream 42.2 \n", + "24 Chinstrap penguin (Pygoscelis antarctica) Dream 51.7 \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", + "0 18.9 188.0 4300.0 MALE \n", + "1 18.7 181.0 3750.0 MALE \n", + "2 14.6 212.0 4725.0 FEMALE \n", + "3 16.7 187.0 3350.0 FEMALE \n", + "4 19.0 197.0 4775.0 MALE \n", + "5 15.3 219.0 5200.0 MALE \n", + "6 21.1 195.0 4400.0 MALE \n", + "7 13.8 215.0 4750.0 FEMALE \n", + "8 13.5 210.0 4550.0 FEMALE \n", + "9 15.2 216.0 5000.0 FEMALE \n", + "10 15.6 221.0 5100.0 MALE \n", + "11 18.6 190.0 3700.0 FEMALE \n", + "12 15.0 223.0 5550.0 MALE \n", + "13 20.0 190.0 4250.0 MALE \n", + "14 17.9 190.0 3450.0 FEMALE \n", + "15 15.7 208.0 5350.0 MALE \n", + "16 21.1 198.0 4400.0 MALE \n", + "17 15.4 215.0 5150.0 MALE \n", + "18 20.0 197.0 3300.0 MALE \n", + "19 18.1 178.0 3900.0 MALE \n", + "20 18.8 203.0 4100.0 MALE \n", + "21 17.9 187.0 3200.0 FEMALE \n", + "22 13.9 210.0 4200.0 FEMALE \n", + "23 18.5 180.0 3550.0 FEMALE \n", + "24 20.3 194.0 3775.0 MALE \n", "...\n", "\n", "[344 rows x 7 columns]" ] }, - "execution_count": 12, + "execution_count": 1, "metadata": {}, "output_type": "execute_result" } @@ -450,32 +418,16 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 2, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "a9ad907fa6e64a61a9dce420bc7d2beb", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job 3537a10a-641a-4d40-ae47-449c641b1bc5 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 34101409-7c65-4045-ad52-c6ba24dc9cbb is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -483,13 +435,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "514e68d5b0b4452a9ccdff947848541a", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 2af0b0d6-c11b-499e-8d25-a2c628b2853b is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 74190ac2-21a2-47b0-bc21-ef5373565f17 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -527,294 +477,294 @@ " \n", " \n", " 0\n", - " Dream\n", - " 36.6\n", - " 18.4\n", - " 184.0\n", - " 3475.0\n", - " FEMALE\n", - " \n", - " \n", - " 1\n", - " Dream\n", - " 39.8\n", - " 19.1\n", - " 184.0\n", - " 4650.0\n", + " Biscoe\n", + " 40.1\n", + " 18.9\n", + " 188.0\n", + " 4300.0\n", " MALE\n", " \n", " \n", - " 2\n", - " Dream\n", - " 40.9\n", - " 18.9\n", - " 184.0\n", - " 3900.0\n", + " 1\n", + " Torgersen\n", + " 39.1\n", + " 18.7\n", + " 181.0\n", + " 3750.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", - " 5\n", - " Dream\n", + " Biscoe\n", " 43.2\n", - " 18.5\n", - " 192.0\n", - " 4100.0\n", - " MALE\n", - " \n", - " \n", - " 9\n", - " Dream\n", - " 40.2\n", - " 20.1\n", - " 200.0\n", - " 3975.0\n", + " 19.0\n", + " 197.0\n", + " 4775.0\n", " MALE\n", " \n", " \n", - " 10\n", - " Dream\n", - " 40.8\n", - " 18.9\n", - " 208.0\n", - " 4300.0\n", + " 6\n", + " Biscoe\n", + " 41.3\n", + " 21.1\n", + " 195.0\n", + " 4400.0\n", " MALE\n", " \n", " \n", " 11\n", " Dream\n", - " 39.0\n", - " 18.7\n", - " 185.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", + " 38.1\n", + " 18.6\n", + " 190.0\n", + " 3700.0\n", " FEMALE\n", " \n", " \n", - " 14\n", - " Dream\n", - " 34.0\n", - " 17.1\n", - " 185.0\n", - " 3400.0\n", - " FEMALE\n", + " 13\n", + " Biscoe\n", + " 37.8\n", + " 20.0\n", + " 190.0\n", + " 4250.0\n", + " MALE\n", " \n", " \n", - " 15\n", - " Dream\n", - " 37.0\n", - " 16.5\n", - " 185.0\n", - " 3400.0\n", + " 14\n", + " Biscoe\n", + " 35.0\n", + " 17.9\n", + " 190.0\n", + " 3450.0\n", " FEMALE\n", " \n", " \n", - " 18\n", - " Dream\n", - " 39.7\n", - " 17.9\n", - " 193.0\n", - " 4250.0\n", + " 16\n", + " Torgersen\n", + " 34.6\n", + " 21.1\n", + " 198.0\n", + " 4400.0\n", " MALE\n", " \n", " \n", " 19\n", " Dream\n", - " 37.8\n", + " 37.2\n", " 18.1\n", - " 193.0\n", - " 3750.0\n", + " 178.0\n", + " 3900.0\n", " MALE\n", " \n", " \n", - " 22\n", - " Dream\n", - " 40.2\n", - " 17.1\n", - " 193.0\n", - " 3400.0\n", + " 21\n", + " Biscoe\n", + " 40.5\n", + " 17.9\n", + " 187.0\n", + " 3200.0\n", " FEMALE\n", " \n", " \n", " 23\n", " Dream\n", - " 36.8\n", + " 42.2\n", " 18.5\n", - " 193.0\n", - " 3500.0\n", + " 180.0\n", + " 3550.0\n", " FEMALE\n", " \n", " \n", - " 26\n", + " 30\n", " Dream\n", - " 41.5\n", - " 18.5\n", - " 201.0\n", - " 4000.0\n", + " 39.2\n", + " 21.1\n", + " 196.0\n", + " 4150.0\n", " MALE\n", " \n", " \n", - " 31\n", - " Dream\n", - " 33.1\n", - " 16.1\n", - " 178.0\n", - " 2900.0\n", - " FEMALE\n", + " 32\n", + " Torgersen\n", + " 42.9\n", + " 17.6\n", + " 196.0\n", + " 4700.0\n", + " MALE\n", " \n", " \n", - " 32\n", + " 38\n", " Dream\n", - " 37.2\n", - " 18.1\n", - " 178.0\n", + " 41.1\n", + " 17.5\n", + " 190.0\n", " 3900.0\n", " MALE\n", " \n", " \n", - " 33\n", - " Dream\n", - " 39.5\n", - " 16.7\n", - " 178.0\n", - " 3250.0\n", + " 40\n", + " Torgersen\n", + " 38.6\n", + " 21.2\n", + " 191.0\n", + " 3800.0\n", + " MALE\n", + " \n", + " \n", + " 42\n", + " Biscoe\n", + " 35.5\n", + " 16.2\n", + " 195.0\n", + " 3350.0\n", " FEMALE\n", " \n", " \n", - " 35\n", + " 44\n", " Dream\n", - " 36.0\n", - " 18.5\n", + " 39.2\n", + " 18.6\n", + " 190.0\n", + " 4250.0\n", + " MALE\n", + " \n", + " \n", + " 45\n", + " Torgersen\n", + " 35.2\n", + " 15.9\n", " 186.0\n", - " 3100.0\n", + " 3050.0\n", " FEMALE\n", " \n", " \n", - " 36\n", + " 46\n", " Dream\n", + " 43.2\n", + " 18.5\n", + " 192.0\n", + " 4100.0\n", + " MALE\n", + " \n", + " \n", + " 49\n", + " Biscoe\n", " 39.6\n", - " 18.1\n", + " 17.7\n", " 186.0\n", - " 4450.0\n", - " MALE\n", + " 3500.0\n", + " FEMALE\n", " \n", " \n", - " 38\n", - " Dream\n", - " 41.3\n", + " 53\n", + " Biscoe\n", + " 45.6\n", " 20.3\n", - " 194.0\n", - " 3550.0\n", + " 191.0\n", + " 4600.0\n", " MALE\n", " \n", " \n", - " 41\n", - " Dream\n", - " 35.7\n", - " 18.0\n", - " 202.0\n", - " 3550.0\n", + " 58\n", + " Torgersen\n", + " 40.9\n", + " 16.8\n", + " 191.0\n", + " 3700.0\n", " FEMALE\n", " \n", " \n", - " 51\n", - " Dream\n", - " 38.1\n", - " 17.6\n", - " 187.0\n", - " 3425.0\n", + " 60\n", + " Torgersen\n", + " 40.3\n", + " 18.0\n", + " 195.0\n", + " 3250.0\n", " FEMALE\n", " \n", " \n", - " 53\n", + " 62\n", " Dream\n", " 36.0\n", - " 17.1\n", - " 187.0\n", - " 3700.0\n", + " 18.5\n", + " 186.0\n", + " 3100.0\n", " FEMALE\n", " \n", + " \n", + " 63\n", + " Torgersen\n", + " 39.3\n", + " 20.6\n", + " 190.0\n", + " 3650.0\n", + " MALE\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 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", + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "0 Biscoe 40.1 18.9 188.0 \n", + "1 Torgersen 39.1 18.7 181.0 \n", + "4 Biscoe 43.2 19.0 197.0 \n", + "6 Biscoe 41.3 21.1 195.0 \n", + "11 Dream 38.1 18.6 190.0 \n", + "13 Biscoe 37.8 20.0 190.0 \n", + "14 Biscoe 35.0 17.9 190.0 \n", + "16 Torgersen 34.6 21.1 198.0 \n", + "19 Dream 37.2 18.1 178.0 \n", + "21 Biscoe 40.5 17.9 187.0 \n", + "23 Dream 42.2 18.5 180.0 \n", + "30 Dream 39.2 21.1 196.0 \n", + "32 Torgersen 42.9 17.6 196.0 \n", + "38 Dream 41.1 17.5 190.0 \n", + "40 Torgersen 38.6 21.2 191.0 \n", + "42 Biscoe 35.5 16.2 195.0 \n", + "44 Dream 39.2 18.6 190.0 \n", + "45 Torgersen 35.2 15.9 186.0 \n", + "46 Dream 43.2 18.5 192.0 \n", + "49 Biscoe 39.6 17.7 186.0 \n", + "53 Biscoe 45.6 20.3 191.0 \n", + "58 Torgersen 40.9 16.8 191.0 \n", + "60 Torgersen 40.3 18.0 195.0 \n", + "62 Dream 36.0 18.5 186.0 \n", + "63 Torgersen 39.3 20.6 190.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", + " body_mass_g sex \n", + "0 4300.0 MALE \n", + "1 3750.0 MALE \n", + "4 4775.0 MALE \n", + "6 4400.0 MALE \n", + "11 3700.0 FEMALE \n", + "13 4250.0 MALE \n", + "14 3450.0 FEMALE \n", + "16 4400.0 MALE \n", + "19 3900.0 MALE \n", + "21 3200.0 FEMALE \n", + "23 3550.0 FEMALE \n", + "30 4150.0 MALE \n", + "32 4700.0 MALE \n", + "38 3900.0 MALE \n", + "40 3800.0 MALE \n", + "42 3350.0 FEMALE \n", + "44 4250.0 MALE \n", + "45 3050.0 FEMALE \n", + "46 4100.0 MALE \n", + "49 3500.0 FEMALE \n", + "53 4600.0 MALE \n", + "58 3700.0 FEMALE \n", + "60 3250.0 FEMALE \n", + "62 3100.0 FEMALE \n", + "63 3650.0 MALE \n", "...\n", "\n", "[146 rows x 6 columns]" ] }, - "execution_count": 13, + "execution_count": 2, "metadata": {}, "output_type": "execute_result" } @@ -843,18 +793,16 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 3, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "81f9aa34c7234bd88b6b7a4bc77d4b4e", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 0808457b-a0df-4a37-b7a5-8885f4a4588c is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 288f0daa-a51e-45b4-86bf-d054467c4a99 is DONE. 28.9 kB processed. " ] }, "metadata": {}, @@ -881,7 +829,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 4, "metadata": {}, "outputs": [ { @@ -897,7 +845,7 @@ " ('linreg', LinearRegression(fit_intercept=False))])" ] }, - "execution_count": 15, + "execution_count": 4, "metadata": {}, "output_type": "execute_result" } @@ -936,9 +884,63 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 5, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "Query job e9bfa6a5-a53f-4d8b-ae8c-cc8cd55d0947 is DONE. 28.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job d8d553cf-3d36-49aa-b18b-9a05576a1fb0 is DONE. 28.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 75ef0083-9a4f-4ffb-a6c6-d82974a1659f is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "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": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "pipeline.fit(X_train, y_train)" ] @@ -953,18 +955,16 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 6, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "fcf406d36c0d4915b318cd30c0f3df25", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 55c5a9ce-8159-4a1a-99a4-af3a906640ba is DONE. 29.3 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 81196f97-304b-4d77-bb0f-8fc8adb8fe75 is RUNNING. " ] }, "metadata": {}, @@ -972,13 +972,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "41399a6b1d4f45328bacc6c868cefdf6", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 3e41c470-de70-4f13-89d9-c5564d0b2836 is DONE. 232 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job b417f27a-387d-4eb2-8d6d-287327ef0471 is DONE. 232 Bytes processed. " ] }, "metadata": {}, @@ -986,13 +984,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "e3c17676eab448c0942c0c32689ba4b5", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job ed2f9042-a737-4d13-bd21-8c3d29cd61a2 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job b7f89a61-d76a-47be-8b83-917d69f255a2 is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -1000,13 +996,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "6c903861564b412aad9d9decad26560c", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 815d16b5-0a5d-42be-a766-1cff5b8f22f2 is DONE. 28.9 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 9619c393-90b3-4fea-a197-d09389e9486c is DONE. 31.7 kB processed. " ] }, "metadata": {}, @@ -1014,13 +1008,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "2c2534cd90e64c81be45753b81b1be46", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 37a38dc6-5073-4544-a1e3-da145a843922 is DONE. 29.4 kB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job e5854451-ffb4-4a28-a25f-3bdd68e9edae is DONE. 32.2 kB processed. " ] }, "metadata": {}, @@ -1029,10 +1021,10 @@ { "data": { "text/plain": [ - "0.6757452736197735" + "0.2655729213572775" ] }, - "execution_count": 17, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -1040,9 +1032,9 @@ "source": [ "from bigframes.ml.metrics import r2_score\n", "\n", - "pred_y = pipeline.predict(X_test)\n", + "y_pred = pipeline.predict(X_test)[\"predicted_body_mass_g\"]\n", "\n", - "r2_score(y_test, pred_y)" + "r2_score(y_test, y_pred)" ] }, { @@ -1055,18 +1047,16 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 7, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "9295d6a3ff834f7a91a43d3f4ef4a61c", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Load job 7b46750c-70b4-468d-87ba-9f84f579f2a6 is DONE. Open Job" + ], "text/plain": [ - "HTML(value='Load job d4c2f933-3514-4901-bcd7-888ee66eba82 is RUNNING. " ] }, "metadata": {}, @@ -1097,32 +1087,16 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 8, "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "b7eb82b3b5fc4a8e97468070a3e76300", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "HTML(value='Query job e4ffd919-6f69-4382-a7e5-db37c7c1fefa is RUNNING. Open Job" + ], "text/plain": [ - "HTML(value='Query job 6b3e3285-79e9-4137-bf3b-7b7185ef76a5 is DONE. 24 Bytes processed. " ] }, "metadata": {}, @@ -1130,13 +1104,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "edc7bc6434bd4be4926626a235aab65a", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 207cb787-cf8a-43ea-8e73-644d3f58b11a is DONE. 24 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 173c4194-e194-43d2-8359-7bec83d3c861 is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -1144,13 +1116,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "55a8cbd9b1ab47eeab6e1c305847630f", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job c5dc5075-cac0-4947-9e9f-06aa9cc5bd2a is DONE. 0 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 53ba2332-590c-488d-9505-23aebaaad9cb is DONE. 48 Bytes processed. " ] }, "metadata": {}, @@ -1158,13 +1128,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "463a5b072148474db629b9346fa3a6d1", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 2ca4a569-7186-48ed-b3e4-004dca704798 is DONE. 282 Bytes processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 66e4a8e0-4cae-4e9d-86e0-17dc24f6cfbb is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -1192,41 +1160,83 @@ " \n", " \n", " predicted_body_mass_g\n", + " species\n", + " island\n", + " culmen_length_mm\n", + " culmen_depth_mm\n", + " flipper_length_mm\n", + " sex\n", " \n", " \n", " tag_number\n", " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", " \n", " \n", " \n", " \n", " 1633\n", - " 3965.994361\n", + " 4017.203152\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Torgersen\n", + " 39.5\n", + " 18.8\n", + " 196.0\n", + " MALE\n", " \n", " \n", " 1672\n", - " 3246.312058\n", + " 3127.601519\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Torgersen\n", + " 38.5\n", + " 17.2\n", + " 181.0\n", + " FEMALE\n", " \n", " \n", " 1690\n", - " 3456.404062\n", + " 3386.101231\n", + " Adelie Penguin (Pygoscelis adeliae)\n", + " Dream\n", + " 37.9\n", + " 18.1\n", + " 188.0\n", + " FEMALE\n", " \n", " \n", "\n", - "

3 rows × 1 columns

\n", - "[3 rows x 1 columns in total]" + "

3 rows × 7 columns

\n", + "[3 rows x 7 columns in total]" ], "text/plain": [ - " predicted_body_mass_g\n", - "tag_number \n", - "1633 3965.994361\n", - "1672 3246.312058\n", - "1690 3456.404062\n", + " predicted_body_mass_g species \\\n", + "tag_number \n", + "1633 4017.203152 Adelie Penguin (Pygoscelis adeliae) \n", + "1672 3127.601519 Adelie Penguin (Pygoscelis adeliae) \n", + "1690 3386.101231 Adelie Penguin (Pygoscelis adeliae) \n", + "\n", + " island culmen_length_mm culmen_depth_mm flipper_length_mm \\\n", + "tag_number \n", + "1633 Torgersen 39.5 18.8 196.0 \n", + "1672 Torgersen 38.5 17.2 181.0 \n", + "1690 Dream 37.9 18.1 188.0 \n", "\n", - "[3 rows x 1 columns]" + " sex \n", + "tag_number \n", + "1633 MALE \n", + "1672 FEMALE \n", + "1690 FEMALE \n", + "\n", + "[3 rows x 7 columns]" ] }, - "execution_count": 19, + "execution_count": 8, "metadata": {}, "output_type": "execute_result" } @@ -1240,28 +1250,53 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## 4. Save in BigQuery" + "## 6. Save in BigQuery" ] }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 9, "metadata": {}, "outputs": [ { "data": { + "text/html": [ + "Copy job d1def4a4-1da1-43a9-8ae5-4459444d993d is DONE. Open Job" + ], "text/plain": [ - "Pipeline(steps=[('preproc',\n", - " ColumnTransformer(transformers=[('onehot', OneHotEncoder(),\n", - " ['island', 'species', 'sex']),\n", - " ('scaler', StandardScaler(),\n", - " ['culmen_depth_mm',\n", - " 'culmen_length_mm',\n", - " 'flipper_length_mm'])])),\n", - " ('linreg', LinearRegression(fit_intercept=False))])" + "" ] }, - "execution_count": 20, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "Pipeline(steps=[('transform',\n", + " ColumnTransformer(transformers=[('ont_hot_encoder',\n", + " OneHotEncoder(max_categories=1000001,\n", + " min_frequency=0),\n", + " 'island'),\n", + " ('standard_scaler',\n", + " StandardScaler(),\n", + " 'culmen_length_mm'),\n", + " ('standard_scaler',\n", + " StandardScaler(),\n", + " 'culmen_depth_mm'),\n", + " ('standard_scaler',\n", + " StandardScaler(),\n", + " 'flipper_length_mm'),\n", + " ('ont_hot_encoder',\n", + " OneHotEncoder(max_categories=1000001,\n", + " min_frequency=0),\n", + " 'sex')])),\n", + " ('estimator',\n", + " LinearRegression(fit_intercept=False,\n", + " optimize_strategy='NORMAL_EQUATION'))])" + ] + }, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -1269,6 +1304,13 @@ "source": [ "pipeline.to_gbq(\"bigframes-dev.bigframes_demo_us.penguin_model\", replace=True)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -1287,7 +1329,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.9" }, "orig_nbformat": 4, "vscode": { diff --git a/tests/system/large/ml/test_cluster.py b/tests/system/large/ml/test_cluster.py index f01116665f1..9244c4b9f19 100644 --- a/tests/system/large/ml/test_cluster.py +++ b/tests/system/large/ml/test_cluster.py @@ -98,7 +98,9 @@ def test_cluster_configure_fit_score_predict( score_result, score_expected, check_exact=False, rtol=0.1 ) - result = model.predict(new_penguins).to_pandas() + predictions = model.predict(new_penguins).to_pandas() + assert predictions.shape == (4, 9) + result = predictions[["CENTROID_ID"]] expected = pd.DataFrame( {"CENTROID_ID": [2, 3, 1, 2]}, dtype="Int64", diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py index a8613dfeb9b..b98d7a757ca 100644 --- a/tests/system/large/ml/test_ensemble.py +++ b/tests/system/large/ml/test_ensemble.py @@ -179,7 +179,7 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id): ) -@pytest.mark.flaky(retries=2, delay=120) +# @pytest.mark.flaky(retries=2, delay=120) def test_xgbclassifier_dart_booster_multiple_params( penguins_df_default_index, dataset_id ): diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 3e569540587..2929baf3f7b 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -545,7 +545,9 @@ def test_pipeline_standard_scaler_kmeans_fit_score_predict( score_result, score_expected, check_exact=False, rtol=0.1 ) - result = pl.predict(new_penguins).to_pandas().sort_index() + predictions = pl.predict(new_penguins).to_pandas().sort_index() + assert predictions.shape == (6, 9) + result = predictions[["CENTROID_ID"]] expected = pd.DataFrame( {"CENTROID_ID": [1, 2, 1, 2, 1, 2]}, dtype="Int64", diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py index 266a38e3ee1..a9fec0bbce0 100644 --- a/tests/system/small/ml/test_cluster.py +++ b/tests/system/small/ml/test_cluster.py @@ -62,7 +62,9 @@ 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() + predictions = penguins_kmeans_model.predict(new_penguins).to_pandas() + assert predictions.shape == (4, 9) + result = predictions[["CENTROID_ID"]] expected = pd.DataFrame( {"CENTROID_ID": [2, 3, 1, 2]}, dtype="Int64", diff --git a/tests/system/small/ml/test_ensemble.py b/tests/system/small/ml/test_ensemble.py index bba083d98d9..55d9fef661a 100644 --- a/tests/system/small/ml/test_ensemble.py +++ b/tests/system/small/ml/test_ensemble.py @@ -98,7 +98,9 @@ def test_xgbregressor_model_score_series( def test_xgbregressor_model_predict( penguins_xgbregressor_model: bigframes.ml.ensemble.XGBRegressor, new_penguins_df ): - result = penguins_xgbregressor_model.predict(new_penguins_df).to_pandas() + predictions = penguins_xgbregressor_model.predict(new_penguins_df).to_pandas() + assert predictions.shape == (3, 8) + result = predictions[["predicted_body_mass_g"]] expected = pandas.DataFrame( {"predicted_body_mass_g": ["4293.1538089", "3410.0271", "3357.944"]}, dtype="Float64", @@ -220,7 +222,9 @@ def test_xgbclassifier_model_score_series( def test_xgbclassifier_model_predict( penguins_xgbclassifier_model: bigframes.ml.ensemble.XGBClassifier, new_penguins_df ): - result = penguins_xgbclassifier_model.predict(new_penguins_df).to_pandas() + predictions = penguins_xgbclassifier_model.predict(new_penguins_df).to_pandas() + assert predictions.shape == (3, 9) + result = predictions[["predicted_sex"]] expected = pandas.DataFrame( {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, dtype="string[pyarrow]", @@ -363,7 +367,11 @@ 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).to_pandas() + predictions = penguins_randomforest_regressor_model.predict( + new_penguins_df + ).to_pandas() + assert predictions.shape == (3, 8) + result = predictions[["predicted_body_mass_g"]] expected = pandas.DataFrame( {"predicted_body_mass_g": ["3897.341797", "3458.385742", "3458.385742"]}, dtype="Float64", @@ -490,7 +498,11 @@ 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).to_pandas() + predictions = penguins_randomforest_classifier_model.predict( + new_penguins_df + ).to_pandas() + assert predictions.shape == (3, 9) + result = predictions[["predicted_sex"]] expected = pandas.DataFrame( {"predicted_sex": ["MALE", "MALE", "FEMALE"]}, dtype="string[pyarrow]", diff --git a/tests/system/small/ml/test_forecasting.py b/tests/system/small/ml/test_forecasting.py index 55079c94cfd..948db59650d 100644 --- a/tests/system/small/ml/test_forecasting.py +++ b/tests/system/small/ml/test_forecasting.py @@ -22,6 +22,8 @@ def test_model_predict(time_series_arima_plus_model): utc = pytz.utc predictions = time_series_arima_plus_model.predict().to_pandas() + assert predictions.shape == (3, 8) + result = predictions[["forecast_timestamp", "forecast_value"]] expected = pd.DataFrame( { "forecast_timestamp": [ @@ -38,7 +40,7 @@ def test_model_predict(time_series_arima_plus_model): ) pd.testing.assert_frame_equal( - predictions, + 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 d3055670668..9008e85a0b5 100644 --- a/tests/system/small/ml/test_imported.py +++ b/tests/system/small/ml/test_imported.py @@ -32,7 +32,9 @@ def test_tensorflow_create_model_default_session(imported_tensorflow_model_path) 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() + predictions = imported_tensorflow_model.predict(df).to_pandas() + assert predictions.shape == (3, 2) + result = predictions[["dense_1"]] # 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. @@ -72,7 +74,9 @@ def test_onnx_create_model_default_session(imported_onnx_model_path): def test_onnx_model_predict(imported_onnx_model, onnx_iris_df): - result = imported_onnx_model.predict(onnx_iris_df).to_pandas() + predictions = imported_onnx_model.predict(onnx_iris_df).to_pandas() + assert predictions.shape == (3, 7) + result = predictions[["label", "probabilities"]] value1 = np.array([0.9999993443489075, 0.0, 0.0]) value2 = np.array([0.0, 0.0, 0.9999993443489075]) expected = pd.DataFrame( diff --git a/tests/system/small/ml/test_linear_model.py b/tests/system/small/ml/test_linear_model.py index 3a8232ed9ec..218c1074ab4 100644 --- a/tests/system/small/ml/test_linear_model.py +++ b/tests/system/small/ml/test_linear_model.py @@ -91,13 +91,15 @@ def test_linear_reg_model_score_series( def test_linear_reg_model_predict(penguins_linear_model, new_penguins_df): predictions = penguins_linear_model.predict(new_penguins_df).to_pandas() + assert predictions.shape == (3, 8) + result = predictions[["predicted_body_mass_g"]] 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(), + result.sort_index(), expected, check_exact=False, rtol=0.1, @@ -224,13 +226,15 @@ def test_logistic_model_score_series( def test_logsitic_model_predict(penguins_logistic_model, new_penguins_df): predictions = penguins_logistic_model.predict(new_penguins_df).to_pandas() + assert predictions.shape == (3, 9) + result = predictions[["predicted_sex"]] 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(), + result.sort_index(), expected, check_exact=False, rtol=0.1, diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 79d3c40317d..306098548ec 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import TestCase - import numpy as np import pytest @@ -48,7 +46,7 @@ def test_create_text_generator_model_default_session(bq_connection, llm_text_pan 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 df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -72,7 +70,7 @@ def test_create_text_generator_32k_model_default_session( 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 df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -97,7 +95,7 @@ def test_create_text_generator_model_default_connection(llm_text_pandas_df): ) df = model.predict(llm_text_df).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -109,7 +107,7 @@ def test_text_generator_predict_default_params_success( palm2_text_generator_model, llm_text_df ): df = palm2_text_generator_model.predict(llm_text_df).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -120,7 +118,7 @@ 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 df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -132,7 +130,7 @@ def test_text_generator_predict_arbitrary_col_label_success( ): llm_text_df = llm_text_df.rename(columns={"prompt": "arbitrary"}) df = palm2_text_generator_model.predict(llm_text_df).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -145,7 +143,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 ).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns series = df["ml_generate_text_llm_result"] assert all(series.str.len() > 20) @@ -196,7 +194,7 @@ def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df ): df = palm2_embedding_generator_model.predict(llm_text_df).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "text_embedding" in df.columns series = df["text_embedding"] value = series[0] @@ -209,7 +207,7 @@ def test_embedding_generator_multilingual_predict_success( palm2_embedding_generator_multilingual_model, llm_text_df ): df = palm2_embedding_generator_multilingual_model.predict(llm_text_df).to_pandas() - TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert df.shape == (3, 4) assert "text_embedding" in df.columns series = df["text_embedding"] value = series[0] @@ -222,7 +220,7 @@ 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 df.shape == (3, 4) assert "text_embedding" in df.columns series = df["text_embedding"] value = series[0] diff --git a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py index 5369d3662d7..be6c5e7c521 100644 --- a/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py +++ b/third_party/bigframes_vendored/sklearn/cluster/_kmeans.py @@ -20,19 +20,7 @@ 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 (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: - bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing the - class labels for each sample. - """ - raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + pass class KMeans(_BaseKMeans): @@ -73,7 +61,7 @@ def predict( DataFrame of shape (n_samples, n_features). New data to predict. Returns: - bigframes.dataframe.DataFrame: DataFrame of the cluster each sample belongs to. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted labels. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/linear_model/_base.py b/third_party/bigframes_vendored/sklearn/linear_model/_base.py index 8dc3b6280a4..ab946e58610 100644 --- a/third_party/bigframes_vendored/sklearn/linear_model/_base.py +++ b/third_party/bigframes_vendored/sklearn/linear_model/_base.py @@ -16,7 +16,6 @@ # 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 bigframes import constants from third_party.bigframes_vendored.sklearn.base import ( @@ -35,7 +34,7 @@ def predict(self, X): Series or DataFrame of shape (n_samples, n_features). Samples. Returns: - bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,). Returns predicted values. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -50,8 +49,7 @@ def predict(self, X): which we want to get the predictions. Returns: - bigframes.dataframe.DataFrame: DataFrame of shape (n_samples,), containing - the class labels for each sample. + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. """ 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 b7b43b85a3e..dfd0ba73563 100644 --- a/third_party/bigframes_vendored/xgboost/sklearn.py +++ b/third_party/bigframes_vendored/xgboost/sklearn.py @@ -18,7 +18,7 @@ def predict(self, X): 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, n_input_columns + n_prediction_columns). Returns predicted values. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) From a7298317ea2604faa6ae31817f1f729d7e0b9818 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 16 Nov 2023 14:44:14 -0800 Subject: [PATCH 167/197] fix: invalid JSON type of the notebook (#215) 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_kmeans.ipynb | 1064 +---------------- 1 file changed, 33 insertions(+), 1031 deletions(-) diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb index ae038136393..8d759509254 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -139,17 +139,9 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Updated property [core/project].\n" - ] - } - ], + "outputs": [], "source": [ "# set your project ID below\n", "PROJECT_ID = \"\" # @param {type:\"string\"}\n", @@ -170,7 +162,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -264,7 +256,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": { "id": "R7STCS8xB5d2" }, @@ -296,7 +288,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": { "id": "zDSwoBo1CU3G" }, @@ -307,101 +299,11 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": { "id": "tYDoaKgJChiq" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 9f096761-e3b5-4d58-a9f7-485ced67afca is DONE. 2.3 GB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job ee8fecb1-2e30-407d-9e2e-9e76061da9e7 is DONE. 2.3 GB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "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", - "
consumer_complaint_narrative
0I signed a contract as a condition of employme...
1First, I want to disclose that XXXX and XXXX b...
2Frequent calls from Focused Receivables Manage...
3I recently contacted Enhanced Recovery Company...
4This began when I subscribed to XXXX XXXX inte...
\n", - "

5 rows × 1 columns

\n", - "
[5 rows x 1 columns in total]" - ], - "text/plain": [ - " consumer_complaint_narrative\n", - "0 I signed a contract as a condition of employme...\n", - "1 First, I want to disclose that XXXX and XXXX b...\n", - "2 Frequent calls from Focused Receivables Manage...\n", - "3 I recently contacted Enhanced Recovery Company...\n", - "4 This began when I subscribed to XXXX XXXX inte...\n", - "\n", - "[5 rows x 1 columns]" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "issues_df = input_df[[\"consumer_complaint_narrative\"]].dropna()\n", "issues_df.head(n=5) # View the first five complaints" @@ -417,7 +319,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": { "id": "OltYSUEcsSOW" }, @@ -439,24 +341,11 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": { "id": "li38q8FzDDMu" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 52d2e961-7896-497c-8b03-ab7374737679 is DONE. 0 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "from bigframes.ml.llm import PaLM2TextEmbeddingGenerator\n", "\n", @@ -465,125 +354,11 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": null, "metadata": { "id": "cOuSOQ5FDewD" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job d093d51a-8eda-442f-80cd-568cb76e00b3 is DONE. 10.6 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 6419df65-3e96-41a7-a7b5-3d058e18763a is DONE. 80.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 917f09ea-c468-4363-a856-b1091e5f775f is DONE. 80.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 5c9679e7-192c-40b5-a14b-edc0fa113eaa is DONE. 61.5 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "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", - "
text_embedding
422[-0.012013785541057587, 0.003669967409223318, ...
616[-0.014948881231248379, -0.04672442376613617, ...
833[-0.01951478235423565, -0.027120858430862427, ...
1370[-0.03140445053577423, -0.048797041177749634, ...
1430[-0.02244548313319683, -0.03336532413959503, 0...
\n", - "

5 rows × 1 columns

\n", - "
[5 rows x 1 columns in total]" - ], - "text/plain": [ - " text_embedding\n", - "422 [-0.012013785541057587, 0.003669967409223318, ...\n", - "616 [-0.014948881231248379, -0.04672442376613617, ...\n", - "833 [-0.01951478235423565, -0.027120858430862427, ...\n", - "1370 [-0.03140445053577423, -0.048797041177749634, ...\n", - "1430 [-0.02244548313319683, -0.03336532413959503, 0...\n", - "\n", - "[5 rows x 1 columns]" - ] - }, - "execution_count": 29, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Will take ~3 minutes to compute the embeddings\n", "predicted_embeddings = model.predict(downsampled_issues_df)\n", @@ -593,263 +368,14 @@ }, { "cell_type": "code", - "execution_count": 30, + "execution_count": null, "metadata": { "id": "4H_etYfsEOFP" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job ce9cb0f9-4b0d-40a1-81f3-d6e60dd6c684 is DONE. 160.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job aa692a30-5706-46ad-8029-faf2fac66234 is DONE. 72.2 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "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", - "
consumer_complaint_narrativetext_embedding
2580664Hello, my name is XXXX XXXX, and I am writing ...[0.0003211698785889894, -0.01816680282354355, ...
1806973This is XXXX XXXX and I am submitting this com...[-0.009485247544944286, -0.025846892967820168,...
2055053XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX...[-0.010950954630970955, -0.0249345600605011, 0...
2515231When I reinvestigated my credit report, I real...[-0.009660656563937664, -0.05793113633990288, ...
2633049Checking my credit report XX/XX/2018 with all ...[-0.0022159104701131582, -0.03330004960298538,...
3117273I contacted TransUnion and spoke a credit rep ...[-0.015955328941345215, -0.006488671060651541,...
698814XXXX XXXX XXXX. makes daily calls to me cell c...[0.005397460889071226, -0.01276913657784462, 0...
267826Can we please reopen Case : XXXX? \n", - "\n", - "Wells Farg...[0.004065403249114752, -0.0005381882656365633,...
54019My rights under 15 USC 1681 have been violated...[0.013823015615344048, -0.02010691538453102, 0...
141050To whom it may concern : My personal informati...[0.008104532025754452, -0.01856449618935585, 0...
2962076I have had a CashApp account since last year, ...[-0.0003019514260813594, -0.03750108182430267,...
2481105that some of the information was erroneous. Th...[-0.014868081547319889, -0.0443895161151886, -...
431562I have disputed the referenced accounts to the...[-0.0020524838473647833, -0.04830990731716156,...
1953029On, XX/XX/22, I attempted to complete a transa...[-0.01599179394543171, -0.0074900356121361256,...
2395979Subject : XXXX XXXX XXXX compensation, refund,...[-0.0035950862802565098, -0.014652969315648079...
455524I paid off my mortgage on XX/XX/2019. The comp...[-0.01100730150938034, -0.03495829552412033, 0...
2155924This kind of account is placed as a charged of...[-0.028635455295443535, -0.028604287654161453,...
1069497This is one of many issues I have had with Wel...[0.008871790021657944, -0.028502725064754486, ...
3181689I have disputed this account with MONTEREY FIN...[-0.004721717908978462, -0.03673810139298439, ...
274268Lender is not updating my loan status in the V...[-0.009221495129168034, -0.0289347805082798, 0...
1671305XXXX is a peer to peer lending conmpany that u...[-0.02911308966577053, -0.01850792020559311, -...
886026( DISPUTE CODE - XXXX ) My personal informatio...[-0.007220877334475517, -0.016615957021713257,...
1044431I filed a complaint against PNC this year and ...[0.002848619595170021, -0.035117778927087784, ...
1938481I applied for a modification and was approved....[-0.03114932030439377, -0.0421406552195549, 0....
1987834Ive been Disputting my XXXX XXXX I opened this...[-0.009406660683453083, -0.020967338234186172,...
\n", - "

25 rows × 2 columns

\n", - "
[10000 rows x 2 columns in total]" - ], - "text/plain": [ - " consumer_complaint_narrative \\\n", - "2580664 Hello, my name is XXXX XXXX, and I am writing ... \n", - "1806973 This is XXXX XXXX and I am submitting this com... \n", - "2055053 XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX... \n", - "2515231 When I reinvestigated my credit report, I real... \n", - "2633049 Checking my credit report XX/XX/2018 with all ... \n", - "3117273 I contacted TransUnion and spoke a credit rep ... \n", - "698814 XXXX XXXX XXXX. makes daily calls to me cell c... \n", - "267826 Can we please reopen Case : XXXX? \n", - "\n", - "Wells Farg... \n", - "54019 My rights under 15 USC 1681 have been violated... \n", - "141050 To whom it may concern : My personal informati... \n", - "2962076 I have had a CashApp account since last year, ... \n", - "2481105 that some of the information was erroneous. Th... \n", - "431562 I have disputed the referenced accounts to the... \n", - "1953029 On, XX/XX/22, I attempted to complete a transa... \n", - "2395979 Subject : XXXX XXXX XXXX compensation, refund,... \n", - "455524 I paid off my mortgage on XX/XX/2019. The comp... \n", - "2155924 This kind of account is placed as a charged of... \n", - "1069497 This is one of many issues I have had with Wel... \n", - "3181689 I have disputed this account with MONTEREY FIN... \n", - "274268 Lender is not updating my loan status in the V... \n", - "1671305 XXXX is a peer to peer lending conmpany that u... \n", - "886026 ( DISPUTE CODE - XXXX ) My personal informatio... \n", - "1044431 I filed a complaint against PNC this year and ... \n", - "1938481 I applied for a modification and was approved.... \n", - "1987834 Ive been Disputting my XXXX XXXX I opened this... \n", - "\n", - " text_embedding \n", - "2580664 [0.0003211698785889894, -0.01816680282354355, ... \n", - "1806973 [-0.009485247544944286, -0.025846892967820168,... \n", - "2055053 [-0.010950954630970955, -0.0249345600605011, 0... \n", - "2515231 [-0.009660656563937664, -0.05793113633990288, ... \n", - "2633049 [-0.0022159104701131582, -0.03330004960298538,... \n", - "3117273 [-0.015955328941345215, -0.006488671060651541,... \n", - "698814 [0.005397460889071226, -0.01276913657784462, 0... \n", - "267826 [0.004065403249114752, -0.0005381882656365633,... \n", - "54019 [0.013823015615344048, -0.02010691538453102, 0... \n", - "141050 [0.008104532025754452, -0.01856449618935585, 0... \n", - "2962076 [-0.0003019514260813594, -0.03750108182430267,... \n", - "2481105 [-0.014868081547319889, -0.0443895161151886, -... \n", - "431562 [-0.0020524838473647833, -0.04830990731716156,... \n", - "1953029 [-0.01599179394543171, -0.0074900356121361256,... \n", - "2395979 [-0.0035950862802565098, -0.014652969315648079... \n", - "455524 [-0.01100730150938034, -0.03495829552412033, 0... \n", - "2155924 [-0.028635455295443535, -0.028604287654161453,... \n", - "1069497 [0.008871790021657944, -0.028502725064754486, ... \n", - "3181689 [-0.004721717908978462, -0.03673810139298439, ... \n", - "274268 [-0.009221495129168034, -0.0289347805082798, 0... \n", - "1671305 [-0.02911308966577053, -0.01850792020559311, -... \n", - "886026 [-0.007220877334475517, -0.016615957021713257,... \n", - "1044431 [0.002848619595170021, -0.035117778927087784, ... \n", - "1938481 [-0.03114932030439377, -0.0421406552195549, 0.... \n", - "1987834 [-0.009406660683453083, -0.020967338234186172,... \n", - "...\n", - "\n", - "[10000 rows x 2 columns]" - ] - }, - "execution_count": 30, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Join the complaints with their embeddings in the same DataFrame\n", - "combined_df = downsampled_issues_df.join(predicted_embeddings, how=\"left\")\n", - "combined_df" + "combined_df = downsampled_issues_df.join(predicted_embeddings)" ] }, { @@ -872,7 +398,7 @@ }, { "cell_type": "code", - "execution_count": 31, + "execution_count": null, "metadata": { "id": "AhNTnEC5FRz2" }, @@ -893,152 +419,14 @@ }, { "cell_type": "code", - "execution_count": 32, + "execution_count": null, "metadata": { "id": "6poSxh-fGJF7" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 65eb317d-59f1-4d10-acd1-4b7f3778114c is DONE. 61.7 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 156e445e-cc01-4b30-84cc-ac1c98a69b81 is DONE. 0 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 5befc212-f4a3-4e33-b1b2-01e809acdcbd is DONE. 61.9 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job bd271178-8b8d-45dc-ac57-7f0194d0daac is DONE. 80.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job bbfb9cca-622d-4bf5-9fc0-6d9a85287d41 is DONE. 80.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job a5f30b32-9fb0-42b4-b426-d8484f008bdb is DONE. 160.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "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", - "
CENTROID_ID
4222
6163
8335
13707
14303
\n", - "

5 rows × 1 columns

\n", - "
[5 rows x 1 columns in total]" - ], - "text/plain": [ - " CENTROID_ID\n", - "422 2\n", - "616 3\n", - "833 5\n", - "1370 7\n", - "1430 3\n", - "\n", - "[5 rows x 1 columns]" - ] - }, - "execution_count": 32, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Use KMeans clustering to calculate our groups. Will take ~3 minutes.\n", - "cluster_model.fit(combined_df[\"text_embedding\"])\n", + "cluster_model.fit(combined_df[[\"text_embedding\"]])\n", "clustered_result = cluster_model.predict(combined_df[[\"text_embedding\"]])\n", "# Notice the CENTROID_ID column, which is the ID number of the group that\n", "# each complaint belongs to.\n", @@ -1047,123 +435,13 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 7a41196e-ea67-44ac-95a7-7dce620d6d21 is DONE. 320.0 kB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 8008b482-1a0d-461f-a215-4676d9d918dc is DONE. 72.4 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "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", - "
consumer_complaint_narrativetext_embeddingCENTROID_ID
2580664Hello, my name is XXXX XXXX, and I am writing ...[0.0003211698785889894, -0.01816680282354355, ...2
1806973This is XXXX XXXX and I am submitting this com...[-0.009485247544944286, -0.025846892967820168,...5
2055053XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX...[-0.010950954630970955, -0.0249345600605011, 0...3
2515231When I reinvestigated my credit report, I real...[-0.009660656563937664, -0.05793113633990288, ...5
2633049Checking my credit report XX/XX/2018 with all ...[-0.0022159104701131582, -0.03330004960298538,...3
\n", - "

5 rows × 3 columns

\n", - "
[5 rows x 3 columns in total]" - ], - "text/plain": [ - " consumer_complaint_narrative \\\n", - "2580664 Hello, my name is XXXX XXXX, and I am writing ... \n", - "1806973 This is XXXX XXXX and I am submitting this com... \n", - "2055053 XXXX XXXX XXXX, XXXX. ( address : XXXX XXXX XX... \n", - "2515231 When I reinvestigated my credit report, I real... \n", - "2633049 Checking my credit report XX/XX/2018 with all ... \n", - "\n", - " text_embedding CENTROID_ID \n", - "2580664 [0.0003211698785889894, -0.01816680282354355, ... 2 \n", - "1806973 [-0.009485247544944286, -0.025846892967820168,... 5 \n", - "2055053 [-0.010950954630970955, -0.0249345600605011, 0... 3 \n", - "2515231 [-0.009660656563937664, -0.05793113633990288, ... 5 \n", - "2633049 [-0.0022159104701131582, -0.03330004960298538,... 3 \n", - "\n", - "[5 rows x 3 columns]" - ] - }, - "execution_count": 33, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Join the group number to the complaints and their text embeddings\n", "combined_clustered_result = combined_df.join(clustered_result)\n", - "\n", - "combined_clustered_result.head(n=5)" + "combined_clustered_result.head(n=5) " ] }, { @@ -1194,36 +472,11 @@ }, { "cell_type": "code", - "execution_count": 34, + "execution_count": null, "metadata": { "id": "2E7wXM_jGqo6" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 50c7c0dd-94a2-494e-a37f-6a838a518f6c is DONE. 11.0 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job d96c847f-c292-4804-bd05-fd643c41c7a5 is DONE. 11.0 MB processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "# Using bigframes, with syntax identical to pandas,\n", "# filter out the first and second groups\n", @@ -1240,100 +493,11 @@ }, { "cell_type": "code", - "execution_count": 36, + "execution_count": null, "metadata": { "id": "ZNDiueI9IP5e" }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "comment list 1:\n", - "1. XXXX is a peer to peer lending conmpany that uses borrowers crypto to collateralize loans from investors ( like myself ). I've been investing with them for almost XXXX years and currently have {$240000.00} tied up in lending products with XXXX. \n", - "As of XXXX days ago we received an email saying all business operations have been ceased and no withdrawals or deposits will be allowed. They said they'll update customers within 10 days, but no one can reach anyone at the company to find out any more details as they are not answering calls nor returning emails. It also appears the company has scrubbed its XXXX page and the XXXX pages of top executives. \n", - "\n", - "All collateral and client 's investment funds are supposedly held at or processed through XXXX XXXX XXXX ( registered SEC company ). XXXX XXXX keeps telling us to contact XXXX and won't give us any information, so we have no way to find out what's happening with our funds/collateral or if everything is gone. We have a XXXX channel up where people are gathering evidence, documentation, etc. This is probably the best place to start to get a broad view of what's happening. Details below. \n", - "\n", - "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX CONST LLC ( Business ID : XXXX ) FoXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX 'Cease of Operations ' email received by all investors XXXX XX/XX/2022 at XXXX : \" Dear XXXX Users, Given the collapses of several cryptocurrencies so far this year and the rapidly deteriorating market conditions that have been prompting heavy withdrawals across all XXXX lending and XXXX exchange platforms recently, we are sad to inform you that we are unable to continue to operate our business as usual. As such, we are limiting our business activities, including pausing user withdrawals as allowed under our Terms of XXXX. \n", - "No deposit or investment request will be processed at this time. \n", - "\n", - "Our team is working diligently towards our objective of maximizing value for all of our Users, and our top priority continues to be to protect your interests. As we explore all options available to us, we will provide updates to you as we go. \n", - "\n", - "We hope to communicate with you within the next XXXX business days on the next steps to address the situation. We appreciate your patience in this trying time. \n", - "\n", - "Sincerely yoursXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", - "2. Submitted XX/XX/XXXX\n", - "Typed XX/XX/XXXX:\n", - "\n", - "XX/XX/XXXX\n", - "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX, XXXX XXXX\n", - "PH:. XXXX\n", - "PH: XXXX\n", - "EM:\n", - "XXXX\n", - "XXXX\n", - "XXXX XXXX \n", - "XXXX XXXX\n", - "Date of Birth XX/XX/XXXX\n", - "SS#: XXXX\n", - "TO:\n", - "* Consumer Financial Protection Brueau\n", - "* Department of Veteran Affairs, Office of the Inspector General\n", - "My name is XXXX XXXX XXXX, I've received more than one email from Discover Card in my XXXX XXXX, past emails from Discover Card were unautherized deletions.\n", - "From: Discover Card XXXX\n", - "To: You XXXX\n", - "Date: XX/XX/XXXX, XXXX XXXX XXXX From: Discover Card XXXX>\n", - "To Recipient \n", - "Date Mon, XX/XX/XXXX XXXX XXXX\n", - "I dont and havent ever had a Discover Checking, Savings, Business Accounts nor Loans of any kind through any Bank called Discover. The 1st time I was contacted by Discover Card I resided alone from XX/XX/XXXX to XX/XX/XXXXat XXXX XXXX XXXX at XXXX XXXX XXXX XXXX XXXX in XXXX, XXXX years prior to me moving here to XXXX, XXXX in XX/XX/XXXX. When \n", - "\n", - "\n", - "Discover Card had 1st contacted me in XXXX, XXXX it was associated with my XXXX XXXX XXXX website related online Merchants Account. Not once have I ever applied for or had any Website Merchant Accounts here in XXXX; I only applied for online online Merchant Accounts associated with my XXXX related Accounts I purchased while residing in XXXX, XXXX. Some of my website related information was stolen both in XXXX, XXXX and here in XXXX along with my other property that hasn't been returned to me. I don't and haven't ever had any XXXX XXXX related Agreements,Contracts or Credit Cards offered to Veterans associated with ones businesses. Nor have I ever applied for or had a Business License or Business Permit in any City or State inspite of my diverse interest. Not once have I ever allowed another be it an Paralegal, Payee, Attorney, Employers, Landlords, Veteran Organizations including Vocational Rehabilitation Programs, XXXX( XXXX XXXX XXXX, XXXX XXXX, Entertainment Companies, Banks, Celebrity Personal Assistant Agencies or Celebs, Shelters, Charities, HUD, Housing Arthority, Department of Veteran Affairs, Military, Law Enforcement or anyone else nor their employess to sign any business related Agreements or Contracts on my behalf; not even my family members or friends. \n", - "None of my XXXX XXXX attempts were associated with my Employers, Department of Veteran Affairs,Vocational Rehabilitation Programs Military, Landlords, HUD( Housing Authority),Friends, Family nor did I ever sign related Agreements or Contracts with them. Not once had I ever provided anyone the passwords to be able to sign into my accounts rather were aware of my accounts or not. Yes, my desktop computer that was stolen along with my other property XX/XX/XXXX was registered with my Online Merchant Account. I had paid for my Merchant related Accounts through my same XXXX XXXX XXXX Account I purchased both of my XXXX XXXX XXXX related accounts through. That was 1st once during the Summer of XX/XX/XXXX and 2nd my related website months later, while I resided in XXXX XXXX and I worked for XXXX. I never offered nor did I ever sign any business Contracts or Agreements with XXXX nor my Landlord or their staff associted with any of my online websites or Merchant Accounts. My XXXX XXXX XXXX Compensation was deposited into both of my XXXX XXXX XXXX Accounts at that time. My account was changed during the Summer of XX/XX/XXXXbecause of theft of my Bank Card. None of my Checking,Savings, past Credit Cards or Business related were shared accounts in which others were allowed to \n", - "use to make purchases. I had written checks from my XXXX XXXX XXXX account to pay for my XXXX XXXX XXXX XXXX on the XXXX XXXX here in XXXX in XX/XX/XXXX before it's name changed to XXXX XXXX. Prior to me using my same account open a Checking account in person at XXXX XXXX before it's name was changed to XXXX XXXX. Where my XXXX XXXX XXXX XXXX has been deposited since that time. I had used my XXXX XXXX Checking to pay for my XXXX XXXX XXXX XXXX both before theft of my property XX/XX/XXXX and that was also prior to the theft of my property from my XXXX XXXX XXXX XXXX in XX/XX/XXXX.\n", - "I've stated this many times:\n", - "I paid for my 1st XXXX XXXX XXXX Membership while employed at XXXX using my XXXX XXXX XXXX account XXXX my XXXX XXXX XXXX XXXX was also deposited. That was changed to XXXX because I didn't receive my 1st XXXX XXXX XXXX Card the bank sent to XXXX XXXX residence on XXXX XXXX in XX/XX/XXXX while I was there. In which both my XXXX salary and XXXX XXXX XXXX XXXX were deposited into my account, no money from XXXX XXXX nor anyone else that was at that residence was given to nor were any of my children there. Nor did XXXX or any other person at that residence ever give me my missing Bank Card not even after I moved out and stayed a month at XXXX XXXX XXXX using my replacement card to pay for my Hotel room. Which is the same account I used to pay for XXXX XXXX Membership, XXXX XXXX XXXX, XXXX XXXX Membership fees, and various online Merchant Account activation related fees.\n", - "* XXXX XXXX XXXX.\n", - "XXXX XXXX XXXX XXXX. Membership\n", - "\n", - "# XXXX\n", - "* XXXX XXXX Membership\n", - "# XXXX\n", - "* Total Merchant Services XXXX and XXXX.\n", - "* XXXX XXXX XXXX XXXX XXXX\n", - "* XXXX XXXX changed my $XXXX a month fees to my XXXX XXXX XXXX account #XXXX.\n", - "XX/XX/XXXX - XX/XX/XXXX XXXX XXXX, XXXX.\n", - "\n", - "Rep: XXXX XXXX XXXX, Fl \n", - "XXXX\n", - "XXXX Website \n", - "XXXX\n", - "Software and website owner, I performed Internet advertising and marketing, to promote this software and website. I worked and XXXX from my home XXXX XXXX XXXX XXXX XXXX , XXXX. I purchased XXXX XXXX XXXX-Software Electronic Book CD and was given a website to promote the software on the internet. The XXXX was given a copy of my website owner certificate document submitted to me when I purchased the software marketing program as well copies of my other school transcripts in addition to XXXX XXXX XXXX for example. XXXX, represented the first initials of my children's names. I wasn't ever paid and I'm still owed the money. Nor did my marketing program have anything to do with any schools, college nor university programs nor did I ever offer or sign any agreement to include it such. Nor did my XXXX XXXX XXXX have anything to do with any other employers, Department of Family and Children, Military, Veteran Organizations or Food Stamp programs, Section 8 nor Indianapolis Housing Authority for example; only me.\n", - "Thank you,\n", - "XXXX XXXX\n", - "3. ACCORDING TO 15 U.S. CODE 6803-DISCLOSURE OF INSTITUTION PRIVACY POLICY, AND ACCORDING TO U.S. CODE 6802- OBLIGATIONS WITH RESPECT TO DISCLOSURES OF PERSONAL INFORMATION. ( b ) OPT OUT ( 1 ) IN GENERAL A FINANCIAL INSTITUTION MAY NOT DISCLOSE NONPUBLIC PERSONAL INFORMATION TO A NONAFFILIATED THIRD PARTY ( TRANSUNION, XXXX, AND XXXX. ) UNLESS- ( A ) SUCH FINANCIAL INSTITUTION CLEARLY AND CONSPICUOUSLY DISCLOSES TO THE CONSUMER, IN WRITING OR IN ELECTRONIC FORM OR OTHER FORM PERMITTED BY THE REGULATIONS PRESCRIBED UNDER SECTION 6804 OF THIS TITLE. ALSO ACCORDING TO THE \" XXXX ACT '', FINANCIAL INSTITUTIONS MUST TELL THEIR CUSTOMERS ABOUT THEIR INFORMATION-SHARING PRACTICES AND EXPLAIN TO CUSTOMERS THEIR RIGHT TO \" OPT OUT '' IF THEY DON'T WANT THEIR INFORMATION SHARED WITH CERTAIN THIRD PARTIES. UNDER THE FDCPA, A COLLECTOR MUST PROVIDE YOU WITH INFORMATION ABOUT THE DEBT IN ITS INITIAL COMMUNICATION OR WITHIN FIVE DAYS AFTER THE INITIAL COMMUNICATION. ALSO, THE FDCPA STATES, \" YOU CAN NOT ATTEMPT TO COLLECT AN DEBT WHILE A PERSON ( THE CONSUMER ) SUPRESS VALIDATION. TRANSUNION, XXXX, XXXX, AND THE ACCOUNTS LISTED BELOW HAVE CLEARLY VIOLATED MY RIGHTS : XXXX ACCOUNT # XXXX, XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX ACCOUNT # XXXX, XXXX XXXX XXXX XXXX ACCOUNT # XXXX, AND XXXX ACCOUNT # XXXX. FAILURE TO RESPOND SATISFACTORILY WITH DELETIONS OF ALL THE ABOVE ACCOUNTS WILL RESULT IN LEGAL ACTIONS BEING TAKEN AGAINST, TRANSUNION, XXXX, XXXX, WHICH I'LL BE SEEKING A {$1000.00} PER VIOLATION FOR DEFAMATION OF CHARACTER ( PER SE ) NEGLIGENT ENABLEMENT OF IDENTITY FRAUD. 15 USC 1681 VIOLATIONS FOR WILLFUL NONCOMPLIANCE-616 CIVIL LIABILITY FOR WILLFUL NONCOPLIANCE. THIS IS THE THIRD TIME I'VE SUBMITTED A COMPLAINT, AND THE REPONSE I GET IS \" YOU CAN NOT LOCATE MY CREDIT REPORT! '' THIS IS CLEARLY NEGLIGENCE.\n", - "4. I do not know how this works, but I need it done or somehow corrected. My name is XXXX XXXX, XXXX XXXX XXXX XXXX TN XXXXMy SS XXXX DOB XXXX. I had some issues with my income being affected by the COVID-19PANDEMICSHUTDOWN. I was under the 1 CARESAct, Pub. L. 116-136, section 4021, codified at FCRAsection 623 ( a ) ( 1 ) ( F ) ( i ) ( I ), 15 U.S.C.1681s- 2 ( a ) ( 1 ) ( F ) ( i ) ( I ). I am requesting some accommodations so I care to protect the integrity of my credit file. US DEPT OF ED / XXXX # XXXX, # XXXX accounts are reporting on XXXX, XXXX The was 30,60, 90 DAYS LATEsince requested assistance due to the pandemic. I found a few accounts that I have never done any business with these companies and the accounts do not belong on my report : XXXX XXXX # XXXX, XXXX XXXX XXXX XXXX # XXXX. \n", - "\n", - "I have some issues with the misspelling of my name, my correct spelling is XXXX XXXX. Please remove any other variation of my name they are not correct. The following addresses do not belong to me please delete them : XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXXSC, XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", - "5. I want to know if this is even legal?! How can they disclose information without knowing its a correct email?!\n", - "\n", - "comment list 2:\n", - "1. Hello, my name is XXXX XXXX, and I am writing to delete the following information in my file. The items I need deleted are listed in the report. I am a victim of identity theft and did not make the charge. I ask that the items be deleted to correct my credit report. I reported the theft of my identity to the Federal Trade Commission and I also have enclosed copies of the Federal Trade Commissions Identity Theft Affidavit. Please delete the items as soon as possible. The accounts are being reported currently open and the accounts need to be closed. \n", - "XXXX account number XXXX opened on XX/XX/2022 for the amount {$530.00} XXXX XXXX XXXX account number XXXX opened on XX/XX/2022 for the amount of {$140.00} The accounts are being reported currently open and need to be closed immediately. \n", - "Based on, 15 U.S. Code 1681c2 a consumer reporting agency shall block the reporting of any information in the file of a consumer that the consumer identifies as information that resulted from an alleged identity theft, not later than 4 business days after the date of receipt. This account should not be furnished on my consumer report. As a consumer I am demanding the deletion of the accounts listed IMMEDIATELY.\n", - "2. To whom it may concern : My personal information was breach in the internet as result accounts had been open in my name, I was advise to fill out an Id theft report to help me deal with this situation, I have listed each one of the accounts that do not belong to me. This is my second request to remove unverified items in my report, but XXXX keep rposting these account with out providing any type of original document as the FCRA provide, you need to provide me with original documents or remove these account immediately.\n", - "3. Ive been Disputting my XXXX XXXX I opened this account and someone got my information and used my card, I contacted XXXX over and over, they removed the negative reporting from my XXXX report but still reporting it negative on my XXXX and Expean this is very unfair to me because Im a victim of identity theft\n", - "4. Today, XX/XX/2021, I received three items in the mail, one envelope containing an unsolicited debit card from Navy Federal credit Union and the other two, with a letter each describing The Important Rights on two accounts should these accounts become delinquent under New York law. \n", - "\n", - "First of all, I never applied for these accounts with Navy Federal, not have I authorized anyone to do so on my behalf. I immediately contacted Navy Federal via phone and was told I was most likely a victim of identity theft and that I should monitor my credit and use a credit monitoring service. I was also asked for my email and mailing information in order to receive a letter from them regarding this issue. \n", - "\n", - "My main concern is having someone using my identity to illegally open bank accounts and commit fraud, destroying my credit and finances in the process. This bank is in another state from where I reside. I have not lived in Virginia nor do I intend to do so in the foreseeable future.\n", - "5. My personal information ( including my SSN, Drivers License Info, Addresses, and more ) was stolen from a hacking, and Equifax did n't tell the public about the hack until more than a month after the hacking. During this time, three Equifax executives were caught inside trading. It really shows how Equifax cares about other people!\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "# Build plain-text prompts to send to PaLM 2. Use only 5 complaints from each group.\n", "prompt1 = 'comment list 1:\\n'\n", @@ -1352,100 +516,11 @@ }, { "cell_type": "code", - "execution_count": 37, + "execution_count": null, "metadata": { "id": "BfHGJLirzSvH" }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Please highlight the most obvious difference betweenthe two lists of comments:\n", - "comment list 1:\n", - "1. XXXX is a peer to peer lending conmpany that uses borrowers crypto to collateralize loans from investors ( like myself ). I've been investing with them for almost XXXX years and currently have {$240000.00} tied up in lending products with XXXX. \n", - "As of XXXX days ago we received an email saying all business operations have been ceased and no withdrawals or deposits will be allowed. They said they'll update customers within 10 days, but no one can reach anyone at the company to find out any more details as they are not answering calls nor returning emails. It also appears the company has scrubbed its XXXX page and the XXXX pages of top executives. \n", - "\n", - "All collateral and client 's investment funds are supposedly held at or processed through XXXX XXXX XXXX ( registered SEC company ). XXXX XXXX keeps telling us to contact XXXX and won't give us any information, so we have no way to find out what's happening with our funds/collateral or if everything is gone. We have a XXXX channel up where people are gathering evidence, documentation, etc. This is probably the best place to start to get a broad view of what's happening. Details below. \n", - "\n", - "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX CONST LLC ( Business ID : XXXX ) FoXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX 'Cease of Operations ' email received by all investors XXXX XX/XX/2022 at XXXX : \" Dear XXXX Users, Given the collapses of several cryptocurrencies so far this year and the rapidly deteriorating market conditions that have been prompting heavy withdrawals across all XXXX lending and XXXX exchange platforms recently, we are sad to inform you that we are unable to continue to operate our business as usual. As such, we are limiting our business activities, including pausing user withdrawals as allowed under our Terms of XXXX. \n", - "No deposit or investment request will be processed at this time. \n", - "\n", - "Our team is working diligently towards our objective of maximizing value for all of our Users, and our top priority continues to be to protect your interests. As we explore all options available to us, we will provide updates to you as we go. \n", - "\n", - "We hope to communicate with you within the next XXXX business days on the next steps to address the situation. We appreciate your patience in this trying time. \n", - "\n", - "Sincerely yoursXXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", - "2. Submitted XX/XX/XXXX\n", - "Typed XX/XX/XXXX:\n", - "\n", - "XX/XX/XXXX\n", - "XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX, XXXX XXXX\n", - "PH:. XXXX\n", - "PH: XXXX\n", - "EM:\n", - "XXXX\n", - "XXXX\n", - "XXXX XXXX \n", - "XXXX XXXX\n", - "Date of Birth XX/XX/XXXX\n", - "SS#: XXXX\n", - "TO:\n", - "* Consumer Financial Protection Brueau\n", - "* Department of Veteran Affairs, Office of the Inspector General\n", - "My name is XXXX XXXX XXXX, I've received more than one email from Discover Card in my XXXX XXXX, past emails from Discover Card were unautherized deletions.\n", - "From: Discover Card XXXX\n", - "To: You XXXX\n", - "Date: XX/XX/XXXX, XXXX XXXX XXXX From: Discover Card XXXX>\n", - "To Recipient \n", - "Date Mon, XX/XX/XXXX XXXX XXXX\n", - "I dont and havent ever had a Discover Checking, Savings, Business Accounts nor Loans of any kind through any Bank called Discover. The 1st time I was contacted by Discover Card I resided alone from XX/XX/XXXX to XX/XX/XXXXat XXXX XXXX XXXX at XXXX XXXX XXXX XXXX XXXX in XXXX, XXXX years prior to me moving here to XXXX, XXXX in XX/XX/XXXX. When \n", - "\n", - "\n", - "Discover Card had 1st contacted me in XXXX, XXXX it was associated with my XXXX XXXX XXXX website related online Merchants Account. Not once have I ever applied for or had any Website Merchant Accounts here in XXXX; I only applied for online online Merchant Accounts associated with my XXXX related Accounts I purchased while residing in XXXX, XXXX. Some of my website related information was stolen both in XXXX, XXXX and here in XXXX along with my other property that hasn't been returned to me. I don't and haven't ever had any XXXX XXXX related Agreements,Contracts or Credit Cards offered to Veterans associated with ones businesses. Nor have I ever applied for or had a Business License or Business Permit in any City or State inspite of my diverse interest. Not once have I ever allowed another be it an Paralegal, Payee, Attorney, Employers, Landlords, Veteran Organizations including Vocational Rehabilitation Programs, XXXX( XXXX XXXX XXXX, XXXX XXXX, Entertainment Companies, Banks, Celebrity Personal Assistant Agencies or Celebs, Shelters, Charities, HUD, Housing Arthority, Department of Veteran Affairs, Military, Law Enforcement or anyone else nor their employess to sign any business related Agreements or Contracts on my behalf; not even my family members or friends. \n", - "None of my XXXX XXXX attempts were associated with my Employers, Department of Veteran Affairs,Vocational Rehabilitation Programs Military, Landlords, HUD( Housing Authority),Friends, Family nor did I ever sign related Agreements or Contracts with them. Not once had I ever provided anyone the passwords to be able to sign into my accounts rather were aware of my accounts or not. Yes, my desktop computer that was stolen along with my other property XX/XX/XXXX was registered with my Online Merchant Account. I had paid for my Merchant related Accounts through my same XXXX XXXX XXXX Account I purchased both of my XXXX XXXX XXXX related accounts through. That was 1st once during the Summer of XX/XX/XXXX and 2nd my related website months later, while I resided in XXXX XXXX and I worked for XXXX. I never offered nor did I ever sign any business Contracts or Agreements with XXXX nor my Landlord or their staff associted with any of my online websites or Merchant Accounts. My XXXX XXXX XXXX Compensation was deposited into both of my XXXX XXXX XXXX Accounts at that time. My account was changed during the Summer of XX/XX/XXXXbecause of theft of my Bank Card. None of my Checking,Savings, past Credit Cards or Business related were shared accounts in which others were allowed to \n", - "use to make purchases. I had written checks from my XXXX XXXX XXXX account to pay for my XXXX XXXX XXXX XXXX on the XXXX XXXX here in XXXX in XX/XX/XXXX before it's name changed to XXXX XXXX. Prior to me using my same account open a Checking account in person at XXXX XXXX before it's name was changed to XXXX XXXX. Where my XXXX XXXX XXXX XXXX has been deposited since that time. I had used my XXXX XXXX Checking to pay for my XXXX XXXX XXXX XXXX both before theft of my property XX/XX/XXXX and that was also prior to the theft of my property from my XXXX XXXX XXXX XXXX in XX/XX/XXXX.\n", - "I've stated this many times:\n", - "I paid for my 1st XXXX XXXX XXXX Membership while employed at XXXX using my XXXX XXXX XXXX account XXXX my XXXX XXXX XXXX XXXX was also deposited. That was changed to XXXX because I didn't receive my 1st XXXX XXXX XXXX Card the bank sent to XXXX XXXX residence on XXXX XXXX in XX/XX/XXXX while I was there. In which both my XXXX salary and XXXX XXXX XXXX XXXX were deposited into my account, no money from XXXX XXXX nor anyone else that was at that residence was given to nor were any of my children there. Nor did XXXX or any other person at that residence ever give me my missing Bank Card not even after I moved out and stayed a month at XXXX XXXX XXXX using my replacement card to pay for my Hotel room. Which is the same account I used to pay for XXXX XXXX Membership, XXXX XXXX XXXX, XXXX XXXX Membership fees, and various online Merchant Account activation related fees.\n", - "* XXXX XXXX XXXX.\n", - "XXXX XXXX XXXX XXXX. Membership\n", - "\n", - "# XXXX\n", - "* XXXX XXXX Membership\n", - "# XXXX\n", - "* Total Merchant Services XXXX and XXXX.\n", - "* XXXX XXXX XXXX XXXX XXXX\n", - "* XXXX XXXX changed my $XXXX a month fees to my XXXX XXXX XXXX account #XXXX.\n", - "XX/XX/XXXX - XX/XX/XXXX XXXX XXXX, XXXX.\n", - "\n", - "Rep: XXXX XXXX XXXX, Fl \n", - "XXXX\n", - "XXXX Website \n", - "XXXX\n", - "Software and website owner, I performed Internet advertising and marketing, to promote this software and website. I worked and XXXX from my home XXXX XXXX XXXX XXXX XXXX , XXXX. I purchased XXXX XXXX XXXX-Software Electronic Book CD and was given a website to promote the software on the internet. The XXXX was given a copy of my website owner certificate document submitted to me when I purchased the software marketing program as well copies of my other school transcripts in addition to XXXX XXXX XXXX for example. XXXX, represented the first initials of my children's names. I wasn't ever paid and I'm still owed the money. Nor did my marketing program have anything to do with any schools, college nor university programs nor did I ever offer or sign any agreement to include it such. Nor did my XXXX XXXX XXXX have anything to do with any other employers, Department of Family and Children, Military, Veteran Organizations or Food Stamp programs, Section 8 nor Indianapolis Housing Authority for example; only me.\n", - "Thank you,\n", - "XXXX XXXX\n", - "3. ACCORDING TO 15 U.S. CODE 6803-DISCLOSURE OF INSTITUTION PRIVACY POLICY, AND ACCORDING TO U.S. CODE 6802- OBLIGATIONS WITH RESPECT TO DISCLOSURES OF PERSONAL INFORMATION. ( b ) OPT OUT ( 1 ) IN GENERAL A FINANCIAL INSTITUTION MAY NOT DISCLOSE NONPUBLIC PERSONAL INFORMATION TO A NONAFFILIATED THIRD PARTY ( TRANSUNION, XXXX, AND XXXX. ) UNLESS- ( A ) SUCH FINANCIAL INSTITUTION CLEARLY AND CONSPICUOUSLY DISCLOSES TO THE CONSUMER, IN WRITING OR IN ELECTRONIC FORM OR OTHER FORM PERMITTED BY THE REGULATIONS PRESCRIBED UNDER SECTION 6804 OF THIS TITLE. ALSO ACCORDING TO THE \" XXXX ACT '', FINANCIAL INSTITUTIONS MUST TELL THEIR CUSTOMERS ABOUT THEIR INFORMATION-SHARING PRACTICES AND EXPLAIN TO CUSTOMERS THEIR RIGHT TO \" OPT OUT '' IF THEY DON'T WANT THEIR INFORMATION SHARED WITH CERTAIN THIRD PARTIES. UNDER THE FDCPA, A COLLECTOR MUST PROVIDE YOU WITH INFORMATION ABOUT THE DEBT IN ITS INITIAL COMMUNICATION OR WITHIN FIVE DAYS AFTER THE INITIAL COMMUNICATION. ALSO, THE FDCPA STATES, \" YOU CAN NOT ATTEMPT TO COLLECT AN DEBT WHILE A PERSON ( THE CONSUMER ) SUPRESS VALIDATION. TRANSUNION, XXXX, XXXX, AND THE ACCOUNTS LISTED BELOW HAVE CLEARLY VIOLATED MY RIGHTS : XXXX ACCOUNT # XXXX, XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX XXXX ACCOUNT # XXXXXXXX XXXX XXXX XXXX ACCOUNT # XXXX, XXXX XXXX XXXX XXXX ACCOUNT # XXXX, AND XXXX ACCOUNT # XXXX. FAILURE TO RESPOND SATISFACTORILY WITH DELETIONS OF ALL THE ABOVE ACCOUNTS WILL RESULT IN LEGAL ACTIONS BEING TAKEN AGAINST, TRANSUNION, XXXX, XXXX, WHICH I'LL BE SEEKING A {$1000.00} PER VIOLATION FOR DEFAMATION OF CHARACTER ( PER SE ) NEGLIGENT ENABLEMENT OF IDENTITY FRAUD. 15 USC 1681 VIOLATIONS FOR WILLFUL NONCOMPLIANCE-616 CIVIL LIABILITY FOR WILLFUL NONCOPLIANCE. THIS IS THE THIRD TIME I'VE SUBMITTED A COMPLAINT, AND THE REPONSE I GET IS \" YOU CAN NOT LOCATE MY CREDIT REPORT! '' THIS IS CLEARLY NEGLIGENCE.\n", - "4. I do not know how this works, but I need it done or somehow corrected. My name is XXXX XXXX, XXXX XXXX XXXX XXXX TN XXXXMy SS XXXX DOB XXXX. I had some issues with my income being affected by the COVID-19PANDEMICSHUTDOWN. I was under the 1 CARESAct, Pub. L. 116-136, section 4021, codified at FCRAsection 623 ( a ) ( 1 ) ( F ) ( i ) ( I ), 15 U.S.C.1681s- 2 ( a ) ( 1 ) ( F ) ( i ) ( I ). I am requesting some accommodations so I care to protect the integrity of my credit file. US DEPT OF ED / XXXX # XXXX, # XXXX accounts are reporting on XXXX, XXXX The was 30,60, 90 DAYS LATEsince requested assistance due to the pandemic. I found a few accounts that I have never done any business with these companies and the accounts do not belong on my report : XXXX XXXX # XXXX, XXXX XXXX XXXX XXXX # XXXX. \n", - "\n", - "I have some issues with the misspelling of my name, my correct spelling is XXXX XXXX. Please remove any other variation of my name they are not correct. The following addresses do not belong to me please delete them : XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXXSC, XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX XXXX\n", - "5. I want to know if this is even legal?! How can they disclose information without knowing its a correct email?!\n", - "comment list 2:\n", - "1. Hello, my name is XXXX XXXX, and I am writing to delete the following information in my file. The items I need deleted are listed in the report. I am a victim of identity theft and did not make the charge. I ask that the items be deleted to correct my credit report. I reported the theft of my identity to the Federal Trade Commission and I also have enclosed copies of the Federal Trade Commissions Identity Theft Affidavit. Please delete the items as soon as possible. The accounts are being reported currently open and the accounts need to be closed. \n", - "XXXX account number XXXX opened on XX/XX/2022 for the amount {$530.00} XXXX XXXX XXXX account number XXXX opened on XX/XX/2022 for the amount of {$140.00} The accounts are being reported currently open and need to be closed immediately. \n", - "Based on, 15 U.S. Code 1681c2 a consumer reporting agency shall block the reporting of any information in the file of a consumer that the consumer identifies as information that resulted from an alleged identity theft, not later than 4 business days after the date of receipt. This account should not be furnished on my consumer report. As a consumer I am demanding the deletion of the accounts listed IMMEDIATELY.\n", - "2. To whom it may concern : My personal information was breach in the internet as result accounts had been open in my name, I was advise to fill out an Id theft report to help me deal with this situation, I have listed each one of the accounts that do not belong to me. This is my second request to remove unverified items in my report, but XXXX keep rposting these account with out providing any type of original document as the FCRA provide, you need to provide me with original documents or remove these account immediately.\n", - "3. Ive been Disputting my XXXX XXXX I opened this account and someone got my information and used my card, I contacted XXXX over and over, they removed the negative reporting from my XXXX report but still reporting it negative on my XXXX and Expean this is very unfair to me because Im a victim of identity theft\n", - "4. Today, XX/XX/2021, I received three items in the mail, one envelope containing an unsolicited debit card from Navy Federal credit Union and the other two, with a letter each describing The Important Rights on two accounts should these accounts become delinquent under New York law. \n", - "\n", - "First of all, I never applied for these accounts with Navy Federal, not have I authorized anyone to do so on my behalf. I immediately contacted Navy Federal via phone and was told I was most likely a victim of identity theft and that I should monitor my credit and use a credit monitoring service. I was also asked for my email and mailing information in order to receive a letter from them regarding this issue. \n", - "\n", - "My main concern is having someone using my identity to illegally open bank accounts and commit fraud, destroying my credit and finances in the process. This bank is in another state from where I reside. I have not lived in Virginia nor do I intend to do so in the foreseeable future.\n", - "5. My personal information ( including my SSN, Drivers License Info, Addresses, and more ) was stolen from a hacking, and Equifax did n't tell the public about the hack until more than a month after the hacking. During this time, three Equifax executives were caught inside trading. It really shows how Equifax cares about other people!\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "# The plain English request we will make of PaLM 2\n", "prompt = (\n", @@ -1465,42 +540,20 @@ }, { "cell_type": "code", - "execution_count": 38, + "execution_count": null, "metadata": { "id": "mL5P0_3X04dE" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 66e3af22-91cb-400a-92c3-69e7cd12ee01 is DONE. 0 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "from bigframes.ml.llm import PaLM2TextGenerator\n", "\n", -<<<<<<< HEAD "q_a_model = PaLM2TextGenerator()" -======= - "# Create a BigQuery Cloud resource connection\n", - "CONN_NAME = \"bqdf-llm\"\n", - "session = bf.get_global_session()\n", - "\n", - "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", - "q_a_model = PaLM2TextGenerator(session=session, connection_name=connection)" ->>>>>>> origin/lmm-kmeans-notebook ] }, { "cell_type": "code", - "execution_count": 39, + "execution_count": null, "metadata": { "id": "ICWHsqAW1FNk" }, @@ -1512,58 +565,11 @@ }, { "cell_type": "code", - "execution_count": 40, + "execution_count": null, "metadata": { "id": "gB7e1LXU1pst" }, - "outputs": [ - { - "data": { - "text/html": [ - "Query job 653add17-29be-408c-8882-064217f8556e is DONE. 0 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job 8fd16954-853a-45fd-80bc-65b1242429e2 is DONE. 8 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/html": [ - "Query job d9929bcb-26ce-4844-b68e-f4a980b90ede is DONE. 171 Bytes processed. Open Job" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "' The first comment list is about people complaining about companies or services, while the second comment list is about people reporting identity theft or fraud.'" - ] - }, - "execution_count": 40, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Send the request for PaLM 2 to generate a response to our prompt\n", "major_difference = q_a_model.predict(df)\n", @@ -1585,11 +591,7 @@ "source": [ "# Summary and next steps\n", "\n", -<<<<<<< HEAD "You've used the ML and LLM capabilities of BigQuery DataFrames to help analyze and understand a large dataset of unstructured feedback.\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", ->>>>>>> origin/lmm-kmeans-notebook "\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)." ] From 81125f9505ad98e89939769a8e1fcf30518705f0 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Thu, 16 Nov 2023 15:58:14 -0800 Subject: [PATCH 168/197] feat: send warnings on LLM prediction partial failures (#216) 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/ml/llm.py | 78 ++++++++++++++++++++++++++++----------------- 1 file changed, 49 insertions(+), 29 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 93e2ba825f4..78f3369daf6 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -17,6 +17,7 @@ from __future__ import annotations from typing import cast, Literal, Optional, Union +import warnings import bigframes from bigframes import clients, constants @@ -24,15 +25,22 @@ from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd -_REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT = "text-bison" -_REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT = "text-bison-32k" -_TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" +_TEXT_GENERATOR_BISON_ENDPOINT = "text-bison" +_TEXT_GENERATOR_BISON_32K_ENDPOINT = "text-bison-32k" +_TEXT_GENERATOR_ENDPOINTS = ( + _TEXT_GENERATOR_BISON_ENDPOINT, + _TEXT_GENERATOR_BISON_32K_ENDPOINT, +) -_REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT = "textembedding-gecko" -_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT = ( - "textembedding-gecko-multilingual" +_EMBEDDING_GENERATOR_GECKO_ENDPOINT = "textembedding-gecko" +_EMBEDDING_GENERATOR_GECKO_MULTILINGUAL_ENDPOINT = "textembedding-gecko-multilingual" +_EMBEDDING_GENERATOR_ENDPOINTS = ( + _EMBEDDING_GENERATOR_GECKO_ENDPOINT, + _EMBEDDING_GENERATOR_GECKO_MULTILINGUAL_ENDPOINT, ) -_EMBED_TEXT_RESULT_COLUMN = "text_embedding" + +_ML_GENERATE_TEXT_STATUS = "ml_generate_text_status" +_ML_EMBED_TEXT_STATUS = "ml_embed_text_status" class PaLM2TextGenerator(base.Predictor): @@ -90,18 +98,16 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - if self.model_name == _REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT: - options = { - "endpoint": _REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT, - } - elif self.model_name == _REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT: - options = { - "endpoint": _REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT, - } - else: + + if self.model_name not in _TEXT_GENERATOR_ENDPOINTS: raise ValueError( - f"Model name {self.model_name} is not supported. We only support {_REMOTE_TEXT_GENERATOR_MODEL_ENDPOINT} and {_REMOTE_TEXT_GENERATOR_32K_MODEL_ENDPOINT}." + f"Model name {self.model_name} is not supported. We only support {', '.join(_TEXT_GENERATOR_ENDPOINTS)}." ) + + options = { + "endpoint": self.model_name, + } + return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -182,7 +188,16 @@ def predict( "top_p": top_p, "flatten_json_output": True, } - return self._bqml_model.generate_text(X, options) + + df = self._bqml_model.generate_text(X, options) + + if (df[_ML_GENERATE_TEXT_STATUS] != "").any(): + warnings.warn( + f"Some predictions failed. Check column {_ML_GENERATE_TEXT_STATUS} for detailed status. You may want to filter the failed rows and retry.", + RuntimeWarning, + ) + + return df class PaLM2TextEmbeddingGenerator(base.Predictor): @@ -241,19 +256,15 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - if self.model_name == "textembedding-gecko": - options = { - "endpoint": _REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT, - } - elif self.model_name == _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT: - options = { - "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT, - } - else: + + if self.model_name not in _EMBEDDING_GENERATOR_ENDPOINTS: raise ValueError( - f"Model name {self.model_name} is not supported. We only support {_REMOTE_EMBEDDING_GENERATOR_MODEL_ENDPOINT} and {_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_ENDPOINT}." + f"Model name {self.model_name} is not supported. We only support {', '.join(_EMBEDDING_GENERATOR_ENDPOINTS)}." ) + options = { + "endpoint": self.model_name, + } return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -284,4 +295,13 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: options = { "flatten_json_output": True, } - return self._bqml_model.generate_text_embedding(X, options) + + df = self._bqml_model.generate_text_embedding(X, options) + + if (df[_ML_EMBED_TEXT_STATUS] != "").any(): + warnings.warn( + f"Some predictions failed. Check column {_ML_EMBED_TEXT_STATUS} for detailed status. You may want to filter the failed rows and retry.", + RuntimeWarning, + ) + + return df From 52dfad281def82548751a276ce42b087dbb09f9a Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 17 Nov 2023 21:42:14 +0000 Subject: [PATCH 169/197] docs: code samples for `Series.where` and `Series.mask` (#217) 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) - [x] Appropriate docs were updated - `Series.where`: https://screenshot.googleplex.com/9XWHpMnwrzVPF9G - `Series.mask`: https://screenshot.googleplex.com/4cPvvzoVaVzoCDD Fixes internal issue 310981880 🦕 --- .../bigframes_vendored/pandas/core/series.py | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index c6d98075f57..01175dc0ef5 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1696,6 +1696,49 @@ def kurt(self): def where(self, cond, other): """Replace values where the condition is False. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([10, 11, 12, 13, 14]) + >>> s + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + dtype: Int64 + + You can filter the values in the Series based on a condition. The values + matching the condition would be kept, and not matching would be replaced. + The default replacement value is ``NA``. + + >>> s.where(s % 2 == 0) + 0 10 + 1 + 2 12 + 3 + 4 14 + dtype: Int64 + + You can specify a custom replacement value for non-matching values. + + >>> s.where(s % 2 == 0, -1) + 0 10 + 1 -1 + 2 12 + 3 -1 + 4 14 + dtype: Int64 + >>> s.where(s % 2 == 0, 100*s) + 0 10 + 1 1100 + 2 12 + 3 1300 + 4 14 + dtype: Int64 + Args: cond (bool Series/DataFrame, array-like, or callable): Where cond is True, keep the original value. Where False, replace @@ -1720,6 +1763,77 @@ def where(self, cond, other): def mask(self, cond, other): """Replace values where the condition is True. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([10, 11, 12, 13, 14]) + >>> s + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + dtype: Int64 + + You can mask the values in the Series based on a condition. The values + matching the condition would be masked. + + >>> s.mask(s % 2 == 0) + 0 + 1 11 + 2 + 3 13 + 4 + dtype: Int64 + + You can specify a custom mask value. + + >>> s.mask(s % 2 == 0, -1) + 0 -1 + 1 11 + 2 -1 + 3 13 + 4 -1 + dtype: Int64 + >>> s.mask(s % 2 == 0, 100*s) + 0 1000 + 1 11 + 2 1200 + 3 13 + 4 1400 + dtype: Int64 + + You can also use a remote function to evaluate the mask condition. This + is useful in situation such as the following, where the mask + condition is evaluated based on a complicated business logic which cannot + be expressed in form of a Series. + + >>> @bpd.remote_function([str], bool, reuse=False) + ... def should_mask(name): + ... hash = 0 + ... for char_ in name: + ... hash += ord(char_) + ... return hash % 2 == 0 + + >>> s = bpd.Series(["Alice", "Bob", "Caroline"]) + >>> s + 0 Alice + 1 Bob + 2 Caroline + dtype: string + >>> s.mask(should_mask) + 0 + 1 Bob + 2 Caroline + dtype: string + >>> s.mask(should_mask, "REDACTED") + 0 REDACTED + 1 Bob + 2 Caroline + dtype: string + Args: cond (bool Series/DataFrame, array-like, or callable): Where cond is False, keep the original value. Where True, replace From a18d40e808ee0822d21715cc3e8f794c418aeebc Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 17 Nov 2023 14:42:15 -0800 Subject: [PATCH 170/197] fix: avoid unnecessary row_number() on sort key for io (#211) 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/core/__init__.py | 12 ++++++++---- bigframes/core/compile/compiled.py | 29 ++++++++++++++++++++--------- bigframes/dataframe.py | 12 ++++-------- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index b476961bdc6..e19fec8f3f7 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -125,14 +125,18 @@ def to_sql( col_id_overrides: typing.Mapping[str, str] = {}, sorted: bool = False, ) -> str: - if sorted or offset_column: - return self._compile_ordered().to_sql( - offset_column=offset_column, + array_value = self + if offset_column: + array_value = self.promote_offsets(offset_column) + if sorted: + return array_value._compile_ordered().to_sql( col_id_overrides=col_id_overrides, sorted=sorted, ) else: - return self._compile_unordered().to_sql(col_id_overrides=col_id_overrides) + return array_value._compile_unordered().to_sql( + col_id_overrides=col_id_overrides + ) def start_query( self, diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index 78050ed4f0c..461c2c005a6 100644 --- a/bigframes/core/compile/compiled.py +++ b/bigframes/core/compile/compiled.py @@ -1031,31 +1031,42 @@ def _reproject_to_table(self) -> OrderedIR: def to_sql( self, - 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 = ibis_bigquery.Backend().compile( self._to_ibis_expr( - ordering_mode="offset_col" - if (offset_column or sorted) - else "unordered", - order_col_name=offsets_id, + ordering_mode="unordered", col_id_overrides=col_id_overrides, + expose_hidden_cols=sorted, ) ) if sorted: + output_columns = [ + col_id_overrides.get(col) if (col in col_id_overrides) else col + for col in self.column_ids + ] + selection = ", ".join(map(lambda col_id: f"`{col_id}`", output_columns)) + order_by_clause = self._ordering_clause(self._ordering.all_ordering_columns) + sql = textwrap.dedent( - f"SELECT * EXCEPT (`{offsets_id}`)\n" + f"SELECT {selection}\n" "FROM (\n" f"{sql}\n" ")\n" - f"ORDER BY `{offsets_id}`\n" + f"{order_by_clause}\n" ) return typing.cast(str, sql) + def _ordering_clause(self, ordering: Iterable[OrderingColumnReference]) -> str: + parts = [] + for col_ref in ordering: + asc_desc = "ASC" if col_ref.direction.is_ascending else "DESC" + null_clause = "NULLS LAST" if col_ref.na_last else "NULLS FIRST" + part = f"`{col_ref.column_id}` {asc_desc} {null_clause}" + parts.append(part) + return f"ORDER BY {' ,'.join(parts)}" + def _to_ibis_expr( self, *, diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 57b4ca42cfc..1f1275e217b 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2577,14 +2577,10 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str: } if ordering_id is not None: - return array_value.to_sql( - offset_column=ordering_id, - col_id_overrides=id_overrides, - ) - else: - return array_value.to_sql( - col_id_overrides=id_overrides, - ) + array_value = array_value.promote_offsets(ordering_id) + return array_value.to_sql( + col_id_overrides=id_overrides, + ) def _run_io_query( self, From 010486c3494e05d714da6cc7d51514518d9ae1ea Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Fri, 17 Nov 2023 15:38:14 -0800 Subject: [PATCH 171/197] docs: add code samples for df reshaping, function, merge, and join methods (#203) 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 #310245117 --
[310245117](https://b.corp.google.com/issues/310245117) 🦕 --- .../bigframes_vendored/pandas/core/frame.py | 218 +++++++++++++++++- 1 file changed, 217 insertions(+), 1 deletion(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index b35d0f3b2e8..8033c064d7b 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2121,6 +2121,59 @@ def groupby( used to group large amounts of data and compute operations on these groups. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'Animal': ['Falcon', 'Falcon', + ... 'Parrot', 'Parrot'], + ... 'Max Speed': [380., 370., 24., 26.]}) + >>> df + Animal Max Speed + 0 Falcon 380.0 + 1 Falcon 370.0 + 2 Parrot 24.0 + 3 Parrot 26.0 + + [4 rows x 2 columns] + + >>> df.groupby(['Animal'])['Max Speed'].mean() + Animal + Falcon 375.0 + Parrot 25.0 + Name: Max Speed, dtype: Float64 + + We can also choose to include NA in group keys or not by setting `dropna`: + + >>> df = bpd.DataFrame([[1, 2, 3],[1, None, 4], [2, 1, 3], [1, 2, 2]], + ... columns=["a", "b", "c"]) + >>> df.groupby(by=["b"]).sum() + a c + b + 1.0 2 3 + 2.0 2 5 + + [2 rows x 2 columns] + + >>> df.groupby(by=["b"], dropna=False).sum() + a c + b + 1.0 2 3 + 2.0 2 5 + 1 4 + + [3 rows x 2 columns] + + We can also choose to return object with group labels or not by setting `as_index`: + + >>> df.groupby(by=["b"], as_index=False).sum() + b a c + 0 1.0 2 3 + 1 2.0 2 5 + + [2 rows x 3 columns] + Args: by (str, Sequence[str]): A label or list of labels may be passed to group by the columns @@ -2224,7 +2277,7 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: Python function wrapped by ``remote_function`` decorator, returns a single value from a single value. na_action (Optional[str], default None): - ``{None, 'ignore'}``, default None. If ‘ignore’, propagate NaN + ``{None, 'ignore'}``, default None. If `ignore`, propagate NaN values, without passing them to func. Returns: @@ -2240,6 +2293,74 @@ def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame: Join columns with `other` DataFrame on index + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Join two DataFrames by specifying how to handle the operation: + + >>> df1 = bpd.DataFrame({'col1': ['foo', 'bar'], 'col2': [1, 2]}, index=[10, 11]) + >>> df1 + col1 col2 + 10 foo 1 + 11 bar 2 + + [2 rows x 2 columns] + + >>> df2 = bpd.DataFrame({'col3': ['foo', 'baz'], 'col4': [3, 4]}, index=[11, 22]) + >>> df2 + col3 col4 + 11 foo 3 + 22 baz 4 + + [2 rows x 2 columns] + + >>> df1.join(df2) + col1 col2 col3 col4 + 10 foo 1 + 11 bar 2 foo 3 + + [2 rows x 4 columns] + + >>> df1.join(df2, how="left") + col1 col2 col3 col4 + 10 foo 1 + 11 bar 2 foo 3 + + [2 rows x 4 columns] + + >>> df1.join(df2, how="right") + col1 col2 col3 col4 + 11 bar 2 foo 3 + 22 baz 4 + + [2 rows x 4 columns] + + >>> df1.join(df2, how="outer") + col1 col2 col3 col4 + 10 foo 1 + 11 bar 2 foo 3 + 22 baz 4 + + [3 rows x 4 columns] + + >>> df1.join(df2, how="inner") + col1 col2 col3 col4 + 11 bar 2 foo 3 + + [1 rows x 4 columns] + + + Another option to join using the key columns is to use the on parameter: + + >>> df1.join(df2, on="col1", how="right") + col1 col2 col3 col4 + 11 foo 3 + 22 baz 4 + + [2 rows x 4 columns] + Args: other: DataFrame with an Index similar to the Index of this one. @@ -2292,6 +2413,78 @@ def merge( rows will be matched against each other. This is different from usual SQL join behaviour and can lead to unexpected results. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + Merge DataFrames df1 and df2 by specifiying type of merge: + + >>> df1 = bpd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]}) + >>> df1 + a b + 0 foo 1 + 1 bar 2 + + [2 rows x 2 columns] + + >>> df2 = bpd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]}) + >>> df2 + a c + 0 foo 3 + 1 baz 4 + + [2 rows x 2 columns] + + >>> df1.merge(df2, how="inner", on="a") + a b c + 0 foo 1 3 + + [1 rows x 3 columns] + + >>> df1.merge(df2, how='left', on='a') + a b c + 0 foo 1 3 + 1 bar 2 + + [2 rows x 3 columns] + + Merge df1 and df2 on the lkey and rkey columns. The value columns have + the default suffixes, _x and _y, appended. + + >>> df1 = bpd.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [1, 2, 3, 5]}) + >>> df1 + lkey value + 0 foo 1 + 1 bar 2 + 2 baz 3 + 3 foo 5 + + [4 rows x 2 columns] + + >>> df2 = bpd.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [5, 6, 7, 8]}) + >>> df2 + rkey value + 0 foo 5 + 1 bar 6 + 2 baz 7 + 3 foo 8 + + [4 rows x 2 columns] + + >>> df1.merge(df2, left_on='lkey', right_on='rkey') + lkey value_x rkey value_y + 0 foo 1 foo 5 + 1 foo 1 foo 8 + 2 bar 2 bar 6 + 3 baz 3 baz 7 + 4 foo 5 foo 5 + 5 foo 5 foo 8 + + [6 rows x 4 columns] + Args: right: Object to merge with. @@ -2342,6 +2535,29 @@ def apply(self, func, *, args=(), **kwargs): the DataFrame's index (``axis=0``) the final return type is inferred from the return type of the applied function. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df + col1 col2 + 0 1 3 + 1 2 4 + + [2 rows x 2 columns] + + >>> def sqaure(x): + ... return x * x + >>> df1 = df.apply(sqaure) + >>> df + col1 col2 + 0 1 3 + 1 2 4 + + [2 rows x 2 columns] + Args: func (function): Function to apply to each column or row. From dd78acb174545ba292776a642afcec46f8ee4a2a Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Mon, 20 Nov 2023 10:36:14 -0800 Subject: [PATCH 172/197] fix: dedup special character (#209) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: dedup special character * 🦉 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 --- tests/unit/core/test_bf_utils.py | 6 +++--- third_party/bigframes_vendored/pandas/io/common.py | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/unit/core/test_bf_utils.py b/tests/unit/core/test_bf_utils.py index fc34f35d9c2..10ce1fd09e3 100644 --- a/tests/unit/core/test_bf_utils.py +++ b/tests/unit/core/test_bf_utils.py @@ -25,7 +25,7 @@ def test_get_standardized_ids_columns(): "0", utils.UNNAMED_COLUMN_ID, "duplicate", - "duplicate.1", + "duplicate_1", "with_space", ] assert idx_ids == [] @@ -37,13 +37,13 @@ def test_get_standardized_ids_indexes(): col_ids, idx_ids = utils.get_standardized_ids(col_labels, idx_labels) - assert col_ids == ["duplicate.2"] + assert col_ids == ["duplicate_2"] assert idx_ids == [ "string", "0", utils.UNNAMED_INDEX_ID, "duplicate", - "duplicate.1", + "duplicate_1", "with_space", ] diff --git a/third_party/bigframes_vendored/pandas/io/common.py b/third_party/bigframes_vendored/pandas/io/common.py index 506984e64da..e186f02b5bd 100644 --- a/third_party/bigframes_vendored/pandas/io/common.py +++ b/third_party/bigframes_vendored/pandas/io/common.py @@ -13,13 +13,13 @@ def dedup_names( """ 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. + Currently the renaming is done by appending a underscore 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'] + ['x', 'y', 'x_1', 'x_2'] ``` """ names = list(names) # so we can index @@ -34,9 +34,9 @@ def dedup_names( if is_potential_multiindex: # for mypy assert isinstance(col, tuple) - col = col[:-1] + (f"{col[-1]}.{cur_count}",) + col = col[:-1] + (f"{col[-1]}_{cur_count}",) else: - col = f"{col}.{cur_count}" + col = f"{col}_{cur_count}" cur_count = counts[col] names[i] = col From c88d38e69682f4c620174086b8f16f4780c04811 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Mon, 20 Nov 2023 23:02:15 +0000 Subject: [PATCH 173/197] docs: add code samples for `index` and `column` properties (#212) 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) - [x] Appropriate docs were updated - `DataFrame.columns`: https://screenshot.googleplex.com/3Bwdb482FBfEsi2 - `DataFrame.index`: https://screenshot.googleplex.com/4iJymH3FxMn8Hhb - `Series.index`: https://screenshot.googleplex.com/7MXQcuASbQ3c8s5 Fixes internal issue 310260952 🦕 --- .../bigframes_vendored/pandas/core/frame.py | 79 ++++++++++++++++++- .../bigframes_vendored/pandas/core/series.py | 49 +++++++++++- 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 8033c064d7b..f448ad79394 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -3106,6 +3106,47 @@ def index(self): index is used for label-based access and alignment, and can be accessed or modified using this attribute. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can access the index of a DataFrame via ``index`` property. + + >>> df = bpd.DataFrame({'Name': ['Alice', 'Bob', 'Aritra'], + ... 'Age': [25, 30, 35], + ... 'Location': ['Seattle', 'New York', 'Kona']}, + ... index=([10, 20, 30])) + >>> df + Name Age Location + 10 Alice 25 Seattle + 20 Bob 30 New York + 30 Aritra 35 Kona + + [3 rows x 3 columns] + >>> df.index # doctest: +ELLIPSIS + + >>> df.index.values + array([10, 20, 30], dtype=object) + + Let's try setting a new index for the dataframe and see that reflect via + ``index`` property. + + >>> df1 = df.set_index(["Name", "Location"]) + >>> df1 + Age + Name Location + Alice Seattle 25 + Bob New York 30 + Aritra Kona 35 + + [3 rows x 1 columns] + >>> df1.index # doctest: +ELLIPSIS + + >>> df1.index.values + array([('Alice', 'Seattle'), ('Bob', 'New York'), ('Aritra', 'Kona')], + dtype=object) + Returns: The index labels of the DataFrame. """ @@ -3113,7 +3154,43 @@ def index(self): @property def columns(self): - "The column labels of the DataFrame." + """The column labels of the DataFrame. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can access the column labels of a DataFrame via ``columns`` property. + + >>> df = bpd.DataFrame({'Name': ['Alice', 'Bob', 'Aritra'], + ... 'Age': [25, 30, 35], + ... 'Location': ['Seattle', 'New York', 'Kona']}, + ... index=([10, 20, 30])) + >>> df + Name Age Location + 10 Alice 25 Seattle + 20 Bob 30 New York + 30 Aritra 35 Kona + + [3 rows x 3 columns] + >>> df.columns + Index(['Name', 'Age', 'Location'], dtype='object') + + You can also set new labels for columns. + + >>> df.columns = ["NewName", "NewAge", "NewLocation"] + >>> df + NewName NewAge NewLocation + 10 Alice 25 Seattle + 20 Bob 30 New York + 30 Aritra 35 Kona + + [3 rows x 3 columns] + >>> df.columns + Index(['NewName', 'NewAge', 'NewLocation'], dtype='object') + + """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def value_counts( diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 01175dc0ef5..a86765a4125 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -44,7 +44,54 @@ def struct(self): @property def index(self): - """The index (axis labels) of the Series.""" + """The index (axis labels) of the Series. + + The index of a Series is used to label and identify each element of the + underlying data. The index can be thought of as an immutable ordered set + (technically a multi-set, as it may contain duplicate labels), and is + used to index and align data. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + You can access the index of a Series via ``index`` property. + + >>> df = bpd.DataFrame({'Name': ['Alice', 'Bob', 'Aritra'], + ... 'Age': [25, 30, 35], + ... 'Location': ['Seattle', 'New York', 'Kona']}, + ... index=([10, 20, 30])) + >>> s = df["Age"] + >>> s + 10 25 + 20 30 + 30 35 + Name: Age, dtype: Int64 + >>> s.index # doctest: +ELLIPSIS + + >>> s.index.values + array([10, 20, 30], dtype=object) + + Let's try setting a multi-index case reflect via ``index`` property. + + >>> df1 = df.set_index(["Name", "Location"]) + >>> s1 = df1["Age"] + >>> s1 + Name Location + Alice Seattle 25 + Bob New York 30 + Aritra Kona 35 + Name: Age, dtype: Int64 + >>> s1.index # doctest: +ELLIPSIS + + >>> s1.index.values + array([('Alice', 'Seattle'), ('Bob', 'New York'), ('Aritra', 'Kona')], + dtype=object) + + Returns: + The index labels of the Series. + """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @property From 1d1477158b5a8e84d099e07c6f566182a1abd7fc Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 21 Nov 2023 16:16:25 +0000 Subject: [PATCH 174/197] test: re-enable `system_prerelease` tests (#221) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: re-enable `system_prerelease` tests * exclude ibis from prerelease install list * install explicit 6.2.0 version for ibis in pre prelease * add unit_prerelease to pre and post submit e2e tests * Update noxfile.py * 🦉 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: Tim Swast Co-authored-by: Owl Bot --- .kokoro/continuous/e2e.cfg | 2 +- .kokoro/presubmit/e2e.cfg | 2 +- noxfile.py | 50 ++++++++++++++++++++++++++++++++------ 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/.kokoro/continuous/e2e.cfg b/.kokoro/continuous/e2e.cfg index d875f360603..2f93a582124 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_noextras e2e notebook samples" + value: "unit_prerelease system_prerelease system_noextras e2e notebook samples" } diff --git a/.kokoro/presubmit/e2e.cfg b/.kokoro/presubmit/e2e.cfg index d875f360603..2f93a582124 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_noextras e2e notebook samples" + value: "unit_prerelease system_prerelease system_noextras e2e notebook samples" } diff --git a/noxfile.py b/noxfile.py index da61232fc79..8d6d641fc1e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -494,6 +494,11 @@ def prerelease(session: nox.sessions.Session, tests_path): CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" ) + # Ignore officially released versions of certain packages specified in + # testing/constraints-*.txt and install a more recent, pre-release versions + # directly + already_installed = set() + # PyArrow prerelease packages are published to an alternative PyPI host. # https://arrow.apache.org/docs/python/install.html#installing-nightly-packages session.install( @@ -504,6 +509,8 @@ def prerelease(session: nox.sessions.Session, tests_path): "--upgrade", "pyarrow", ) + already_installed.add("pyarrow") + session.install( "--extra-index-url", "https://pypi.anaconda.org/scipy-wheels-nightly/simple", @@ -512,16 +519,48 @@ def prerelease(session: nox.sessions.Session, tests_path): "--upgrade", "pandas", ) + already_installed.add("pandas") + + # TODO(shobs): + # Commit https://github.com/ibis-project/ibis/commit/c20ba7feab6bdea6c299721310e04dbc10551cc2 + # introduced breaking change that removed the following: + # ibis.expr.rules.column + # ibis.expr.rules.value + # ibis.expr.rules.any + # Let's exclude ibis head from prerelease install list for now. Instead, use + # a working ibis-framework version resolved via setup.by (currently resolves + # to version 6.2.0 due to version requirement "6.2.0,<7.0.0dev"). + # We should enable the head back once bigframes support a version that + # includes the above commit. + # session.install( + # "--upgrade", + # "-e", # Use -e so that py.typed file is included. + # "git+https://github.com/ibis-project/ibis.git#egg=ibis-framework", + # ) + session.install("--no-deps", "ibis-framework==6.2.0") + already_installed.add("ibis-framework") + + # Workaround https://github.com/googleapis/python-db-dtypes-pandas/issues/178 + session.install("--no-deps", "db-dtypes") + already_installed.add("db-dtypes") + + # Ensure we catch breaking changes in the client libraries early. + session.install( + "--upgrade", + "-e", + "git+https://github.com/googleapis/python-bigquery.git#egg=google-cloud-bigquery", + ) + already_installed.add("google-cloud-bigquery") session.install( "--upgrade", - "-e", # Use -e so that py.typed file is included. - "git+https://github.com/ibis-project/ibis.git#egg=ibis-framework", + "-e", + "git+https://github.com/googleapis/python-bigquery-storage.git#egg=google-cloud-bigquery-storage", ) - # Workaround https://github.com/googleapis/python-db-dtypes-pandas/issues/178 - session.install("--no-deps", "db-dtypes") + already_installed.add("google-cloud-bigquery-storage") # Workaround to install pandas-gbq >=0.15.0, which is required by test only. session.install("--no-deps", "pandas-gbq") + already_installed.add("pandas-gbq") session.install( *set(UNIT_TEST_STANDARD_DEPENDENCIES + SYSTEM_TEST_STANDARD_DEPENDENCIES), @@ -541,9 +580,6 @@ def prerelease(session: nox.sessions.Session, tests_path): 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( From ed8876d3439a3b45b65e8789737c3c2e3a7f1adb Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 21 Nov 2023 14:02:31 -0800 Subject: [PATCH 175/197] feat: add the recent api method for ML component (#225) * feat: add the recent api method for ML component --- bigframes/ml/cluster.py | 2 ++ bigframes/ml/compose.py | 2 ++ bigframes/ml/decomposition.py | 2 ++ bigframes/ml/ensemble.py | 5 +++++ bigframes/ml/forecasting.py | 2 ++ bigframes/ml/imported.py | 3 +++ bigframes/ml/linear_model.py | 3 +++ bigframes/ml/llm.py | 4 +++- bigframes/ml/pipeline.py | 2 ++ bigframes/ml/preprocessing.py | 7 +++++++ tests/unit/session/test_io_bigquery.py | 3 +++ 11 files changed, 34 insertions(+), 1 deletion(-) diff --git a/bigframes/ml/cluster.py b/bigframes/ml/cluster.py index c9f52ba0b6b..6b79d356a23 100644 --- a/bigframes/ml/cluster.py +++ b/bigframes/ml/cluster.py @@ -22,11 +22,13 @@ from google.cloud import bigquery import bigframes +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.cluster._kmeans +@log_adapter.class_logger class KMeans( base.UnsupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.cluster._kmeans.KMeans, diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index bf046ff6914..ace876dd2df 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -22,6 +22,7 @@ from typing import List, Optional, Tuple, Union from bigframes import constants +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, preprocessing, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.compose._column_transformer @@ -36,6 +37,7 @@ ] +@log_adapter.class_logger class ColumnTransformer( base.Transformer, third_party.bigframes_vendored.sklearn.compose._column_transformer.ColumnTransformer, diff --git a/bigframes/ml/decomposition.py b/bigframes/ml/decomposition.py index 7cda7a69933..ef777cb33ac 100644 --- a/bigframes/ml/decomposition.py +++ b/bigframes/ml/decomposition.py @@ -22,11 +22,13 @@ from google.cloud import bigquery import bigframes +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.decomposition._pca +@log_adapter.class_logger class PCA( base.UnsupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.decomposition._pca.PCA, diff --git a/bigframes/ml/ensemble.py b/bigframes/ml/ensemble.py index fcb3fe53438..1cc9fb37395 100644 --- a/bigframes/ml/ensemble.py +++ b/bigframes/ml/ensemble.py @@ -22,6 +22,7 @@ from google.cloud import bigquery import bigframes +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.ensemble._forest @@ -47,6 +48,7 @@ } +@log_adapter.class_logger class XGBRegressor( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.xgboost.sklearn.XGBRegressor, @@ -202,6 +204,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor: return new_model.session.read_gbq_model(model_name) +@log_adapter.class_logger class XGBClassifier( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.xgboost.sklearn.XGBClassifier, @@ -356,6 +359,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier: return new_model.session.read_gbq_model(model_name) +@log_adapter.class_logger class RandomForestRegressor( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor, @@ -521,6 +525,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegresso return new_model.session.read_gbq_model(model_name) +@log_adapter.class_logger class RandomForestClassifier( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier, diff --git a/bigframes/ml/forecasting.py b/bigframes/ml/forecasting.py index cf23854fa01..995201062bc 100644 --- a/bigframes/ml/forecasting.py +++ b/bigframes/ml/forecasting.py @@ -21,10 +21,12 @@ from google.cloud import bigquery import bigframes +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd +@log_adapter.class_logger class ARIMAPlus(base.SupervisedTrainablePredictor): """Time Series ARIMA Plus model.""" diff --git a/bigframes/ml/imported.py b/bigframes/ml/imported.py index f6afc9aa382..4ae0a8ea4d0 100644 --- a/bigframes/ml/imported.py +++ b/bigframes/ml/imported.py @@ -21,10 +21,12 @@ from google.cloud import bigquery import bigframes +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd +@log_adapter.class_logger class TensorFlowModel(base.Predictor): """Imported TensorFlow model. @@ -101,6 +103,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> TensorFlowModel: return new_model.session.read_gbq_model(model_name) +@log_adapter.class_logger class ONNXModel(base.Predictor): """Imported Open Neural Network Exchange (ONNX) model. diff --git a/bigframes/ml/linear_model.py b/bigframes/ml/linear_model.py index 433d9fbc380..5ee87b88505 100644 --- a/bigframes/ml/linear_model.py +++ b/bigframes/ml/linear_model.py @@ -23,6 +23,7 @@ import bigframes import bigframes.constants as constants +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.linear_model._base @@ -46,6 +47,7 @@ } +@log_adapter.class_logger class LinearRegression( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.linear_model._base.LinearRegression, @@ -178,6 +180,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression: return new_model.session.read_gbq_model(model_name) +@log_adapter.class_logger class LogisticRegression( base.SupervisedTrainablePredictor, third_party.bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression, diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 78f3369daf6..5beb54a32d4 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.core import blocks, log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd @@ -43,6 +43,7 @@ _ML_EMBED_TEXT_STATUS = "ml_embed_text_status" +@log_adapter.class_logger class PaLM2TextGenerator(base.Predictor): """PaLM2 text generator LLM model. @@ -200,6 +201,7 @@ def predict( return df +@log_adapter.class_logger class PaLM2TextEmbeddingGenerator(base.Predictor): """PaLM2 text embedding generator LLM model. diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index ad0b3fae111..4ae2bfe5552 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -24,11 +24,13 @@ import bigframes import bigframes.constants as constants +from bigframes.core import log_adapter from bigframes.ml import base, compose, forecasting, loader, preprocessing, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.pipeline +@log_adapter.class_logger class Pipeline( base.BaseEstimator, third_party.bigframes_vendored.sklearn.pipeline.Pipeline, diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 5f44d402184..a403e57e715 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -20,6 +20,7 @@ import typing from typing import Any, cast, List, Literal, Optional, Tuple, Union +from bigframes.core import log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd import third_party.bigframes_vendored.sklearn.preprocessing._data @@ -28,6 +29,7 @@ import third_party.bigframes_vendored.sklearn.preprocessing._label +@log_adapter.class_logger class StandardScaler( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._data.StandardScaler, @@ -111,6 +113,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +@log_adapter.class_logger class MaxAbsScaler( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._data.MaxAbsScaler, @@ -194,6 +197,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +@log_adapter.class_logger class MinMaxScaler( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._data.MinMaxScaler, @@ -277,6 +281,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +@log_adapter.class_logger class KBinsDiscretizer( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._discretization.KBinsDiscretizer, @@ -395,6 +400,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +@log_adapter.class_logger class OneHotEncoder( base.Transformer, third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder, @@ -524,6 +530,7 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: ) +@log_adapter.class_logger class LabelEncoder( base.LabelTransformer, third_party.bigframes_vendored.sklearn.preprocessing._label.LabelEncoder, diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index e1481d3f05c..c87835c4126 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -59,6 +59,7 @@ def test_create_job_configs_labels_length_limit_not_met(): def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit(): + log_adapter.get_and_reset_api_methods() cur_labels = { "bigframes-api": "read_pandas", "source": "bigquery-dataframes-temp", @@ -87,6 +88,7 @@ def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit(): def test_create_job_configs_labels_length_limit_met_and_labels_is_none(): + log_adapter.get_and_reset_api_methods() df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]}) # Test running methods more than the labels' length limit for i in range(66): @@ -102,6 +104,7 @@ def test_create_job_configs_labels_length_limit_met_and_labels_is_none(): def test_create_job_configs_labels_length_limit_met(): + log_adapter.get_and_reset_api_methods() cur_labels = { "bigframes-api": "read_pandas", "source": "bigquery-dataframes-temp", From d7957fad071d223ef8f6fb8f3de395c865ff60aa Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Tue, 21 Nov 2023 21:26:52 -0800 Subject: [PATCH 176/197] docs: code samples for dataframe.any, dataframe.all and dataframe.prod (#223) * docs: code samples for dataframe.any, dataframe.all and dataframe.prod * Update examples * update example output --- .../bigframes_vendored/pandas/core/frame.py | 85 ++++++++++++++++++- 1 file changed, 84 insertions(+), 1 deletion(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index f448ad79394..b771be30418 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2584,6 +2584,33 @@ def any(self, *, axis=0, bool_only: bool = False): along a Dataframe axis that is True or equivalent (e.g. non-zero or non-empty). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [True, True], "B": [False, False]}) + >>> df + A B + 0 True False + 1 True False + + [2 rows x 2 columns] + + Checking if each column contains at least one True element(the default behavior without an explicit axis parameter). + + >>> df.any() + A True + B False + dtype: boolean + + Checking if each row contains at least one True element. + + >>> df.any(axis=1) + 0 True + 1 True + dtype: boolean + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2604,6 +2631,33 @@ def all(self, axis=0, *, bool_only: bool = False): along a DataFrame axis that is False or equivalent (e.g. zero or empty). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [True, True], "B": [False, False]}) + >>> df + A B + 0 True False + 1 True False + + [2 rows x 2 columns] + + Checking if all values in each column are True(the default behavior without an explicit axis parameter). + + >>> df.all() + A True + B False + dtype: boolean + + Checking across rows to see if all values are True. + + >>> df.all(axis=1) + 0 False + 1 False + dtype: boolean + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2620,8 +2674,37 @@ def prod(self, axis=0, *, numeric_only: bool = False): """ Return the product of the values over the requested axis. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 2, 3], "B": [4.5, 5.5, 6.5]}) + >>> df + A B + 0 1 4.5 + 1 2 5.5 + 2 3 6.5 + + [3 rows x 2 columns] + + Calculating the product of each column(the default behavior without an explicit axis parameter). + + >>> df.prod() + A 6.0 + B 160.875 + dtype: Float64 + + Calculating the product of each row. + + >>> df.prod(axis=1) + 0 4.5 + 1 11.0 + 2 19.5 + dtype: Float64 + Args: - aßxis ({index (0), columns (1)}): + 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): From 71844b03cdbfe684320c186a0488c8c7fb4fcd6e Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 22 Nov 2023 23:46:14 +0000 Subject: [PATCH 177/197] docs: make the code samples reflect default bq connection usage (#206) 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 internal issue 305136837 🦕 --- .../getting_started_bq_dataframes.ipynb | 4 +- .../remote_functions/remote_function.ipynb | 514 ++++++++---------- samples/snippets/remote_function.py | 6 +- 3 files changed, 241 insertions(+), 283 deletions(-) diff --git a/notebooks/getting_started/getting_started_bq_dataframes.ipynb b/notebooks/getting_started/getting_started_bq_dataframes.ipynb index 6cc6acc9935..18be5e48fd1 100644 --- a/notebooks/getting_started/getting_started_bq_dataframes.ipynb +++ b/notebooks/getting_started/getting_started_bq_dataframes.ipynb @@ -802,7 +802,7 @@ "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", + "> 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-default-connection`. 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", @@ -817,7 +817,7 @@ }, "outputs": [], "source": [ - "@bf.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", + "@bf.remote_function([float], str)\n", "def get_bucket(num):\n", " if not num: return \"NA\"\n", " boundary = 4000\n", diff --git a/notebooks/remote_functions/remote_function.ipynb b/notebooks/remote_functions/remote_function.ipynb index 06be0e7293e..063c1738b4d 100644 --- a/notebooks/remote_functions/remote_function.ipynb +++ b/notebooks/remote_functions/remote_function.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 2, + "execution_count": 19, "id": "3613b1cd", "metadata": {}, "outputs": [], @@ -16,24 +16,16 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 20, "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 25.4 s, sys: 2.5 s, total: 27.9 s\n", - "Wall time: 2min 31s\n" + "CPU times: user 2.34 s, sys: 307 ms, total: 2.65 s\n", + "Wall time: 17.8 s\n" ] }, { @@ -141,7 +133,7 @@ "9 154 Sure, but what about a solution using O(1) mem... 8" ] }, - "execution_count": 3, + "execution_count": 20, "metadata": {}, "output_type": "execute_result" } @@ -160,7 +152,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 21, "id": "fd8a04a3", "metadata": {}, "outputs": [], @@ -191,7 +183,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 22, "id": "2b5e4568", "metadata": {}, "outputs": [ @@ -199,8 +191,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 4.22 s, sys: 18.2 ms, total: 4.24 s\n", - "Wall time: 4.26 s\n" + "CPU times: user 3.32 s, sys: 0 ns, total: 3.32 s\n", + "Wall time: 3.32 s\n" ] }, { @@ -319,7 +311,7 @@ "9 154 Sure, but what about a solution using O(1) mem... 8 19" ] }, - "execution_count": 5, + "execution_count": 22, "metadata": {}, "output_type": "execute_result" } @@ -333,65 +325,25 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 23, "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. " ] }, "metadata": {}, @@ -399,13 +351,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "09706700e8dd4cf39f65a0d58371c1eb", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job a283cb39-41b1-44cd-a6c3-f2a2c6a55b25 is DONE. 17.2 GB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job 4c1d9d3e-be25-4818-b74d-6214164d99ab is DONE. 0 Bytes processed. " ] }, "metadata": {}, @@ -440,62 +390,62 @@ " \n", " \n", " 0\n", - " 11012908\n", - " you're welcome! according to the docs it shoul...\n", + " 11231597\n", + " In your update, why are some of the system fun...\n", " 0\n", " \n", " \n", " 1\n", - " 11013760\n", - " You *should* be concerned with the disk being ...\n", - " 0\n", + " 49684807\n", + " what you have tried so far . ??\n", + " 1\n", " \n", " \n", " 2\n", - " 11013784\n", - " have you looked at `Integrate` or `NIntegrate`?\n", + " 7623925\n", + " @Michael: It should work. Perhaps you looked i...\n", " 0\n", " \n", " \n", " 3\n", - " 11015512\n", - " sorry, is a typo. The variable name is dist. (...\n", + " 34046685\n", + " Will it work with SQL compact? Please excuse m...\n", " 0\n", " \n", " \n", " 4\n", - " 11016238\n", - " Pfff, I'm having trouble with that formula too...\n", + " 6426146\n", + " do you know the equation to your pdf?\n", " 0\n", " \n", " \n", " 5\n", - " 11016276\n", - " Thanks thinksteep! Does this mean that by usin...\n", + " 60686114\n", + " m sorry but at least you have to think about it.\n", " 0\n", " \n", " \n", " 6\n", - " 11016551\n", - " Jason, thanks for the reply. I've been workin...\n", + " 16631986\n", + " i think also making disable this by only jquer...\n", " 0\n", " \n", " \n", " 7\n", - " 11017973\n", - " I assume an `off` of 0.5 would put be exactly ...\n", + " 16498565\n", + " I am including these files on my header of the...\n", " 0\n", " \n", " \n", " 8\n", - " 11018225\n", - " Thank you very much. I do worry too much abou...\n", + " 26601001\n", + " wrong answer, you didn't understand the logic\n", " 0\n", " \n", " \n", " 9\n", - " 11018370\n", - " @IanClelland, I edited my question a bit. The ...\n", + " 73255842\n", + " Call the setOnClickListener before return row.\n", " 0\n", " \n", " \n", @@ -505,21 +455,21 @@ ], "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", + "0 11231597 In your update, why are some of the system fun... 0\n", + "1 49684807 what you have tried so far . ?? 1\n", + "2 7623925 @Michael: It should work. Perhaps you looked i... 0\n", + "3 34046685 Will it work with SQL compact? Please excuse m... 0\n", + "4 6426146 do you know the equation to your pdf? 0\n", + "5 60686114 m sorry but at least you have to think about it. 0\n", + "6 16631986 i think also making disable this by only jquer... 0\n", + "7 16498565 I am including these files on my header of the... 0\n", + "8 26601001 wrong answer, you didn't understand the logic 0\n", + "9 73255842 Call the setOnClickListener before return row. 0\n", "\n", "[10 rows x 3 columns]" ] }, - "execution_count": 6, + "execution_count": 23, "metadata": {}, "output_type": "execute_result" } @@ -539,7 +489,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 24, "id": "55ed241e", "metadata": {}, "outputs": [ @@ -549,8 +499,9 @@ "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", + "remote_function(input_types: 'List[type]', output_type: 'type', dataset: 'Optional[str]' = None, bigquery_connection: 'Optional[str]' = None, reuse: 'bool' = True, name: 'Optional[str]' = None, packages: 'Optional[Sequence[str]]' = None)\n", + " Decorator to turn a user defined function into a BigQuery remote function. Check out\n", + " the code samples at: https://cloud.google.com/bigquery/docs/remote-functions#bigquery-dataframes.\n", " \n", " .. note::\n", " Please make sure following is setup before using this API:\n", @@ -576,7 +527,7 @@ " * 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", + " * Service Account User (roles/iam.serviceAccountUser) on the service account `PROJECT_NUMBER-compute@developer.gserviceaccount.com`\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", @@ -602,15 +553,25 @@ " 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", + " 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", + " function and corresponding cloud function (if any) that was\n", + " 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", + " name (str, Optional):\n", + " Explicit name of the persisted BigQuery remote function. Use it with\n", + " caution, because two users working in the same project and dataset\n", + " could overwrite each other's remote functions if they use the same\n", + " persistent name.\n", + " packages (str[], Optional):\n", + " Explicit name of the external package dependencies. Each dependency\n", + " is added to the `requirements.txt` as is, and can be of the form\n", + " supported in https://pip.pypa.io/en/stable/reference/requirements-file-format/.\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", @@ -631,49 +592,16 @@ }, { "cell_type": "code", - "execution_count": 8, - "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": 9, + "execution_count": 25, "id": "fbc27f81", "metadata": {}, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "[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" - ] - }, { "name": "stdout", "output_type": "stream", "text": [ "\n", - "Wall time: 89.0601 s\n" + "Wall time: 76.2628 s\n" ] } ], @@ -684,7 +612,7 @@ "\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", + "@pd.remote_function([int], int, reuse=False)\n", "def nth_prime(n):\n", " prime_numbers = [2,3]\n", " i=3\n", @@ -712,7 +640,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 26, "id": "c1c9355f", "metadata": {}, "outputs": [ @@ -720,33 +648,17 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 16.8 ms, sys: 61 µs, total: 16.8 ms\n", - "Wall time: 17 ms\n" + "CPU times: user 55.8 ms, sys: 182 µs, total: 56 ms\n", + "Wall time: 54.5 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. Open Job" + ], "text/plain": [ - "HTML(value='Query job 4f8d5734-8070-4630-8a59-c05a31d60476 is RUNNING. " ] }, "metadata": {}, @@ -754,13 +666,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "116d6ef3d6b247d3aaafef5fe6b970de", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job c0a2c187-364d-4978-97bc-30352828f624 is DONE. 17.2 GB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job ec057f9e-726b-44f0-a5c0-24c05c7ecfeb is RUNNING. " ] }, "metadata": {}, @@ -796,71 +706,71 @@ " \n", " \n", " 0\n", - " 11012908\n", - " you're welcome! according to the docs it shoul...\n", + " 11231597\n", + " In your update, why are some of the system fun...\n", " 0\n", " -1\n", " \n", " \n", " 1\n", - " 11013760\n", - " You *should* be concerned with the disk being ...\n", - " 0\n", - " -1\n", + " 49684807\n", + " what you have tried so far . ??\n", + " 1\n", + " 2\n", " \n", " \n", " 2\n", - " 11013784\n", - " have you looked at `Integrate` or `NIntegrate`?\n", + " 7623925\n", + " @Michael: It should work. Perhaps you looked i...\n", " 0\n", " -1\n", " \n", " \n", " 3\n", - " 11015512\n", - " sorry, is a typo. The variable name is dist. (...\n", + " 34046685\n", + " Will it work with SQL compact? Please excuse m...\n", " 0\n", " -1\n", " \n", " \n", " 4\n", - " 11016238\n", - " Pfff, I'm having trouble with that formula too...\n", + " 6426146\n", + " do you know the equation to your pdf?\n", " 0\n", " -1\n", " \n", " \n", " 5\n", - " 11016276\n", - " Thanks thinksteep! Does this mean that by usin...\n", + " 60686114\n", + " m sorry but at least you have to think about it.\n", " 0\n", " -1\n", " \n", " \n", " 6\n", - " 11016551\n", - " Jason, thanks for the reply. I've been workin...\n", + " 16631986\n", + " i think also making disable this by only jquer...\n", " 0\n", " -1\n", " \n", " \n", " 7\n", - " 11017973\n", - " I assume an `off` of 0.5 would put be exactly ...\n", + " 16498565\n", + " I am including these files on my header of the...\n", " 0\n", " -1\n", " \n", " \n", " 8\n", - " 11018225\n", - " Thank you very much. I do worry too much abou...\n", + " 26601001\n", + " wrong answer, you didn't understand the logic\n", " 0\n", " -1\n", " \n", " \n", " 9\n", - " 11018370\n", - " @IanClelland, I edited my question a bit. The ...\n", + " 73255842\n", + " Call the setOnClickListener before return row.\n", " 0\n", " -1\n", " \n", @@ -871,21 +781,21 @@ ], "text/plain": [ " id text score n_prime\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", + "0 11231597 In your update, why are some of the system fun... 0 -1\n", + "1 49684807 what you have tried so far . ?? 1 2\n", + "2 7623925 @Michael: It should work. Perhaps you looked i... 0 -1\n", + "3 34046685 Will it work with SQL compact? Please excuse m... 0 -1\n", + "4 6426146 do you know the equation to your pdf? 0 -1\n", + "5 60686114 m sorry but at least you have to think about it. 0 -1\n", + "6 16631986 i think also making disable this by only jquer... 0 -1\n", + "7 16498565 I am including these files on my header of the... 0 -1\n", + "8 26601001 wrong answer, you didn't understand the logic 0 -1\n", + "9 73255842 Call the setOnClickListener before return row. 0 -1\n", "\n", "[10 rows x 4 columns]" ] }, - "execution_count": 10, + "execution_count": 26, "metadata": {}, "output_type": "execute_result" } @@ -900,7 +810,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 27, "id": "2701cb81", "metadata": {}, "outputs": [ @@ -908,8 +818,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "bigframes-dev.bigframes_temp_us.bigframes_b0feb1fbaf8188b64d7e70118d93c5d4\n", - "projects/bigframes-dev/locations/us-central1/functions/bigframes-b0feb1fbaf8188b64d7e70118d93c5d4\n" + "shobs-test.bigframes_temp_us.bigframes_343b7b4bb93ca8747dae20c22bdaec8b_p27heyce\n", + "projects/shobs-test/locations/us-central1/functions/bigframes-343b7b4bb93ca8747dae20c22bdaec8b-p27heyce\n" ] } ], @@ -922,7 +832,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 28, "id": "920fa18e", "metadata": {}, "outputs": [ @@ -937,6 +847,42 @@ " \n", " Then it can be applied to a DataFrame or Series.\n", " \n", + " .. note::\n", + " The return type of the function must be explicitly specified in the\n", + " function's original definition even if not otherwise required.\n", + " \n", + " BigQuery Utils provides many public functions under the ``bqutil`` project on Google Cloud Platform project\n", + " (See: https://github.com/GoogleCloudPlatform/bigquery-utils/tree/master/udfs#using-the-udfs).\n", + " You can checkout Community UDFs to use community-contributed functions.\n", + " (See: https://github.com/GoogleCloudPlatform/bigquery-utils/tree/master/udfs/community#community-udfs).\n", + " \n", + " **Examples:**\n", + " \n", + " Use the ``cw_lower_case_ascii_only`` function from Community UDFs.\n", + " (https://github.com/GoogleCloudPlatform/bigquery-utils/blob/master/udfs/community/cw_lower_case_ascii_only.sqlx)\n", + " \n", + " >>> import bigframes.pandas as bpd\n", + " >>> bpd.options.display.progress_bar = None\n", + " \n", + " >>> df = bpd.DataFrame({'id': [1, 2, 3], 'name': ['AURÉLIE', 'CÉLESTINE', 'DAPHNÉ']})\n", + " >>> df\n", + " id name\n", + " 0 1 AURÉLIE\n", + " 1 2 CÉLESTINE\n", + " 2 3 DAPHNÉ\n", + " \n", + " [3 rows x 2 columns]\n", + " \n", + " >>> func = bpd.read_gbq_function(\"bqutil.fn.cw_lower_case_ascii_only\")\n", + " >>> df1 = df.assign(new_name=df['name'].apply(func))\n", + " >>> df1\n", + " id name new_name\n", + " 0 1 AURÉLIE aurÉlie\n", + " 1 2 CÉLESTINE cÉlestine\n", + " 2 3 DAPHNÉ daphnÉ\n", + " \n", + " [3 rows x 3 columns]\n", + " \n", " Args:\n", " function_name (str):\n", " the function's name in BigQuery in the format\n", @@ -965,7 +911,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 29, "id": "a6c9da0a", "metadata": {}, "outputs": [], @@ -978,7 +924,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 30, "id": "d7e7de7f", "metadata": {}, "outputs": [ @@ -986,33 +932,17 @@ "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" + "CPU times: user 70.8 ms, sys: 3.49 ms, total: 74.3 ms\n", + "Wall time: 75.2 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. Open Job" + ], "text/plain": [ - "HTML(value='Query job 02e3bf43-a387-41c7-85c7-4a5366251de7 is RUNNING. " ] }, "metadata": {}, @@ -1020,13 +950,11 @@ }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "faf93766ce1e489183c86a9daf5ce7d1", - "version_major": 2, - "version_minor": 0 - }, + "text/html": [ + "Query job 4d3da7ed-42e6-4b2b-b656-ac9ef6d2e871 is DONE. 17.2 GB processed. Open Job" + ], "text/plain": [ - "HTML(value='Query job fa4329e8-2918-44c4-96c5-d8591364abc9 is RUNNING. " ] }, "metadata": {}, @@ -1063,80 +991,80 @@ " \n", " \n", " 0\n", - " 11012908\n", - " you're welcome! according to the docs it shoul...\n", + " 11231597\n", + " In your update, why are some of the system fun...\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 1\n", - " 11013760\n", - " You *should* be concerned with the disk being ...\n", - " 0\n", - " -1\n", - " -1\n", + " 49684807\n", + " what you have tried so far . ??\n", + " 1\n", + " 2\n", + " 2\n", " \n", " \n", " 2\n", - " 11013784\n", - " have you looked at `Integrate` or `NIntegrate`?\n", + " 7623925\n", + " @Michael: It should work. Perhaps you looked i...\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 3\n", - " 11015512\n", - " sorry, is a typo. The variable name is dist. (...\n", + " 34046685\n", + " Will it work with SQL compact? Please excuse m...\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 4\n", - " 11016238\n", - " Pfff, I'm having trouble with that formula too...\n", + " 6426146\n", + " do you know the equation to your pdf?\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 5\n", - " 11016276\n", - " Thanks thinksteep! Does this mean that by usin...\n", + " 60686114\n", + " m sorry but at least you have to think about it.\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 6\n", - " 11016551\n", - " Jason, thanks for the reply. I've been workin...\n", + " 16631986\n", + " i think also making disable this by only jquer...\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 7\n", - " 11017973\n", - " I assume an `off` of 0.5 would put be exactly ...\n", + " 16498565\n", + " I am including these files on my header of the...\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 8\n", - " 11018225\n", - " Thank you very much. I do worry too much abou...\n", + " 26601001\n", + " wrong answer, you didn't understand the logic\n", " 0\n", " -1\n", " -1\n", " \n", " \n", " 9\n", - " 11018370\n", - " @IanClelland, I edited my question a bit. The ...\n", + " 73255842\n", + " Call the setOnClickListener before return row.\n", " 0\n", " -1\n", " -1\n", @@ -1148,20 +1076,20 @@ ], "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", + "0 11231597 In your update, why are some of the system fun... 0 \n", + "1 49684807 what you have tried so far . ?? 1 \n", + "2 7623925 @Michael: It should work. Perhaps you looked i... 0 \n", + "3 34046685 Will it work with SQL compact? Please excuse m... 0 \n", + "4 6426146 do you know the equation to your pdf? 0 \n", + "5 60686114 m sorry but at least you have to think about it. 0 \n", + "6 16631986 i think also making disable this by only jquer... 0 \n", + "7 16498565 I am including these files on my header of the... 0 \n", + "8 26601001 wrong answer, you didn't understand the logic 0 \n", + "9 73255842 Call the setOnClickListener before return row. 0 \n", "\n", " n_prime n_prime_again \n", "0 -1 -1 \n", - "1 -1 -1 \n", + "1 2 2 \n", "2 -1 -1 \n", "3 -1 -1 \n", "4 -1 -1 \n", @@ -1174,7 +1102,7 @@ "[10 rows x 5 columns]" ] }, - "execution_count": 15, + "execution_count": 30, "metadata": {}, "output_type": "execute_result" } @@ -1186,6 +1114,38 @@ "df = df.assign(n_prime_again=df['score'].apply(nth_prime_existing))\n", "df.head(10)" ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "bafab950", + "metadata": {}, + "outputs": [], + "source": [ + "# Clean up GCP assets created as part of bigframes remote_function\n", + "def cleanup_remote_function_assets(remote_udf, ignore_failures=False):\n", + " \"\"\"Clean up the GCP assets behind a bigframes remote function.\"\"\"\n", + "\n", + " session = pd.get_global_session()\n", + "\n", + " # Clean up BQ remote function\n", + " try:\n", + " session.bqclient.delete_routine(remote_udf.bigframes_remote_function)\n", + " except Exception:\n", + " # By default don't raise exception in cleanup\n", + " if not ignore_failures:\n", + " raise\n", + "\n", + " # Clean up cloud function\n", + " try:\n", + " session.cloudfunctionsclient.delete_function(name=remote_udf.bigframes_cloud_function)\n", + " except Exception:\n", + " # By default don't raise exception in cleanup\n", + " if not ignore_failures:\n", + " raise\n", + "\n", + "cleanup_remote_function_assets(nth_prime)" + ] } ], "metadata": { diff --git a/samples/snippets/remote_function.py b/samples/snippets/remote_function.py index 646d7b0c307..61b7dc092ad 100644 --- a/samples/snippets/remote_function.py +++ b/samples/snippets/remote_function.py @@ -38,8 +38,8 @@ def run_remote_function_and_read_gbq_function(project_id: str): # 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`. We will also set `reuse=False` to make sure we don't + # examples we will be letting the default connection `bigframes-default-connection` + # be used. We will also set `reuse=False` to make sure we don't # step over someone else creating remote function in the same project from # the exact same source code at the same time. Let's try a `pandas`-like use # case in which we want to apply a user defined scalar function to every @@ -49,7 +49,6 @@ def run_remote_function_and_read_gbq_function(project_id: str): @bpd.remote_function( [float], str, - bigquery_connection="bigframes-rf-conn", reuse=False, ) def get_bucket(num): @@ -94,7 +93,6 @@ def get_bucket(num): @bpd.remote_function( [str], str, - bigquery_connection="bigframes-rf-conn", reuse=False, packages=["cryptography"], ) From 3a375e87b64b8fb51370bfec8f2cfdbcd8fe960a Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Wed, 22 Nov 2023 16:52:22 -0800 Subject: [PATCH 178/197] docs: add examples for dataframe.min, dataframe.max and dataframe.sum (#227) * docs: add examples for dataframe.min, dataframe.max and dataframe.sum * update spacing --- .../bigframes_vendored/pandas/core/frame.py | 87 ++++++++++++++++++- 1 file changed, 84 insertions(+), 3 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index b771be30418..e41ac905aa7 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2597,7 +2597,7 @@ def any(self, *, axis=0, bool_only: bool = False): [2 rows x 2 columns] - Checking if each column contains at least one True element(the default behavior without an explicit axis parameter). + Checking if each column contains at least one True element (the default behavior without an explicit axis parameter). >>> df.any() A True @@ -2644,7 +2644,7 @@ def all(self, axis=0, *, bool_only: bool = False): [2 rows x 2 columns] - Checking if all values in each column are True(the default behavior without an explicit axis parameter). + Checking if all values in each column are True (the default behavior without an explicit axis parameter). >>> df.all() A True @@ -2688,7 +2688,7 @@ def prod(self, axis=0, *, numeric_only: bool = False): [3 rows x 2 columns] - Calculating the product of each column(the default behavior without an explicit axis parameter). + Calculating the product of each column (the default behavior without an explicit axis parameter). >>> df.prod() A 6.0 @@ -2721,6 +2721,33 @@ def min(self, axis=0, *, numeric_only: bool = False): If you want the *index* of the minimum, use ``idxmin``. This is the equivalent of the ``numpy.ndarray`` method ``argmin``. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Finding the minimum value in each column (the default behavior without an explicit axis parameter). + + >>> df.min() + A 1.0 + B 2.0 + dtype: Float64 + + Finding the minimum value in each row. + + >>> df.min(axis=1) + 0 1.0 + 1 3.0 + dtype: Float64 + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2739,6 +2766,33 @@ def max(self, axis=0, *, numeric_only: bool = False): If you want the *index* of the maximum, use ``idxmax``. This is the equivalent of the ``numpy.ndarray`` method ``argmax``. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Finding the maximum value in each column (the default behavior without an explicit axis parameter). + + >>> df.max() + A 3.0 + B 4.0 + dtype: Float64 + + Finding the maximum value in each row. + + >>> df.max(axis=1) + 0 2.0 + 1 4.0 + dtype: Float64 + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2756,6 +2810,33 @@ def sum(self, axis=0, *, numeric_only: bool = False): This is equivalent to the method ``numpy.sum``. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Calculating the sum of each column (the default behavior without an explicit axis parameter). + + >>> df.sum() + A 4.0 + B 6.0 + dtype: Float64 + + Calculating the sum of each row. + + >>> df.sum(axis=1) + 0 3.0 + 1 7.0 + dtype: Float64 + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. From b62a07a95cd60f995a48825c9874822d0eb02483 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 24 Nov 2023 00:10:18 +0000 Subject: [PATCH 179/197] docs: code samples for `Series.dot` and `DataFrame.dot` (#226) --- bigframes/dataframe.py | 3 +- bigframes/operations/base.py | 2 +- tests/system/small/test_dataframe.py | 23 ++++++ .../bigframes_vendored/pandas/core/frame.py | 71 +++++++++++++++++++ .../bigframes_vendored/pandas/core/series.py | 15 ++++ 5 files changed, 112 insertions(+), 2 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 1f1275e217b..8567296e29e 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2797,7 +2797,8 @@ def get_right_id(id): result = result[other_frame.columns] if isinstance(other, bf_series.Series): - result = result[other.name].rename() + # There should be exactly one column in the result + result = result[result.columns[0]].rename() return result diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index d33befe4da6..85ce1dd9e6d 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -141,7 +141,7 @@ def _apply_binary_op( if isinstance(other, pd.Series): # TODO: Convert to BigQuery DataFrames series raise NotImplementedError( - f"Pandas series not supported supported as operand. {constants.FEEDBACK_LINK}" + f"Pandas series not supported as operand. {constants.FEEDBACK_LINK}" ) if isinstance(other, series.Series): (left, right, block) = self._align(other, how=alignment) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a0cf25807cf..e25e9ce501c 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3493,6 +3493,29 @@ def test_df_dot_operator( ) +def test_df_dot_series_inline(): + left = [[1, 2, 3], [2, 5, 7]] + right = [2, 1, 3] + + bf1 = dataframe.DataFrame(left) + bf2 = series.Series(right) + bf_result = bf1.dot(bf2).to_pandas() + + df1 = pd.DataFrame(left) + df2 = pd.Series(right) + pd_result = df1.dot(df2) + + # Patch pandas dtypes for testing parity + # Pandas result is int64 instead of Int64 (nullable) dtype. + pd_result = pd_result.astype(pd.Int64Dtype()) + pd_result.index = pd_result.index.astype(pd.Int64Dtype()) + + pd.testing.assert_series_equal( + bf_result, + pd_result, + ) + + def test_df_dot_series( matrix_2by3_df, matrix_2by3_pandas_df, matrix_3by4_df, matrix_3by4_pandas_df ): diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index e41ac905aa7..a1aac5d2b54 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -3485,6 +3485,77 @@ def dot(self, other): The dot method for Series computes the inner product, instead of the matrix product here. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> left = bpd.DataFrame([[0, 1, -2, -1], [1, 1, 1, 1]]) + >>> left + 0 1 2 3 + 0 0 1 -2 -1 + 1 1 1 1 1 + + [2 rows x 4 columns] + >>> right = bpd.DataFrame([[0, 1], [1, 2], [-1, -1], [2, 0]]) + >>> right + 0 1 + 0 0 1 + 1 1 2 + 2 -1 -1 + 3 2 0 + + [4 rows x 2 columns] + >>> left.dot(right) + 0 1 + 0 1 4 + 1 2 2 + + [2 rows x 2 columns] + + You can also use the operator ``@`` for the dot product: + + >>> left @ right + 0 1 + 0 1 4 + 1 2 2 + + [2 rows x 2 columns] + + The right input can be a Series, in which case the result will also be a + Series: + + >>> right = bpd.Series([1, 2, -1,0]) + >>> left @ right + 0 4 + 1 2 + dtype: Int64 + + Any user defined index of the left matrix and columns of the right + matrix will reflect in the result. + + >>> left = bpd.DataFrame([[1, 2, 3], [2, 5, 7]], index=["alpha", "beta"]) + >>> left + 0 1 2 + alpha 1 2 3 + beta 2 5 7 + + [2 rows x 3 columns] + >>> right = bpd.DataFrame([[2, 4, 8], [1, 5, 10], [3, 6, 9]], columns=["red", "green", "blue"]) + >>> right + red green blue + 0 2 4 8 + 1 1 5 10 + 2 3 6 9 + + [3 rows x 3 columns] + >>> left.dot(right) + red green blue + alpha 13 32 55 + beta 30 75 129 + + [2 rows x 3 columns] + Args: other (Series or DataFrame): The other object to compute the matrix product with. diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index a86765a4125..1b751ed83b6 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -631,6 +631,21 @@ def dot(self, other) -> Series | np.ndarray: BigQuery Dataframes does not validate this property and will produce incorrect results if indices are not equal. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([0, 1, 2, 3]) + >>> other = bpd.Series([-1, 2, -3, 4]) + >>> s.dot(other) + 8 + + You can also use the operator ``@`` for the dot product: + + >>> s @ other + 8 + Args: other (Series): The other object to compute the dot product with its columns. From f9c6e727e2b901310bb5301da449d616ea85e135 Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Tue, 28 Nov 2023 10:48:21 -0800 Subject: [PATCH 180/197] docs: add examples for dataframe.kurt, dataframe.std, dataframe.count (#232) * docs: add examples for dataframe.kurt, dataframe.std, dataframe.count * update count example * update count example * update examples * update . to : --- .../bigframes_vendored/pandas/core/frame.py | 96 +++++++++++++++++-- 1 file changed, 87 insertions(+), 9 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index a1aac5d2b54..a7018ed3a24 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2597,14 +2597,14 @@ def any(self, *, axis=0, bool_only: bool = False): [2 rows x 2 columns] - Checking if each column contains at least one True element (the default behavior without an explicit axis parameter). + Checking if each column contains at least one True element(the default behavior without an explicit axis parameter): >>> df.any() A True B False dtype: boolean - Checking if each row contains at least one True element. + Checking if each row contains at least one True element: >>> df.any(axis=1) 0 True @@ -2644,14 +2644,14 @@ def all(self, axis=0, *, bool_only: bool = False): [2 rows x 2 columns] - Checking if all values in each column are True (the default behavior without an explicit axis parameter). + Checking if all values in each column are True(the default behavior without an explicit axis parameter): >>> df.all() A True B False dtype: boolean - Checking across rows to see if all values are True. + Checking across rows to see if all values are True: >>> df.all(axis=1) 0 False @@ -2688,14 +2688,14 @@ def prod(self, axis=0, *, numeric_only: bool = False): [3 rows x 2 columns] - Calculating the product of each column (the default behavior without an explicit axis parameter). + Calculating the product of each column(the default behavior without an explicit axis parameter): >>> df.prod() A 6.0 B 160.875 dtype: Float64 - Calculating the product of each row. + Calculating the product of each row: >>> df.prod(axis=1) 0 4.5 @@ -2911,11 +2911,37 @@ 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. + """Return unbiased kurtosis over columns. Kurtosis obtained using Fisher's definition of kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 2, 3, 4, 5], + ... "B": [3, 4, 3, 2, 1], + ... "C": [2, 2, 3, 2, 2]}) + >>> df + A B C + 0 1 3 2 + 1 2 4 2 + 2 3 3 3 + 3 4 2 2 + 4 5 1 2 + + [5 rows x 3 columns] + + Calculating the kurtosis value of each column: + + >>> df.kurt() + A -1.2 + B -0.177515 + C 5.0 + dtype: Float64 + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -2926,10 +2952,36 @@ def kurt(self, *, numeric_only: bool = False): raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def std(self, *, numeric_only: bool = False): - """Return sample standard deviation over requested axis. + """Return sample standard deviation over columns. Normalized by N-1 by default. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 2, 3, 4, 5], + ... "B": [3, 4, 3, 2, 1], + ... "C": [2, 2, 3, 2, 2]}) + >>> df + A B C + 0 1 3 2 + 1 2 4 2 + 2 3 3 3 + 3 4 2 2 + 4 5 1 2 + + [5 rows x 3 columns] + + Calculating the standard deviation of each column: + + >>> df.std() + A 1.581139 + B 1.140175 + C 0.447214 + dtype: Float64 + Args: numeric_only (bool. default False): Default False. Include only float, int, boolean columns. @@ -2941,11 +2993,37 @@ def std(self, *, numeric_only: bool = False): def count(self, *, numeric_only: bool = False): """ - Count non-NA cells for each column or row. + Count non-NA cells for each column. The values `None`, `NaN`, `NaT`, and optionally `numpy.inf` (depending on `pandas.options.mode.use_inf_as_na`) are considered NA. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, None, 3, 4, 5], + ... "B": [1, 2, 3, 4, 5], + ... "C": [None, 3.5, None, 4.5, 5.0]}) + >>> df + A B C + 0 1.0 1 + 1 2 3.5 + 2 3.0 3 + 3 4.0 4 4.5 + 4 5.0 5 5.0 + + [5 rows x 3 columns] + + Counting non-NA values for each column: + + >>> df.count() + A 4.0 + B 5.0 + C 3.0 + dtype: Float64 + Args: numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. From edd0522747eadb74780124fb18ed7face251441d Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Tue, 28 Nov 2023 11:40:28 -0800 Subject: [PATCH 181/197] =?UTF-8?q?docs:=20add=20examples=20for=20datafram?= =?UTF-8?q?e.mean,=20dataframe.median,=20dataframe.va=E2=80=A6=20(#228)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: add examples for dataframe.mean, dataframe.median, dataframe.var and dataframe.skew * column to columns * update var example --- .../bigframes_vendored/pandas/core/frame.py | 105 +++++++++++++++++- 1 file changed, 103 insertions(+), 2 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index a7018ed3a24..2a8972f2e58 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -2852,6 +2852,33 @@ def sum(self, axis=0, *, numeric_only: bool = False): def mean(self, axis=0, *, numeric_only: bool = False): """Return the mean of the values over the requested axis. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Calculating the mean of each column (the default behavior without an explicit axis parameter). + + >>> df.mean() + A 2.0 + B 3.0 + dtype: Float64 + + Calculating the mean of each row. + + >>> df.mean(axis=1) + 0 1.5 + 1 3.5 + dtype: Float64 + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2865,7 +2892,27 @@ def mean(self, axis=0, *, numeric_only: bool = False): 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. + """Return the median of the values over colunms. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Finding the median value of each column. + + >>> df.median() + A 1.0 + B 2.0 + dtype: Float64 Args: numeric_only (bool. default False): @@ -2884,6 +2931,34 @@ def var(self, axis=0, *, numeric_only: bool = False): Normalized by N-1 by default. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [1, 3], "B": [2, 4]}) + >>> df + A B + 0 1 2 + 1 3 4 + + [2 rows x 2 columns] + + Calculating the variance of each column (the default behavior without an explicit axis parameter). + + >>> df.var() + A 2.0 + B 2.0 + dtype: Float64 + + Calculating the variance of each row. + + >>> df.var(axis=1) + 0 0.5 + 1 0.5 + dtype: Float64 + + Args: axis ({index (0), columns (1)}): Axis for the function to be applied on. @@ -2897,10 +2972,36 @@ def var(self, axis=0, *, numeric_only: bool = False): raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def skew(self, *, numeric_only: bool = False): - """Return unbiased skew over requested axis. + """Return unbiased skew over columns. Normalized by N-1. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'A': [1, 2, 3, 4, 5], + ... 'B': [5, 4, 3, 2, 1], + ... 'C': [2, 2, 3, 2, 2]}) + >>> df + A B C + 0 1 5 2 + 1 2 4 2 + 2 3 3 3 + 3 4 2 2 + 4 5 1 2 + + [5 rows x 3 columns] + + Calculating the skewness of each column. + + >>> df.skew() + A 0.0 + B 0.0 + C 2.236068 + dtype: Float64 + Args: numeric_only (bool, default False): Include only float, int, boolean columns. From ae03756f5ee45e0e74e0c0bdd4777e018eba2273 Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Tue, 28 Nov 2023 12:53:51 -0800 Subject: [PATCH 182/197] fix: make to_pandas override enable_downsampling when sampling_method is manually set. (#200) * fix: make to_pandas override enable_downsampling when sampling_method is manually set. * fix: make to_pandas override enable_downsampling when sampling_method is manually set. * fix: make to_pandas override enable_downsampling when sampling_method is manually set. --- bigframes/core/blocks.py | 41 ++++++++++++++++------------ tests/system/small/test_dataframe.py | 11 ++++++++ 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index f1113d938ec..34913872e70 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -389,23 +389,6 @@ def to_pandas( ordered: bool = True, ) -> 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, @@ -453,6 +436,28 @@ def _compute_and_count( ) -> 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. + enable_downsampling = ( + True + if sampling_method is not None + else bigframes.options.sampling.enable_downsampling + ) + + max_download_size = ( + max_download_size or bigframes.options.sampling.max_download_size + ) + + random_state = random_state or bigframes.options.sampling.random_state + + if sampling_method is None: + sampling_method = bigframes.options.sampling.sampling_method or _UNIFORM + 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)}." + ) + expr = self._apply_value_keys_to_expr(value_keys=value_keys) results_iterator, query_job = expr.start_query( @@ -469,7 +474,7 @@ def _compute_and_count( ) if fraction < 1: - if not bigframes.options.sampling.enable_downsampling: + if not 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" diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index e25e9ce501c..9b9567418b7 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3546,3 +3546,14 @@ def test_df_dot_operator_series( bf_result, pd_result, ) + + +def test_to_pandas_downsampling_option_override(session): + df = session.read_gbq("bigframes-dev.bigframes_tests_sys.batting") + download_size = 1 + + df = df.to_pandas(max_download_size=download_size, sampling_method="head") + + total_memory_bytes = df.memory_usage(deep=True).sum() + total_memory_mb = total_memory_bytes / (1024 * 1024) + assert total_memory_mb == pytest.approx(download_size, rel=0.3) From 9d6613d318b558722b7bab12773efdea4bbe9931 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 28 Nov 2023 13:46:18 -0800 Subject: [PATCH 183/197] feat: add info and memory_usage methods to dataframe (#219) --- .pre-commit-config.yaml | 2 +- bigframes/_config/display_options.py | 4 + bigframes/core/indexes/index.py | 13 ++- bigframes/dataframe.py | 84 +++++++++++++++++++ bigframes/dtypes.py | 13 +++ noxfile.py | 1 + setup.py | 1 + tests/system/small/test_dataframe.py | 42 ++++++++++ .../pandas/core/config_init.py | 11 +++ .../bigframes_vendored/pandas/core/frame.py | 66 +++++++++++++++ 10 files changed, 235 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6e0fd8b98fb..517176da892 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -38,4 +38,4 @@ repos: rev: v1.1.1 hooks: - id: mypy - additional_dependencies: [types-requests] + additional_dependencies: [types-requests, types-tabulate] diff --git a/bigframes/_config/display_options.py b/bigframes/_config/display_options.py index ad3ea3f68cc..afa36aa84c2 100644 --- a/bigframes/_config/display_options.py +++ b/bigframes/_config/display_options.py @@ -32,6 +32,10 @@ class DisplayOptions: progress_bar: Optional[str] = "auto" repr_mode: Literal["head", "deferred"] = "head" + max_info_columns: int = 100 + max_info_rows: Optional[int] = 200000 + memory_usage: bool = True + @contextlib.contextmanager def pandas_repr(display_options: DisplayOptions): diff --git a/bigframes/core/indexes/index.py b/bigframes/core/indexes/index.py index 6c66c36062a..fc7cf167d49 100644 --- a/bigframes/core/indexes/index.py +++ b/bigframes/core/indexes/index.py @@ -155,6 +155,14 @@ def _block(self) -> blocks.Block: def T(self) -> Index: return self.transpose() + def _memory_usage(self) -> int: + (n_rows,) = self.shape + return sum( + self.dtypes.map( + lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows + ) + ) + def transpose(self) -> Index: return self @@ -326,7 +334,10 @@ def _apply_aggregation(self, op: agg_ops.AggregateOp) -> typing.Any: def __getitem__(self, key: int) -> typing.Any: if isinstance(key, int): - result_pd_df, _ = self._block.slice(key, key + 1, 1).to_pandas() + if key != -1: + result_pd_df, _ = self._block.slice(key, key + 1, 1).to_pandas() + else: # special case, want [-1:] instead of [-1:0] + result_pd_df, _ = self._block.slice(key).to_pandas() if result_pd_df.empty: raise IndexError("single positional indexer is out-of-bounds") return result_pd_df.index[0] diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 8567296e29e..f7796291b9a 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -18,6 +18,7 @@ import datetime import re +import sys import textwrap import typing from typing import ( @@ -36,6 +37,7 @@ import google.cloud.bigquery as bigquery import numpy import pandas +import tabulate import bigframes import bigframes._config.display_options as display_options @@ -350,6 +352,88 @@ def query_job(self) -> Optional[bigquery.QueryJob]: self._set_internal_query_job(self._compute_dry_run()) return self._query_job + def memory_usage(self, index: bool = True): + n_rows, _ = self.shape + # like pandas, treat all variable-size objects as just 8-byte pointers, ignoring actual object + column_sizes = self.dtypes.map( + lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows + ) + if index: + index_size = pandas.Series([self.index._memory_usage()], index=["Index"]) + column_sizes = pandas.concat([index_size, column_sizes]) + return column_sizes + + def info( + self, + verbose: Optional[bool] = None, + buf=None, + max_cols: Optional[int] = None, + memory_usage: Optional[bool] = None, + show_counts: Optional[bool] = None, + ): + obuf = buf or sys.stdout + + n_rows, n_columns = self.shape + + max_cols = ( + max_cols + if max_cols is not None + else bigframes.options.display.max_info_columns + ) + + show_all_columns = verbose if verbose is not None else (n_columns < max_cols) + + obuf.write(f"{type(self)}\n") + + index_type = "MultiIndex" if self.index.nlevels > 1 else "Index" + + # These accessses are kind of expensive, maybe should try to skip? + first_indice = self.index[0] + last_indice = self.index[-1] + obuf.write(f"{index_type}: {n_rows} entries, {first_indice} to {last_indice}\n") + + dtype_strings = self.dtypes.astype("string") + if show_all_columns: + obuf.write(f"Data columns (total {n_columns} columns):\n") + column_info = self.columns.to_frame(name="Column") + + max_rows = bigframes.options.display.max_info_rows + too_many_rows = n_rows > max_rows if max_rows is not None else False + + if show_counts if show_counts is not None else (not too_many_rows): + non_null_counts = self.count().to_pandas() + column_info["Non-Null Count"] = non_null_counts.map( + lambda x: f"{int(x)} non-null" + ) + + column_info["Dtype"] = dtype_strings + + column_info = column_info.reset_index(drop=True) + column_info.index.name = "#" + + column_info_formatted = tabulate.tabulate(column_info, headers="keys") # type: ignore + obuf.write(column_info_formatted) + obuf.write("\n") + + else: # Just number of columns and first, last + obuf.write( + f"Columns: {n_columns} entries, {self.columns[0]} to {self.columns[-1]}\n" + ) + dtype_counts = dtype_strings.value_counts().sort_index(ascending=True).items() + dtype_counts_formatted = ", ".join( + f"{dtype}({count})" for dtype, count in dtype_counts + ) + obuf.write(f"dtypes: {dtype_counts_formatted}\n") + + show_memory = ( + memory_usage + if memory_usage is not None + else bigframes.options.display.memory_usage + ) + if show_memory: + # TODO: Convert to different units (kb, mb, etc.) + obuf.write(f"memory usage: {self.memory_usage().sum()} bytes\n") + def _set_internal_query_job(self, query_job: bigquery.QueryJob): self._query_job = query_job diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index cd35e380c02..774eb74d06a 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -143,6 +143,19 @@ # "string" and "string[pyarrow] are accepted" BIGFRAMES_STRING_TO_BIGFRAMES["string[pyarrow]"] = pd.StringDtype(storage="pyarrow") +# For the purposes of dataframe.memory_usage +# https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_sizes +DTYPE_BYTE_SIZES = { + pd.BooleanDtype(): 1, + pd.Int64Dtype(): 8, + pd.Float32Dtype(): 8, + pd.StringDtype(): 8, + pd.ArrowDtype(pa.time64("us")): 8, + pd.ArrowDtype(pa.timestamp("us")): 8, + pd.ArrowDtype(pa.timestamp("us", tz="UTC")): 8, + pd.ArrowDtype(pa.date32()): 8, +} + def ibis_dtype_to_bigframes_dtype( ibis_dtype: ibis_dtypes.DataType, diff --git a/noxfile.py b/noxfile.py index 8d6d641fc1e..c1fb53f7947 100644 --- a/noxfile.py +++ b/noxfile.py @@ -228,6 +228,7 @@ def mypy(session): "types-python-dateutil", "types-requests", "types-setuptools", + "types-tabulate", ] ) | set(SYSTEM_TEST_STANDARD_DEPENDENCIES) diff --git a/setup.py b/setup.py index 29eacb74a9a..abf165b3df3 100644 --- a/setup.py +++ b/setup.py @@ -50,6 +50,7 @@ "requests >=2.27.1", "scikit-learn >=1.2.2", "sqlalchemy >=1.4,<3.0dev", + "tabulate >= 0.9", "ipywidgets >=7.7.1", "humanize >= 4.6.0", ] diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 9b9567418b7..9744d3f6e91 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import io import operator import tempfile import typing @@ -255,6 +256,47 @@ def test_drop_with_custom_column_labels(scalars_dfs): assert_pandas_df_equal(bf_result, pd_result) +def test_df_memory_usage(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = scalars_pandas_df.memory_usage() + bf_result = scalars_df.memory_usage() + + pd.testing.assert_series_equal(pd_result, bf_result, rtol=1.5) + + +def test_df_info(scalars_dfs): + expected = ( + "\n" + "Index: 9 entries, 0 to 8\n" + "Data columns (total 13 columns):\n" + " # Column Non-Null Count Dtype\n" + "--- ------------- ---------------- ------------------------------\n" + " 0 bool_col 8 non-null boolean\n" + " 1 bytes_col 6 non-null object\n" + " 2 date_col 7 non-null date32[day][pyarrow]\n" + " 3 datetime_col 6 non-null timestamp[us][pyarrow]\n" + " 4 geography_col 4 non-null geometry\n" + " 5 int64_col 8 non-null Int64\n" + " 6 int64_too 9 non-null Int64\n" + " 7 numeric_col 6 non-null object\n" + " 8 float64_col 7 non-null Float64\n" + " 9 rowindex_2 9 non-null Int64\n" + " 10 string_col 8 non-null string\n" + " 11 time_col 6 non-null time64[us][pyarrow]\n" + " 12 timestamp_col 6 non-null timestamp[us, tz=UTC][pyarrow]\n" + "dtypes: Float64(1), Int64(3), boolean(1), date32[day][pyarrow](1), geometry(1), object(2), string(1), time64[us][pyarrow](1), timestamp[us, tz=UTC][pyarrow](1), timestamp[us][pyarrow](1)\n" + "memory usage: 945 bytes\n" + ) + + scalars_df, _ = scalars_dfs + bf_result = io.StringIO() + + scalars_df.info(buf=bf_result) + + assert expected == bf_result.getvalue() + + def test_drop_index(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs diff --git a/third_party/bigframes_vendored/pandas/core/config_init.py b/third_party/bigframes_vendored/pandas/core/config_init.py index 198654015e0..dfb91dfeb81 100644 --- a/third_party/bigframes_vendored/pandas/core/config_init.py +++ b/third_party/bigframes_vendored/pandas/core/config_init.py @@ -33,6 +33,17 @@ Instead estimated bytes processed will be shown. Dataframe and Series objects can still be computed with methods that explicitly execute and download results. + max_info_columns (int): + max_info_columns is used in DataFrame.info method to decide if + per column information will be printed. + max_info_rows (int or None): + df.info() will usually show null-counts for each column. + For large frames this can be quite slow. max_info_rows and max_info_cols + limit this null check only to frames with smaller dimensions than + specified. + memory_usage (bool): + This specifies if the memory usage of a DataFrame should be displayed when + df.info() is called. Valid values True,False, """ sampling_options_doc = """ diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 2a8972f2e58..099d8b8e663 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -92,6 +92,72 @@ def values(self) -> np.ndarray: """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def info( + self, + verbose: bool | None = None, + buf=None, + max_cols: int | None = None, + memory_usage: bool | None = None, + show_counts: bool | None = None, + ) -> None: + """ + Print a concise summary of a DataFrame. + + This method prints information about a DataFrame including + the index dtypeand columns, non-null values and memory usage. + + Args: + verbose (bool, optional): + Whether to print the full summary. By default, the setting in + ``pandas.options.display.max_info_columns`` is followed. + buf (writable buffer, defaults to sys.stdout): + Where to send the output. By default, the output is printed to + sys.stdout. Pass a writable buffer if you need to further process + the output. + max_cols (int, optional): + When to switch from the verbose to the truncated output. If the + DataFrame has more than `max_cols` columns, the truncated output + is used. By default, the setting in + ``pandas.options.display.max_info_columns`` is used. + memory_usage (bool, optional): + Specifies whether total memory usage of the DataFrame + elements (including the index) should be displayed. By default, + this follows the ``pandas.options.display.memory_usage`` setting. + True always show memory usage. False never shows memory usage. + Memory estimation is made based in column dtype and number of rows + assuming values consume the same memory amount for corresponding dtypes. + show_counts (bool, optional): + Whether to show the non-null counts. By default, this is shown + only if the DataFrame is smaller than + ``pandas.options.display.max_info_rows`` and + ``pandas.options.display.max_info_columns``. A value of True always + shows the counts, and False never shows the counts. + + Returns: + None: This method prints a summary of a DataFrame and returns None.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def memory_usage(self, index: bool = True): + """ + Return the memory usage of each column in bytes. + + The memory usage can optionally include the contribution of + the index and elements of `object` dtype. + + This value is displayed in `DataFrame.info` by default. This can be + suppressed by setting ``pandas.options.display.memory_usage`` to False. + + Args: + index (bool, default True): + Specifies whether to include the memory usage of the DataFrame's + index in returned Series. If ``index=True``, the memory usage of + the index is the first item in the output. + + Returns: + Series: A Series whose index is the original column names and whose values is the memory usage of each column in bytes. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # IO methods (to / from other formats) def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarray: From 69b016eae7ea97d84ceeb22ba09f5472841db072 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 28 Nov 2023 23:42:58 +0000 Subject: [PATCH 184/197] fix: use anonymous dataset to create `remote_function` (#205) * fix: use anonymous dataset to create `remote_function` * update README about anonymous dataset instead of bigframes_temp_location * remove dataset creation step from remote function This is because now the dataset is an anonymous dataset that must have been created previously as part of bigframes session creation. * restore create_dataset, guarded by get_dataset --- README.rst | 7 ++-- bigframes/remote_function.py | 19 ++++++--- bigframes/session/__init__.py | 14 ------- tests/system/large/test_remote_function.py | 47 +++++++++++++++++++++- tests/system/small/test_remote_function.py | 37 +++++------------ 5 files changed, 73 insertions(+), 51 deletions(-) diff --git a/README.rst b/README.rst index 5ddb4a7639a..91dac127518 100644 --- a/README.rst +++ b/README.rst @@ -267,10 +267,9 @@ definition. To view and manage connections, do the following: 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: +in a special type of `hidden dataset `__ +referred to as an anonymous dataset. To view and manage remote functions created +in a user provided dataset, do the following: 1. Go to `BigQuery in the Google Cloud Console `__. 2. Select the project in which you created the remote function. diff --git a/bigframes/remote_function.py b/bigframes/remote_function.py index a39cd033f69..7280ac7d425 100644 --- a/bigframes/remote_function.py +++ b/bigframes/remote_function.py @@ -188,6 +188,7 @@ def create_bq_remote_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( @@ -204,14 +205,22 @@ def create_bq_remote_function( logger.info(f"Creating BQ remote function: {create_function_ddl}") - # Make sure the dataset exists + # Make sure the dataset exists. I.e. if it doesn't exist, go ahead and + # create it 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) + try: + # This check does not require bigquery.datasets.create IAM + # permission. So, if the data set already exists, then user can work + # without having that permission. + self._bq_client.get_dataset(dataset) + except google.api_core.exceptions.NotFound: + # This requires bigquery.datasets.create IAM permission + 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. @@ -610,7 +619,7 @@ def get_routine_reference( raise DatasetMissingError dataset_ref = bigquery.DatasetReference( - bigquery_client.project, session._session_dataset_id + bigquery_client.project, session._anonymous_dataset.dataset_id ) return dataset_ref.routine(routine_ref_str) @@ -778,9 +787,7 @@ def remote_function( dataset, default_project=bigquery_client.project ) else: - dataset_ref = bigquery.DatasetReference.from_string( - session._session_dataset_id, default_project=bigquery_client.project - ) + dataset_ref = session._anonymous_dataset bq_location, cloud_function_region = get_remote_function_locations( bigquery_client.location diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 928123ce749..d2f6137883c 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -198,13 +198,6 @@ def cloudfunctionsclient(self): def resourcemanagerclient(self): return self._clients_provider.resourcemanagerclient - @property - def _session_dataset_id(self): - """A dataset for storing temporary objects local to the session - 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 @@ -229,13 +222,6 @@ def _create_bq_datasets(self): query_destination.dataset_id, ) - # 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 - def close(self): """No-op. Temporary resources are deleted after 7 days.""" diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index 6ed3e6511a2..5cb4df188ca 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -22,7 +22,7 @@ import textwrap from google.api_core.exceptions import NotFound, ResourceExhausted -from google.cloud import functions_v2 +from google.cloud import bigquery, functions_v2 import pandas import pytest import test_utils.prefixer @@ -1210,3 +1210,48 @@ def square(x): cleanup_remote_function_assets( session.bqclient, session.cloudfunctionsclient, square ) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_anonymous_dataset(session, scalars_dfs): + try: + # This usage of remote_function is expected to create the remote + # function in the bigframes session's anonymous dataset. Use reuse=False + # param to make sure parallel instances of the test don't step over each + # other due to the common anonymous dataset. + @session.remote_function([int], int, reuse=False) + def square(x): + return x * x + + assert ( + bigquery.Routine(square.bigframes_remote_function).dataset_id + == session._anonymous_dataset.dataset_id + ) + + 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(bf_result, pd_result) + finally: + # clean up the gcp assets created for the remote function + cleanup_remote_function_assets( + session.bqclient, session.cloudfunctionsclient, square + ) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 3d8532a13b6..960a3841261 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -62,13 +62,12 @@ def bq_cf_connection_location_project_mismatched() -> str: @pytest.fixture(scope="module") -def session_with_bq_connection_and_permanent_dataset( +def session_with_bq_connection( bq_cf_connection, dataset_id_permanent ) -> bigframes.Session: session = bigframes.Session( bigframes.BigQueryOptions(bq_connection=bq_cf_connection) ) - session._session_dataset = bigquery.Dataset(dataset_id_permanent) return session @@ -277,13 +276,11 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_direct_session_param( - session_with_bq_connection_and_permanent_dataset, scalars_dfs -): +def test_remote_function_direct_session_param(session_with_bq_connection, scalars_dfs): @rf.remote_function( [int], int, - session=session_with_bq_connection_and_permanent_dataset, + session=session_with_bq_connection, ) def square(x): return x * x @@ -313,9 +310,7 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_remote_function_via_session_default( - session_with_bq_connection_and_permanent_dataset, scalars_dfs -): +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 @@ -323,7 +318,7 @@ def test_remote_function_via_session_default( # 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_and_permanent_dataset.remote_function([int], int) + @session_with_bq_connection.remote_function([int], int) def square(x): return x * x @@ -391,15 +386,11 @@ def square(x): @pytest.mark.flaky(retries=2, delay=120) -def test_dataframe_applymap( - session_with_bq_connection_and_permanent_dataset, scalars_dfs -): +def test_dataframe_applymap(session_with_bq_connection, scalars_dfs): def add_one(x): return x + 1 - remote_add_one = session_with_bq_connection_and_permanent_dataset.remote_function( - [int], int - )(add_one) + 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"] @@ -422,15 +413,11 @@ def add_one(x): @pytest.mark.flaky(retries=2, delay=120) -def test_dataframe_applymap_na_ignore( - session_with_bq_connection_and_permanent_dataset, scalars_dfs -): +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_and_permanent_dataset.remote_function( - [int], int - )(add_one) + 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"] @@ -451,13 +438,11 @@ def add_one(x): @pytest.mark.flaky(retries=2, delay=120) -def test_series_map(session_with_bq_connection_and_permanent_dataset, scalars_dfs): +def test_series_map(session_with_bq_connection, scalars_dfs): def add_one(x): return x + 1 - remote_add_one = session_with_bq_connection_and_permanent_dataset.remote_function( - [int], int - )(add_one) + remote_add_one = session_with_bq_connection.remote_function([int], int)(add_one) scalars_df, scalars_pandas_df = scalars_dfs From f8917abc094e222e0435891d4d184b77bfe67722 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Tue, 28 Nov 2023 18:24:16 -0800 Subject: [PATCH 185/197] fix: update the llm+kmeans notebook with recent change (#236) 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 internal issue 313682530 🦕 --- .../bq_dataframes_llm_kmeans.ipynb | 47 +++++-------------- 1 file changed, 12 insertions(+), 35 deletions(-) diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb index 8d759509254..5f74046fc0e 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -366,18 +366,6 @@ "predicted_embeddings.head() " ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "4H_etYfsEOFP" - }, - "outputs": [], - "source": [ - "# Join the complaints with their embeddings in the same DataFrame\n", - "combined_df = downsampled_issues_df.join(predicted_embeddings)" - ] - }, { "attachments": {}, "cell_type": "markdown", @@ -426,30 +414,19 @@ "outputs": [], "source": [ "# Use KMeans clustering to calculate our groups. Will take ~3 minutes.\n", - "cluster_model.fit(combined_df[[\"text_embedding\"]])\n", - "clustered_result = cluster_model.predict(combined_df[[\"text_embedding\"]])\n", + "cluster_model.fit(predicted_embeddings[[\"text_embedding\"]])\n", + "clustered_result = cluster_model.predict(predicted_embeddings)\n", "# Notice the CENTROID_ID column, which is the ID number of the group that\n", "# each complaint belongs to.\n", "clustered_result.head(n=5)" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Join the group number to the complaints and their text embeddings\n", - "combined_clustered_result = combined_df.join(clustered_result)\n", - "combined_clustered_result.head(n=5) " - ] - }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ - "Our dataframe combined_clustered_result now has three columns: the complaints, their text embeddings, and an ID from 1-10 (inclusive) indicating which semantically similar group they belong to." + "Our dataframe combined_clustered_result now has three complaint columns: the content, their text embeddings, and an ID from 1-10 (inclusive) indicating which semantically similar group they belong to." ] }, { @@ -480,14 +457,14 @@ "source": [ "# Using bigframes, with syntax identical to pandas,\n", "# filter out the first and second groups\n", - "cluster_1_result = combined_clustered_result[\n", - " combined_clustered_result[\"CENTROID_ID\"] == 1\n", - "][[\"consumer_complaint_narrative\"]]\n", + "cluster_1_result = clustered_result[\n", + " clustered_result[\"CENTROID_ID\"] == 1\n", + "][[\"content\"]]\n", "cluster_1_result_pandas = cluster_1_result.head(5).to_pandas()\n", "\n", - "cluster_2_result = combined_clustered_result[\n", - " combined_clustered_result[\"CENTROID_ID\"] == 2\n", - "][[\"consumer_complaint_narrative\"]]\n", + "cluster_2_result = clustered_result[\n", + " clustered_result[\"CENTROID_ID\"] == 2\n", + "][[\"content\"]]\n", "cluster_2_result_pandas = cluster_2_result.head(5).to_pandas()" ] }, @@ -503,15 +480,15 @@ "prompt1 = 'comment list 1:\\n'\n", "for i in range(5):\n", " prompt1 += str(i + 1) + '. ' + \\\n", - " cluster_1_result_pandas[\"consumer_complaint_narrative\"].iloc[i] + '\\n'\n", + " cluster_1_result_pandas[\"content\"].iloc[i] + '\\n'\n", "\n", "prompt2 = 'comment list 2:\\n'\n", "for i in range(5):\n", " prompt2 += str(i + 1) + '. ' + \\\n", - " cluster_2_result_pandas[\"consumer_complaint_narrative\"].iloc[i] + '\\n'\n", + " cluster_2_result_pandas[\"content\"].iloc[i] + '\\n'\n", "\n", "print(prompt1)\n", - "print(prompt2)\n" + "print(prompt2)" ] }, { From fbc31ab0fb166e6ffab48b3507614fef566ec64c Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 29 Nov 2023 03:52:20 +0000 Subject: [PATCH 186/197] test: fix prerelease tests (#239) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The heads of google-cloud-bigquery and google-cloud-bigquery-storage packages are installed with -e. This is leading to `from google.cloud import bigquery` looking into the folder of the latter and running into "ImportError: cannot import name 'bigquery' from 'google.cloud' (.../google-cloud-bigquery-storage/google/cloud/__init__.py). Removing -e from google-cloud-bigquery installation gets rid of this error. 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 internal issue 313701211 🦕 --- noxfile.py | 1 - 1 file changed, 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index c1fb53f7947..3b10a37fc7d 100644 --- a/noxfile.py +++ b/noxfile.py @@ -548,7 +548,6 @@ def prerelease(session: nox.sessions.Session, tests_path): # Ensure we catch breaking changes in the client libraries early. session.install( "--upgrade", - "-e", "git+https://github.com/googleapis/python-bigquery.git#egg=google-cloud-bigquery", ) already_installed.add("google-cloud-bigquery") From d0d9b84b101eb03c499d85e74dcfc900dedd4137 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Wed, 29 Nov 2023 09:15:41 -0800 Subject: [PATCH 187/197] fix: add df snapshots lookup for `read_gbq` (#229) --- bigframes/pandas/__init__.py | 6 ++ bigframes/session/__init__.py | 56 +++++++++++-------- bigframes/session/_io/bigquery.py | 5 -- tests/system/small/test_session.py | 18 ++++++ tests/unit/session/test_io_bigquery.py | 14 ----- .../bigframes_vendored/pandas/io/gbq.py | 3 + 6 files changed, 59 insertions(+), 43 deletions(-) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index d35f8383665..0c2c1f87aa6 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -486,6 +486,7 @@ def read_gbq( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query_or_table) return global_session.with_default_session( @@ -494,6 +495,7 @@ def read_gbq( index_col=index_col, col_order=col_order, max_results=max_results, + use_cache=use_cache, ) @@ -516,6 +518,7 @@ def read_gbq_query( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query) return global_session.with_default_session( @@ -524,6 +527,7 @@ def read_gbq_query( index_col=index_col, col_order=col_order, max_results=max_results, + use_cache=use_cache, ) @@ -536,6 +540,7 @@ def read_gbq_table( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ) -> bigframes.dataframe.DataFrame: _set_default_session_location_if_possible(query) return global_session.with_default_session( @@ -544,6 +549,7 @@ def read_gbq_table( index_col=index_col, col_order=col_order, max_results=max_results, + use_cache=use_cache, ) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index d2f6137883c..84a6eb56381 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -177,6 +177,7 @@ def __init__( # Now that we're starting the session, don't allow the options to be # changed. context._session_started = True + self._df_snapshot: Dict[bigquery.TableReference, datetime.datetime] = {} @property def bqclient(self): @@ -232,6 +233,7 @@ def read_gbq( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, # 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. @@ -242,6 +244,7 @@ def read_gbq( col_order=col_order, max_results=max_results, api_name="read_gbq", + use_cache=use_cache, ) else: # TODO(swast): Query the snapshot table but mark it as a @@ -253,6 +256,7 @@ def read_gbq( col_order=col_order, max_results=max_results, api_name="read_gbq", + use_cache=use_cache, ) def _query_to_destination( @@ -260,6 +264,7 @@ def _query_to_destination( query: str, index_cols: List[str], api_name: str, + use_cache: bool = True, ) -> Tuple[Optional[bigquery.TableReference], Optional[bigquery.QueryJob]]: # 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. @@ -284,6 +289,7 @@ def _query_to_destination( job_config = bigquery.QueryJobConfig() job_config.labels["bigframes-api"] = api_name job_config.destination = temp_table + job_config.use_query_cache = use_cache try: # Write to temp table to workaround BigQuery 10 GB query results @@ -305,6 +311,7 @@ def read_gbq_query( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ) -> dataframe.DataFrame: """Turn a SQL query into a DataFrame. @@ -362,6 +369,7 @@ def read_gbq_query( col_order=col_order, max_results=max_results, api_name="read_gbq_query", + use_cache=use_cache, ) def _read_gbq_query( @@ -372,6 +380,7 @@ def _read_gbq_query( col_order: Iterable[str] = (), max_results: Optional[int] = None, api_name: str = "read_gbq_query", + use_cache: bool = True, ) -> dataframe.DataFrame: if isinstance(index_col, str): index_cols = [index_col] @@ -379,7 +388,10 @@ def _read_gbq_query( index_cols = list(index_col) destination, query_job = self._query_to_destination( - query, index_cols, api_name=api_name + query, + index_cols, + api_name=api_name, + use_cache=use_cache, ) # If there was no destination table, that means the query must have @@ -403,6 +415,7 @@ def _read_gbq_query( index_col=index_cols, col_order=col_order, max_results=max_results, + use_cache=use_cache, ) def read_gbq_table( @@ -412,6 +425,7 @@ def read_gbq_table( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ) -> dataframe.DataFrame: """Turn a BigQuery table into a DataFrame. @@ -434,6 +448,7 @@ def read_gbq_table( col_order=col_order, max_results=max_results, api_name="read_gbq_table", + use_cache=use_cache, ) def _get_snapshot_sql_and_primary_key( @@ -441,6 +456,7 @@ def _get_snapshot_sql_and_primary_key( table_ref: bigquery.table.TableReference, *, api_name: str, + use_cache: bool = True, ) -> Tuple[ibis_types.Table, Optional[Sequence[str]]]: """Create a read-only Ibis table expression representing a table. @@ -448,19 +464,6 @@ def _get_snapshot_sql_and_primary_key( column(s), then return those too so that ordering generation can be avoided. """ - if table_ref.dataset_id.upper() == "_SESSION": - # _SESSION tables aren't supported by the tables.get REST API. - return ( - self.ibis_client.sql( - f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`" - ), - None, - ) - table_expression = self.ibis_client.table( - table_ref.table_id, - database=f"{table_ref.project}.{table_ref.dataset_id}", - ) - # If there are primary keys defined, the query engine assumes these # columns are unique, even if the constraint is not enforced. We make # the same assumption and use these columns as the total ordering keys. @@ -481,14 +484,18 @@ def _get_snapshot_sql_and_primary_key( job_config = bigquery.QueryJobConfig() job_config.labels["bigframes-api"] = api_name - current_timestamp = list( - self.bqclient.query( - "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", - job_config=job_config, - ).result() - )[0][0] + if use_cache and table_ref in self._df_snapshot.keys(): + snapshot_timestamp = self._df_snapshot[table_ref] + else: + snapshot_timestamp = list( + self.bqclient.query( + "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", + job_config=job_config, + ).result() + )[0][0] + self._df_snapshot[table_ref] = snapshot_timestamp table_expression = self.ibis_client.sql( - bigframes_io.create_snapshot_sql(table_ref, current_timestamp) + bigframes_io.create_snapshot_sql(table_ref, snapshot_timestamp) ) return table_expression, primary_keys @@ -500,12 +507,11 @@ def _read_gbq_table( col_order: Iterable[str] = (), max_results: Optional[int] = None, api_name: str, + use_cache: bool = True, ) -> dataframe.DataFrame: if max_results and max_results <= 0: raise ValueError("`max_results` should be a positive number.") - # 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 ) @@ -513,7 +519,9 @@ def _read_gbq_table( ( table_expression, total_ordering_cols, - ) = self._get_snapshot_sql_and_primary_key(table_ref, api_name=api_name) + ) = self._get_snapshot_sql_and_primary_key( + table_ref, api_name=api_name, use_cache=use_cache + ) for key in col_order: if key not in table_expression.columns: diff --git a/bigframes/session/_io/bigquery.py b/bigframes/session/_io/bigquery.py index dae73301e7b..4770f120894 100644 --- a/bigframes/session/_io/bigquery.py +++ b/bigframes/session/_io/bigquery.py @@ -117,11 +117,6 @@ 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}`" - # 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("_"): diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 7cd9f1dd598..26c5093b352 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -16,6 +16,7 @@ import random import tempfile import textwrap +import time import typing from typing import List @@ -308,6 +309,23 @@ def test_read_gbq_w_script_no_select(session, dataset_id: str): assert df["statement_type"][0] == "SCRIPT" +def test_read_gbq_twice_with_same_timestamp(session, penguins_table_id): + df1 = session.read_gbq(penguins_table_id) + time.sleep(1) + df2 = session.read_gbq(penguins_table_id) + df1.columns = [ + "species1", + "island1", + "culmen_length_mm1", + "culmen_depth_mm1", + "flipper_length_mm1", + "body_mass_g1", + "sex1", + ] + df3 = df1.join(df2) + assert df3 is not None + + 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) diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index c87835c4126..3f3bfbe7d36 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -147,20 +147,6 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): assert "`my-test-project`.`_e8166e0cdb`.`anonbb92cd`" in sql -def test_create_snapshot_sql_doesnt_timetravel_session_tables(): - table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") - - sql = bigframes.session._io.bigquery.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 - - def test_create_temp_table_default_expiration(): """Make sure the created table has an expiration.""" bqclient = mock.create_autospec(bigquery.Client) diff --git a/third_party/bigframes_vendored/pandas/io/gbq.py b/third_party/bigframes_vendored/pandas/io/gbq.py index 2161310b07c..eabb48e600a 100644 --- a/third_party/bigframes_vendored/pandas/io/gbq.py +++ b/third_party/bigframes_vendored/pandas/io/gbq.py @@ -16,6 +16,7 @@ def read_gbq( index_col: Iterable[str] | str = (), col_order: Iterable[str] = (), max_results: Optional[int] = None, + use_cache: bool = True, ): """Loads a DataFrame from BigQuery. @@ -83,6 +84,8 @@ def read_gbq( max_results (Optional[int], default None): If set, limit the maximum number of rows to fetch from the query results. + use_cache (bool, default True): + Whether to cache the query inputs. Default to True. Returns: bigframes.dataframe.DataFrame: A DataFrame representing results of the query or table. From 0bfc4fb117686c734d4a2503d5a6de0e64e9f9b9 Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Wed, 29 Nov 2023 11:44:16 -0800 Subject: [PATCH 188/197] feat: add remote vertex model support (#237) b/299356085 --- bigframes/ml/__init__.py | 1 + bigframes/ml/core.py | 8 ++ bigframes/ml/remote.py | 157 +++++++++++++++++++++++++ bigframes/ml/sql.py | 26 +++- docs/reference/bigframes.ml/index.rst | 2 + docs/reference/bigframes.ml/remote.rst | 7 ++ docs/templates/toc.yml | 6 + tests/system/small/ml/conftest.py | 41 +++++++ tests/system/small/ml/test_core.py | 16 +++ tests/system/small/ml/test_remote.py | 33 ++++++ tests/unit/ml/test_sql.py | 26 ++++ 11 files changed, 319 insertions(+), 4 deletions(-) create mode 100644 bigframes/ml/remote.py create mode 100644 docs/reference/bigframes.ml/remote.rst create mode 100644 tests/system/small/ml/test_remote.py diff --git a/bigframes/ml/__init__.py b/bigframes/ml/__init__.py index 55c8709d8d8..b2c62ff9612 100644 --- a/bigframes/ml/__init__.py +++ b/bigframes/ml/__init__.py @@ -26,4 +26,5 @@ "llm", "forecasting", "imported", + "remote", ] diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py index d8135f7085f..5aad77a3942 100644 --- a/bigframes/ml/core.py +++ b/bigframes/ml/core.py @@ -294,6 +294,8 @@ def create_remote_model( self, session: bigframes.Session, connection_name: str, + input: Mapping[str, str] = {}, + output: Mapping[str, str] = {}, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> BqmlModel: """Create a session-temporary BQML remote model with the CREATE OR REPLACE MODEL statement @@ -301,6 +303,10 @@ def create_remote_model( Args: connection_name: a BQ connection to talk with Vertex AI, of the format ... https://cloud.google.com/bigquery/docs/create-cloud-resource-connection + input: + input schema for general remote models + output: + output schema for general remote models options: a dict of options to configure the model. Generates a BQML OPTIONS clause @@ -311,6 +317,8 @@ def create_remote_model( sql = self._model_creation_sql_generator.create_remote_model( connection_name=connection_name, model_ref=model_ref, + input=input, + output=output, options=options, ) diff --git a/bigframes/ml/remote.py b/bigframes/ml/remote.py new file mode 100644 index 00000000000..d4c34bbd0db --- /dev/null +++ b/bigframes/ml/remote.py @@ -0,0 +1,157 @@ +# 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. + +"""BigFrames general remote models.""" + +from __future__ import annotations + +from typing import Mapping, Optional, Union +import warnings + +import bigframes +from bigframes import clients +from bigframes.core import log_adapter +from bigframes.ml import base, core, globals, utils +import bigframes.pandas as bpd + +_SUPPORTED_DTYPES = ( + "bool", + "string", + "int64", + "float64", + "array", + "array", + "array", + "array", +) + +_REMOTE_MODEL_STATUS = "remote_model_status" + + +@log_adapter.class_logger +class VertexAIModel(base.BaseEstimator): + """Remote model from a Vertex AI https endpoint. User must specify https endpoint, input schema and output schema. + How to deploy a model in Vertex AI https://cloud.google.com/bigquery/docs/bigquery-ml-remote-model-tutorial#Deploy-Model-on-Vertex-AI. + + Args: + endpoint (str): + Vertex AI https endpoint. + input ({column_name: column_type}): + Input schema. Supported types are "bool", "string", "int64", "float64", "array", "array", "array", "array". + output ({column_name: column_type}): + Output label schema. Supported the same types as the input. + 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. BigQuery DataFrame will try to create the connection and attach + permission if the connection isn't fully setup. + """ + + def __init__( + self, + endpoint: str, + input: Mapping[str, str], + output: Mapping[str, str], + session: Optional[bigframes.Session] = None, + connection_name: Optional[str] = None, + ): + self.endpoint = endpoint + self.input = input + self.output = output + 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 = self._bq_connection_manager.resolve_full_connection_name( + connection_name, + default_project=self.session._project, + default_location=self.session._location, + ) + + self._bqml_model_factory = globals.bqml_model_factory() + 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 = { + "endpoint": self.endpoint, + } + + def standardize_type(v: str): + v = v.lower() + v = v.replace("boolean", "bool") + + if v not in _SUPPORTED_DTYPES: + raise ValueError( + f"Data type {v} is not supported. We only support {', '.join(_SUPPORTED_DTYPES)}." + ) + + return v + + self.input = {k: standardize_type(v) for k, v in self.input.items()} + self.output = {k: standardize_type(v) for k, v in self.output.items()} + + return self._bqml_model_factory.create_remote_model( + session=self.session, + connection_name=self.connection_name, + input=self.input, + output=self.output, + options=options, + ) + + def predict( + self, + X: Union[bpd.DataFrame, bpd.Series], + ) -> bpd.DataFrame: + """Predict the result from the input DataFrame. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + Input DataFrame or Series, which needs to comply with the input parameter of the model. + + Returns: + bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. + """ + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.predict(X) + + # unlike LLM models, the general remote model status is null for successful runs. + if (df[_REMOTE_MODEL_STATUS].notna()).any(): + warnings.warn( + f"Some predictions failed. Check column {_REMOTE_MODEL_STATUS} for detailed status. You may want to filter the failed rows and retry.", + RuntimeWarning, + ) + + return df diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index ab051231fb0..1c88eda4ab3 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -57,6 +57,12 @@ def build_expressions(self, *expr_sqls: str) -> str: indent_str = " " return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls) + def build_schema(self, **kwargs: str) -> str: + """Encode a dict of values into a formatted schema type items for SQL""" + indent_str = " " + param_strs = [f"{k} {v}" for k, v in kwargs.items()] + return "\n" + indent_str + f",\n{indent_str}".join(param_strs) + def options(self, **kwargs: Union[str, int, float, Iterable[str]]) -> str: """Encode the OPTIONS clause for BQML""" return f"OPTIONS({self.build_parameters(**kwargs)})" @@ -65,6 +71,14 @@ def struct_options(self, **kwargs: Union[int, float]) -> str: """Encode a BQ STRUCT as options.""" return f"STRUCT({self.build_structs(**kwargs)})" + def input(self, **kwargs: str) -> str: + """Encode a BQML INPUT clause.""" + return f"INPUT({self.build_schema(**kwargs)})" + + def output(self, **kwargs: str) -> str: + """Encode a BQML OUTPUT clause.""" + return f"OUTPUT({self.build_schema(**kwargs)})" + # Connection def connection(self, conn_name: str) -> str: """Encode the REMOTE WITH CONNECTION clause for BQML. conn_name is of the format ...""" @@ -154,15 +168,19 @@ def create_remote_model( self, connection_name: str, model_ref: google.cloud.bigquery.ModelReference, + input: Mapping[str, str] = {}, + output: Mapping[str, str] = {}, options: Mapping[str, Union[str, int, float, Iterable[str]]] = {}, ) -> str: """Encode the CREATE OR REPLACE MODEL statement for BQML remote model.""" - options_sql = self.options(**options) - parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"] + if input: + parts.append(self.input(**input)) + if output: + parts.append(self.output(**output)) parts.append(self.connection(connection_name)) - if options_sql: - parts.append(options_sql) + if options: + parts.append(self.options(**options)) return "\n".join(parts) def create_imported_model( diff --git a/docs/reference/bigframes.ml/index.rst b/docs/reference/bigframes.ml/index.rst index f3cbe1174a7..1975d62e6d9 100644 --- a/docs/reference/bigframes.ml/index.rst +++ b/docs/reference/bigframes.ml/index.rst @@ -30,3 +30,5 @@ API Reference pipeline preprocessing + + remote diff --git a/docs/reference/bigframes.ml/remote.rst b/docs/reference/bigframes.ml/remote.rst new file mode 100644 index 00000000000..7827acfe923 --- /dev/null +++ b/docs/reference/bigframes.ml/remote.rst @@ -0,0 +1,7 @@ +bigframes.ml.remote +=================== + +.. automodule:: bigframes.ml.remote + :members: + :inherited-members: + :undoc-members: diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 9879721d286..58ac1c0efe6 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -108,6 +108,12 @@ - name: PaLM2TextEmbeddingGenerator uid: bigframes.ml.llm.PaLM2TextEmbeddingGenerator name: llm + - items: + - name: Overview + uid: bigframes.ml.remote + - name: VertexAIModel + uid: bigframes.ml.remote.VertexAIModel + name: remote - items: - name: metrics uid: bigframes.ml.metrics diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index c11445b79a5..c4a1272e441 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -29,6 +29,7 @@ imported, linear_model, llm, + remote, ) @@ -247,6 +248,46 @@ def palm2_embedding_generator_multilingual_model( ) +@pytest.fixture(scope="session") +def linear_remote_model_params() -> dict: + # Pre-deployed endpoint of linear reg model in Vertex. + # bigframes-test-linreg2 -> bigframes-test-linreg-endpoint2 + return { + "input": {"culmen_length_mm": "float64"}, + "output": {"predicted_body_mass_g": "array"}, + "endpoint": "https://us-central1-aiplatform.googleapis.com/v1/projects/1084210331973/locations/us-central1/endpoints/3193318217619603456", + } + + +@pytest.fixture(scope="session") +def bqml_linear_remote_model( + session, bq_connection, linear_remote_model_params +) -> core.BqmlModel: + options = { + "endpoint": linear_remote_model_params["endpoint"], + } + return globals.bqml_model_factory().create_remote_model( + session=session, + input=linear_remote_model_params["input"], + output=linear_remote_model_params["output"], + connection_name=bq_connection, + options=options, + ) + + +@pytest.fixture(scope="session") +def linear_remote_vertex_model( + session, bq_connection, linear_remote_model_params +) -> remote.VertexAIModel: + return remote.VertexAIModel( + endpoint=linear_remote_model_params["endpoint"], + input=linear_remote_model_params["input"], + output=linear_remote_model_params["output"], + session=session, + connection_name=bq_connection, + ) + + @pytest.fixture(scope="session") def time_series_bqml_arima_plus_model( session, time_series_arima_plus_model_name diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index be34a4871c6..22cbbb1932f 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -289,6 +289,22 @@ def test_model_predict_with_unnamed_index( ) +def test_remote_model_predict( + bqml_linear_remote_model: core.BqmlModel, new_penguins_df +): + predictions = bqml_linear_remote_model.predict(new_penguins_df).to_pandas() + expected = pd.DataFrame( + {"predicted_body_mass_g": [[3739.54], [3675.79], [3619.54]]}, + 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, + ) + + @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_remote.py b/tests/system/small/ml/test_remote.py new file mode 100644 index 00000000000..e8eb1c85e86 --- /dev/null +++ b/tests/system/small/ml/test_remote.py @@ -0,0 +1,33 @@ +# 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 remote + + +def test_remote_linear_vertex_model_predict( + linear_remote_vertex_model: remote.VertexAIModel, new_penguins_df +): + predictions = linear_remote_vertex_model.predict(new_penguins_df).to_pandas() + expected = pd.DataFrame( + {"predicted_body_mass_g": [[3739.54], [3675.79], [3619.54]]}, + 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, + ) diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index ea16722393d..92230585405 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -190,6 +190,32 @@ def test_create_remote_model_produces_correct_sql( ) +def test_create_remote_model_with_params_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", + model_ref=bigquery.ModelReference.from_string( + "test-proj._anonXYZ.create_remote_model" + ), + input={"column1": "int64"}, + output={"result": "array"}, + options={"option_key1": "option_value1", "option_key2": 2}, + ) + assert ( + sql + == """CREATE OR REPLACE MODEL `test-proj`.`_anonXYZ`.`create_remote_model` +INPUT( + column1 int64) +OUTPUT( + result array) +REMOTE WITH CONNECTION `my_project.us.my_connection` +OPTIONS( + option_key1="option_value1", + option_key2=2)""" + ) + + def test_create_imported_model_produces_correct_sql( model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator, ): From 6c899be2989e24f697d72fe1bb92ebbf7dec84cb Mon Sep 17 00:00:00 2001 From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com> Date: Wed, 29 Nov 2023 13:12:15 -0800 Subject: [PATCH 189/197] chore: release 0.15.0 (#241) Release-As: 0.15.0 From 8089b15feddaeb9c56a8f976b439315fcfed0301 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Wed, 29 Nov 2023 14:46:09 -0800 Subject: [PATCH 190/197] chore(main): release 0.15.0 (#214) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 45 ++++++++++++++++++++++++++++++++++++++++++++ bigframes/version.py | 2 +- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 091967513aa..ef75a017e01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,51 @@ [1]: https://pypi.org/project/bigframes/#history +## [0.15.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.14.1...v0.15.0) (2023-11-29) + + +### ⚠ BREAKING CHANGES + +* model.predict returns all the columns ([#204](https://github.com/googleapis/python-bigquery-dataframes/issues/204)) + +### Features + +* Add info and memory_usage methods to dataframe ([#219](https://github.com/googleapis/python-bigquery-dataframes/issues/219)) ([9d6613d](https://github.com/googleapis/python-bigquery-dataframes/commit/9d6613d318b558722b7bab12773efdea4bbe9931)) +* Add remote vertex model support ([#237](https://github.com/googleapis/python-bigquery-dataframes/issues/237)) ([0bfc4fb](https://github.com/googleapis/python-bigquery-dataframes/commit/0bfc4fb117686c734d4a2503d5a6de0e64e9f9b9)) +* Add the recent api method for ML component ([#225](https://github.com/googleapis/python-bigquery-dataframes/issues/225)) ([ed8876d](https://github.com/googleapis/python-bigquery-dataframes/commit/ed8876d3439a3b45b65e8789737c3c2e3a7f1adb)) +* Model.predict returns all the columns ([#204](https://github.com/googleapis/python-bigquery-dataframes/issues/204)) ([416171a](https://github.com/googleapis/python-bigquery-dataframes/commit/416171a70d91d4a6b71622ba72685147ab7d6186)) +* Send warnings on LLM prediction partial failures ([#216](https://github.com/googleapis/python-bigquery-dataframes/issues/216)) ([81125f9](https://github.com/googleapis/python-bigquery-dataframes/commit/81125f9505ad98e89939769a8e1fcf30518705f0)) + + +### Bug Fixes + +* Add df snapshots lookup for `read_gbq` ([#229](https://github.com/googleapis/python-bigquery-dataframes/issues/229)) ([d0d9b84](https://github.com/googleapis/python-bigquery-dataframes/commit/d0d9b84b101eb03c499d85e74dcfc900dedd4137)) +* Avoid unnecessary row_number() on sort key for io ([#211](https://github.com/googleapis/python-bigquery-dataframes/issues/211)) ([a18d40e](https://github.com/googleapis/python-bigquery-dataframes/commit/a18d40e808ee0822d21715cc3e8f794c418aeebc)) +* Dedup special character ([#209](https://github.com/googleapis/python-bigquery-dataframes/issues/209)) ([dd78acb](https://github.com/googleapis/python-bigquery-dataframes/commit/dd78acb174545ba292776a642afcec46f8ee4a2a)) +* Invalid JSON type of the notebook ([#215](https://github.com/googleapis/python-bigquery-dataframes/issues/215)) ([a729831](https://github.com/googleapis/python-bigquery-dataframes/commit/a7298317ea2604faa6ae31817f1f729d7e0b9818)) +* Make to_pandas override enable_downsampling when sampling_method is manually set. ([#200](https://github.com/googleapis/python-bigquery-dataframes/issues/200)) ([ae03756](https://github.com/googleapis/python-bigquery-dataframes/commit/ae03756f5ee45e0e74e0c0bdd4777e018eba2273)) +* Polish the llm+kmeans notebook ([#208](https://github.com/googleapis/python-bigquery-dataframes/issues/208)) ([e8532b1](https://github.com/googleapis/python-bigquery-dataframes/commit/e8532b1d999d26ea1ebdd30efb8f2c0a93a6a28d)) +* Update the llm+kmeans notebook with recent change ([#236](https://github.com/googleapis/python-bigquery-dataframes/issues/236)) ([f8917ab](https://github.com/googleapis/python-bigquery-dataframes/commit/f8917abc094e222e0435891d4d184b77bfe67722)) +* Use anonymous dataset to create `remote_function` ([#205](https://github.com/googleapis/python-bigquery-dataframes/issues/205)) ([69b016e](https://github.com/googleapis/python-bigquery-dataframes/commit/69b016eae7ea97d84ceeb22ba09f5472841db072)) + + +### Documentation + +* Add code samples for `index` and `column` properties ([#212](https://github.com/googleapis/python-bigquery-dataframes/issues/212)) ([c88d38e](https://github.com/googleapis/python-bigquery-dataframes/commit/c88d38e69682f4c620174086b8f16f4780c04811)) +* Add code samples for df reshaping, function, merge, and join methods ([#203](https://github.com/googleapis/python-bigquery-dataframes/issues/203)) ([010486c](https://github.com/googleapis/python-bigquery-dataframes/commit/010486c3494e05d714da6cc7d51514518d9ae1ea)) +* Add examples for dataframe.kurt, dataframe.std, dataframe.count ([#232](https://github.com/googleapis/python-bigquery-dataframes/issues/232)) ([f9c6e72](https://github.com/googleapis/python-bigquery-dataframes/commit/f9c6e727e2b901310bb5301da449d616ea85e135)) +* Add examples for dataframe.mean, dataframe.median, dataframe.va… ([#228](https://github.com/googleapis/python-bigquery-dataframes/issues/228)) ([edd0522](https://github.com/googleapis/python-bigquery-dataframes/commit/edd0522747eadb74780124fb18ed7face251441d)) +* Add examples for dataframe.min, dataframe.max and dataframe.sum ([#227](https://github.com/googleapis/python-bigquery-dataframes/issues/227)) ([3a375e8](https://github.com/googleapis/python-bigquery-dataframes/commit/3a375e87b64b8fb51370bfec8f2cfdbcd8fe960a)) +* Code samples for `Series.dot` and `DataFrame.dot` ([#226](https://github.com/googleapis/python-bigquery-dataframes/issues/226)) ([b62a07a](https://github.com/googleapis/python-bigquery-dataframes/commit/b62a07a95cd60f995a48825c9874822d0eb02483)) +* Code samples for `Series.where` and `Series.mask` ([#217](https://github.com/googleapis/python-bigquery-dataframes/issues/217)) ([52dfad2](https://github.com/googleapis/python-bigquery-dataframes/commit/52dfad281def82548751a276ce42b087dbb09f9a)) +* Code samples for dataframe.any, dataframe.all and dataframe.prod ([#223](https://github.com/googleapis/python-bigquery-dataframes/issues/223)) ([d7957fa](https://github.com/googleapis/python-bigquery-dataframes/commit/d7957fad071d223ef8f6fb8f3de395c865ff60aa)) +* Make the code samples reflect default bq connection usage ([#206](https://github.com/googleapis/python-bigquery-dataframes/issues/206)) ([71844b0](https://github.com/googleapis/python-bigquery-dataframes/commit/71844b03cdbfe684320c186a0488c8c7fb4fcd6e)) + + +### Miscellaneous Chores + +* Release 0.15.0 ([#241](https://github.com/googleapis/python-bigquery-dataframes/issues/241)) ([6c899be](https://github.com/googleapis/python-bigquery-dataframes/commit/6c899be2989e24f697d72fe1bb92ebbf7dec84cb)) + ## [0.14.1](https://github.com/googleapis/python-bigquery-dataframes/compare/v0.14.0...v0.14.1) (2023-11-16) diff --git a/bigframes/version.py b/bigframes/version.py index 46e57e5b888..920cb95c3d2 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.14.1" +__version__ = "0.15.0" From 66d1839c3e9a3011c7feb13a59d966b64cf8313f Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Thu, 30 Nov 2023 11:39:47 -0800 Subject: [PATCH 191/197] fix: update the llm_kmeans notebook (#247) --- notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb index 5f74046fc0e..69efb11018a 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -371,7 +371,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "We now have the complaints and their text embeddings as two columns in our combined_df. Recall that complaints with numerically similar text embeddings should have similar meanings semantically. We will now group similar complaints together." + "We now have the complaints and their text embeddings as two columns in our predicted_embeddings DataFrame." ] }, { @@ -426,7 +426,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Our dataframe combined_clustered_result now has three complaint columns: the content, their text embeddings, and an ID from 1-10 (inclusive) indicating which semantically similar group they belong to." + "Our DataFrame clustered_result now has an additional column that includes an ID from 1-10 (inclusive) indicating which semantically similar group they belong to." ] }, { @@ -501,7 +501,7 @@ "source": [ "# The plain English request we will make of PaLM 2\n", "prompt = (\n", - " \"Please highlight the most obvious difference between\"\n", + " \"Please highlight the most obvious difference between \"\n", " \"the two lists of comments:\\n\" + prompt1 + prompt2\n", ")\n", "print(prompt)" From 1737acc51b4fdd9b385bbf91a758efd2e7ead11a Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Thu, 30 Nov 2023 20:22:16 -0800 Subject: [PATCH 192/197] feat: add DataFrame.select_dtypes method (#242) 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/dataframe.py | 13 +++++++ tests/system/small/test_dataframe.py | 20 +++++++++++ .../bigframes_vendored/pandas/core/frame.py | 36 +++++++++++++++++++ 3 files changed, 69 insertions(+) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index f7796291b9a..c6b28f1b012 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -434,6 +434,19 @@ def info( # TODO: Convert to different units (kb, mb, etc.) obuf.write(f"memory usage: {self.memory_usage().sum()} bytes\n") + def select_dtypes(self, include=None, exclude=None) -> DataFrame: + # Create empty pandas dataframe with same schema and then leverage actual pandas implementation + as_pandas = pandas.DataFrame( + { + col_id: pandas.Series([], dtype=dtype) + for col_id, dtype in zip(self._block.value_columns, self._block.dtypes) + } + ) + selected_columns = tuple( + as_pandas.select_dtypes(include=include, exclude=exclude).columns + ) + return DataFrame(self._block.select_columns(selected_columns)) + def _set_internal_query_job(self, query_job: bigquery.QueryJob): self._query_job = query_job diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 9744d3f6e91..5940df590c3 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -297,6 +297,26 @@ def test_df_info(scalars_dfs): assert expected == bf_result.getvalue() +@pytest.mark.parametrize( + ("include", "exclude"), + [ + ("Int64", None), + (["int"], None), + ("number", None), + ([pd.Int64Dtype(), pd.BooleanDtype()], None), + (None, [pd.Int64Dtype(), pd.BooleanDtype()]), + ("Int64", ["boolean"]), + ], +) +def test_select_dtypes(scalars_dfs, include, exclude): + scalars_df, scalars_pandas_df = scalars_dfs + + pd_result = scalars_pandas_df.select_dtypes(include=include, exclude=exclude) + bf_result = scalars_df.select_dtypes(include=include, exclude=exclude).to_pandas() + + pd.testing.assert_frame_equal(pd_result, bf_result) + + def test_drop_index(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 099d8b8e663..3bd90be2e4f 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -158,6 +158,42 @@ def memory_usage(self, index: bool = True): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def select_dtypes(self, include=None, exclude=None) -> DataFrame: + """ + Return a subset of the DataFrame's columns based on the column dtypes. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': ["hello", "world"], 'col3': [True, False]}) + >>> df.select_dtypes(include=['Int64']) + col1 + 0 1 + 1 2 + + [2 rows x 1 columns] + + >>> df.select_dtypes(exclude=['Int64']) + col2 col3 + 0 hello True + 1 world False + + [2 rows x 2 columns] + + + Args: + include (scalar or list-like): + A selection of dtypes or strings to be included. + exclude (scalar or list-like): + A selection of dtypes or strings to be excluded. + + Returns: + DataFrame: The subset of the frame including the dtypes in ``include`` and excluding the dtypes in ``exclude``. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + # ---------------------------------------------------------------------- # IO methods (to / from other formats) def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarray: From 0523a31fa0b589f88afe0ad5b447634409ddeb86 Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Fri, 1 Dec 2023 10:06:23 -0800 Subject: [PATCH 193/197] docs: add examples for dataframe.cummin, dataframe.cummax, dataframe.cumsum, dataframe.cumprod (#243) --- .../bigframes_vendored/pandas/core/frame.py | 96 ++++++++++++++++++- 1 file changed, 92 insertions(+), 4 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 3bd90be2e4f..6b5a580e992 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -3367,40 +3367,128 @@ def nunique(self): raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummin(self) -> DataFrame: - """Return cumulative minimum over a DataFrame axis. + """Return cumulative minimum over columns. Returns a DataFrame of the same size containing the cumulative minimum. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + >>> df.cummin() + A B + 0 3 1 + 1 1 1 + 2 1 1 + + [3 rows x 2 columns] + Returns: bigframes.dataframe.DataFrame: Return cumulative minimum of DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cummax(self) -> DataFrame: - """Return cumulative maximum over a DataFrame axis. + """Return cumulative maximum over columns. Returns a DataFrame of the same size containing the cumulative maximum. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + >>> df.cummax() + A B + 0 3 1 + 1 3 2 + 2 3 3 + + [3 rows x 2 columns] + Returns: bigframes.dataframe.DataFrame: Return cumulative maximum of DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumsum(self) -> DataFrame: - """Return cumulative sum over a DataFrame axis. + """Return cumulative sum over columns. Returns a DataFrame of the same size containing the cumulative sum. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + >>> df.cumsum() + A B + 0 3 1 + 1 4 3 + 2 6 6 + + [3 rows x 2 columns] + Returns: bigframes.dataframe.DataFrame: Return cumulative sum of DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) def cumprod(self) -> DataFrame: - """Return cumulative product over a DataFrame axis. + """Return cumulative product over columns. Returns a DataFrame of the same size containing the cumulative product. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + >>> df.cumprod() + A B + 0 3 1 + 1 3 2 + 2 6 6 + + [3 rows x 2 columns] + Returns: bigframes.dataframe.DataFrame: Return cumulative product of DataFrame. """ From 8d81e24677613dcf4d275c27a327384b8c17bc85 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Fri, 1 Dec 2023 11:58:10 -0800 Subject: [PATCH 194/197] feat: add DataFrame from_dict and from_records methods (#244) 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/dataframe.py | 26 +++++++ tests/system/small/test_dataframe.py | 48 ++++++++++++ .../bigframes_vendored/pandas/core/frame.py | 75 +++++++++++++++++++ 3 files changed, 149 insertions(+) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c6b28f1b012..3b0fd7008ab 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2381,6 +2381,32 @@ def _split( blocks = self._block._split(ns=ns, fracs=fracs, random_state=random_state) return [DataFrame(block) for block in blocks] + @classmethod + def from_dict( + cls, + data: dict, + orient: str = "columns", + dtype=None, + columns=None, + ) -> DataFrame: + return cls(pandas.DataFrame.from_dict(data, orient, dtype, columns)) # type: ignore + + @classmethod + def from_records( + cls, + data, + index=None, + exclude=None, + columns=None, + coerce_float: bool = False, + nrows: int | None = None, + ) -> DataFrame: + return cls( + pandas.DataFrame.from_records( + data, index, exclude, columns, coerce_float, nrows + ) + ) + def to_csv( self, path_or_buf: str, sep=",", *, header: bool = True, index: bool = True ) -> None: diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 5940df590c3..9318a5d9d2f 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3309,6 +3309,54 @@ 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_from_dict_columns_orient(): + data = {"a": [1, 2], "b": [3.3, 2.4]} + bf_result = dataframe.DataFrame.from_dict(data, orient="columns").to_pandas() + pd_result = pd.DataFrame.from_dict(data, orient="columns") + assert_pandas_df_equal( + pd_result, bf_result, check_dtype=False, check_index_type=False + ) + + +def test_df_from_dict_index_orient(): + data = {"a": [1, 2], "b": [3.3, 2.4]} + bf_result = dataframe.DataFrame.from_dict( + data, orient="index", columns=["col1", "col2"] + ).to_pandas() + pd_result = pd.DataFrame.from_dict(data, orient="index", columns=["col1", "col2"]) + assert_pandas_df_equal( + pd_result, bf_result, check_dtype=False, check_index_type=False + ) + + +def test_df_from_dict_tight_orient(): + data = { + "index": [("i1", "i2"), ("i3", "i4")], + "columns": ["col1", "col2"], + "data": [[1, 2.6], [3, 4.5]], + "index_names": ["in1", "in2"], + "column_names": ["column_axis"], + } + + bf_result = dataframe.DataFrame.from_dict(data, orient="tight").to_pandas() + pd_result = pd.DataFrame.from_dict(data, orient="tight") + assert_pandas_df_equal( + pd_result, bf_result, check_dtype=False, check_index_type=False + ) + + +def test_df_from_records(): + records = ((1, "a"), (2.5, "b"), (3.3, "c"), (4.9, "d")) + + bf_result = dataframe.DataFrame.from_records( + records, columns=["c1", "c2"] + ).to_pandas() + pd_result = pd.DataFrame.from_records(records, columns=["c1", "c2"]) + assert_pandas_df_equal( + pd_result, bf_result, check_dtype=False, check_index_type=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() diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 6b5a580e992..08fe8e2de0b 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -196,6 +196,81 @@ def select_dtypes(self, include=None, exclude=None) -> DataFrame: # ---------------------------------------------------------------------- # IO methods (to / from other formats) + @classmethod + def from_dict( + cls, + data: dict, + orient="columns", + dtype=None, + columns=None, + ) -> DataFrame: + """ + Construct DataFrame from dict of array-like or dicts. + + Creates DataFrame object from dictionary by columns or by index + allowing dtype specification. + + Args: + data (dict): + Of the form {field : array-like} or {field : dict}. + orient ({'columns', 'index', 'tight'}, default 'columns'): + The "orientation" of the data. If the keys of the passed dict + should be the columns of the resulting DataFrame, pass 'columns' + (default). Otherwise if the keys should be rows, pass 'index'. + If 'tight', assume a dict with keys ['index', 'columns', 'data', + 'index_names', 'column_names']. + dtype (dtype, default None): + Data type to force after DataFrame construction, otherwise infer. + columns (list, default None): + Column labels to use when ``orient='index'``. Raises a ValueError + if used with ``orient='columns'`` or ``orient='tight'``. + + Returns: + DataFrame + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + @classmethod + def from_records( + cls, + data, + index=None, + exclude=None, + columns=None, + coerce_float: bool = False, + nrows: int | None = None, + ) -> DataFrame: + """ + Convert structured or record ndarray to DataFrame. + + Creates a DataFrame object from a structured ndarray, sequence of + tuples or dicts, or DataFrame. + + Args: + data (structured ndarray, sequence of tuples or dicts): + Structured input data. + index (str, list of fields, array-like): + Field of array to use as the index, alternately a specific set of + input labels to use. + exclude (sequence, default None): + Columns or fields to exclude. + columns (sequence, default None): + Column names to use. If the passed data do not have names + associated with them, this argument provides names for the + columns. Otherwise this argument indicates the order of the columns + in the result (any names not found in the data will become all-NA + columns). + coerce_float (bool, default False): + Attempt to convert values of non-string, non-numeric objects (like + decimal.Decimal) to floating point, useful for SQL result sets. + nrows (int, default None): + Number of rows to read if data is an iterator. + + Returns: + DataFrame + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarray: """ Convert the DataFrame to a NumPy array. From c2829e3d976a43c53251c9288266e3a8ec5304c5 Mon Sep 17 00:00:00 2001 From: Ashley Xu <139821907+ashleyxuu@users.noreply.github.com> Date: Mon, 4 Dec 2023 16:18:16 -0800 Subject: [PATCH 195/197] docs: correct the params rendering for `ml.remote` and `ml.ensemble` modules (#248) 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) - [x] Appropriate docs were updated (if necessary) - `ensemble.RandomForestClassifier`: https://screenshot.googleplex.com/4Q88xgdm5hkaYXu - `ensemble.RandomForestRegressor`: https://screenshot.googleplex.com/3CU6pJBjYHQvnDo - `remote.VertexAIModel`: https://screenshot.googleplex.com/8SL2max6GfPMwFe Fixes internal issue 314150462 🦕 --- bigframes/ml/remote.py | 8 +-- docs/templates/toc.yml | 12 ++-- .../sklearn/ensemble/_forest.py | 72 +++++++++---------- 3 files changed, 46 insertions(+), 46 deletions(-) diff --git a/bigframes/ml/remote.py b/bigframes/ml/remote.py index d4c34bbd0db..8da073802d7 100644 --- a/bigframes/ml/remote.py +++ b/bigframes/ml/remote.py @@ -47,10 +47,10 @@ class VertexAIModel(base.BaseEstimator): Args: endpoint (str): Vertex AI https endpoint. - input ({column_name: column_type}): - Input schema. Supported types are "bool", "string", "int64", "float64", "array", "array", "array", "array". - output ({column_name: column_type}): - Output label schema. Supported the same types as the input. + input (Mapping): + Input schema: `{column_name: column_type}`. Supported types are "bool", "string", "int64", "float64", "array", "array", "array", "array". + output (Mapping): + Output label schema: `{column_name: column_type}`. Supported the same types as the input. session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 58ac1c0efe6..b680a5fc1ad 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -108,12 +108,6 @@ - name: PaLM2TextEmbeddingGenerator uid: bigframes.ml.llm.PaLM2TextEmbeddingGenerator name: llm - - items: - - name: Overview - uid: bigframes.ml.remote - - name: VertexAIModel - uid: bigframes.ml.remote.VertexAIModel - name: remote - items: - name: metrics uid: bigframes.ml.metrics @@ -144,6 +138,12 @@ - name: OneHotEncoder uid: bigframes.ml.preprocessing.OneHotEncoder name: preprocessing + - items: + - name: Overview + uid: bigframes.ml.remote + - name: VertexAIModel + uid: bigframes.ml.remote.VertexAIModel + name: remote name: bigframes.ml name: BigQuery DataFrames status: beta diff --git a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py index 6be41bf9aa5..63c62274fd1 100644 --- a/third_party/bigframes_vendored/sklearn/ensemble/_forest.py +++ b/third_party/bigframes_vendored/sklearn/ensemble/_forest.py @@ -47,16 +47,16 @@ def fit(self, X, y): """Build a forest of trees from the training set (X, y). Args: - X: + 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): Series or DataFrame of shape (n_samples,) or (n_samples, n_targets). Target values. Will be cast to X's dtype if necessary. Returns: - Fitted Estimator. + ForestModel: Fitted Estimator. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -73,12 +73,12 @@ def predict(self, X): mean predicted regression targets of the trees in the forest. Args: - X: + 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: - The predicted values. + bigframes.dataframe.DataFrame: The predicted values. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -91,38 +91,38 @@ class RandomForestRegressor(ForestRegressor): to improve the predictive accuracy and control over-fitting. Args: - num_parallel_tree: Optional[int] + num_parallel_tree (Optional[int]): Number of parallel trees constructed during each iteration. Default to 100. Minimum value is 2. - 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. The value should be between 0 and 1. - colsample_bylevel : Optional[float] + 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] + 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] + 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 15. The value should be greater than 0 and less than 1. - subsample : Optional[float] + 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] + 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. - 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". """ @@ -144,7 +144,7 @@ def predict(self, X): which we want to get the predictions. Returns: - The predicted values. + bigframes.dataframe.DataFrame: The predicted values. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -158,37 +158,37 @@ class RandomForestClassifier(ForestClassifier): improve the predictive accuracy and control over-fitting. Args: - num_parallel_tree: Optional[int] + num_parallel_tree (Optional[int]): Number of parallel trees constructed during each iteration. Default to 100. Minimum value is 2. - 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. The value should be between 0 and 1. - colsample_bylevel : Optional[float] + 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] + 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] + 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 15. The value should be greater than 0 and less than 1. - subsample : Optional[float] + 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] + 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. - 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 77074ecbe7f52d1d7d1d1dc537fbe4062b407672 Mon Sep 17 00:00:00 2001 From: Huan Chen <142538604+Genesis929@users.noreply.github.com> Date: Mon, 4 Dec 2023 17:15:14 -0800 Subject: [PATCH 196/197] =?UTF-8?q?docs:=20add=20examples=20for=20datafram?= =?UTF-8?q?e.nunique,=20dataframe.diff,=20dataframe.a=E2=80=A6=20(#251)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: add examples for dataframe.nunique, dataframe.diff, dataframe.agg, dataframe.describe * update spacing * update ordering --- .../bigframes_vendored/pandas/core/frame.py | 114 +++++++++++++++++- 1 file changed, 112 insertions(+), 2 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 08fe8e2de0b..174ab069f68 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -3434,7 +3434,26 @@ def melt(self, id_vars, value_vars, var_name, value_name): def nunique(self): """ - Count number of distinct elements in specified axis. + Count number of distinct elements in each column. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 2]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 2 + + [3 rows x 2 columns] + + >>> df.nunique() + A 3.0 + B 2.0 + dtype: Float64 Returns: bigframes.series.Series: Series with number of distinct elements. @@ -3578,6 +3597,40 @@ def diff( Calculates the difference of a DataFrame element compared with another element in the DataFrame (default is element in previous row). + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + Calculating difference with default periods=1: + + >>> df.diff() + A B + 0 + 1 -2 1 + 2 1 1 + + [3 rows x 2 columns] + + Calculating difference with periods=-1: + + >>> df.diff(periods=-1) + A B + 0 2 -1 + 1 -1 -1 + 2 + + [3 rows x 2 columns] + Args: periods (int, default 1): Periods to shift for calculating difference, accepts negative @@ -3590,7 +3643,37 @@ def diff( def agg(self, func): """ - Aggregate using one or more operations over the specified axis. + Aggregate using one or more operations over columns. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [1, 2, 3]}) + >>> df + A B + 0 3 1 + 1 1 2 + 2 2 3 + + [3 rows x 2 columns] + + Using a single function: + + >>> df.agg('sum') + A 6.0 + B 6.0 + dtype: Float64 + + Using a list of functions: + + >>> df.agg(['sum', 'mean']) + A B + sum 6.0 6.0 + mean 2.0 2.0 + + [2 rows x 2 columns] Args: func (function): @@ -3623,6 +3706,33 @@ def describe(self): upper percentile is ``75``. The ``50`` percentile is the same as the median. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({"A": [3, 1, 2], "B": [0, 2, 8]}) + >>> df + A B + 0 3 0 + 1 1 2 + 2 2 8 + + [3 rows x 2 columns] + + >>> df.describe() + A B + count 3.0 3.0 + mean 2.0 3.333333 + std 1.0 4.163332 + min 1.0 0.0 + 25% 1.0 0.0 + 50% 2.0 2.0 + 75% 3.0 8.0 + max 3.0 8.0 + + [8 rows x 2 columns] + Returns: bigframes.dataframe.DataFrame: Summary statistics of the Series or Dataframe provided. """ From 89a1c67fa5cbb76c1cc6ae24d5f919e22514705c Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 5 Dec 2023 02:14:14 +0000 Subject: [PATCH 197/197] docs: Fix return annotation in API docstrings (#253) 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 internal issue 314367409 🦕 --- .../bigframes_vendored/pandas/core/frame.py | 30 +++++++-------- .../bigframes_vendored/pandas/core/series.py | 38 +++++++++---------- 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 174ab069f68..7168572705a 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -226,7 +226,7 @@ def from_dict( if used with ``orient='columns'`` or ``orient='tight'``. Returns: - DataFrame + DataFrame: DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -267,7 +267,7 @@ def from_records( Number of rows to read if data is an iterator. Returns: - DataFrame + DataFrame: DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -717,7 +717,7 @@ def to_markdown( These parameters will be passed to `tabulate `_. Returns: - DataFrame in Markdown-friendly format. + DataFrame: DataFrame in Markdown-friendly format. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1272,7 +1272,7 @@ def sort_values( if `first`; `last` puts NaNs at the end. Returns: - DataFrame with sorted values. + DataFrame: DataFrame with sorted values. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1282,7 +1282,7 @@ def sort_index( """Sort object by labels (along an axis). Returns: - The original DataFrame sorted by the labels. + DataFrame: The original DataFrame sorted by the labels. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1330,7 +1330,7 @@ def eq(self, other, axis: str | int = "columns") -> DataFrame: (1 or 'columns'). Returns: - Result of the comparison. + DataFrame: Result of the comparison. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1861,7 +1861,7 @@ def rtruediv(self, other, axis: str | int = "columns") -> DataFrame: (1 or 'columns'). For Series input, axis to match Series index on. Returns: - DataFrame result of the arithmetic operation. + DataFrame: DataFrame result of the arithmetic operation. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -2796,7 +2796,7 @@ def any(self, *, axis=0, bool_only: bool = False): Include only boolean columns. Returns: - Series + bigframes.series.Series: Series indicating if any element is True per column. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -2843,7 +2843,7 @@ def all(self, axis=0, *, bool_only: bool = False): Include only boolean columns. Returns: - bigframes.series.Series: Series if all elements are True. + bigframes.series.Series: Series indicating if all elements are True per column. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -3184,7 +3184,7 @@ def skew(self, *, numeric_only: bool = False): Include only float, int, boolean columns. Returns: - Series + Series: Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -3225,7 +3225,7 @@ def kurt(self, *, numeric_only: bool = False): Include only float, int, boolean columns. Returns: - Series + Series: Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -3770,7 +3770,7 @@ def pivot(self, *, columns, index=None, values=None): have hierarchically indexed columns. Returns: - Returns reshaped DataFrame. + DataFrame: Returns reshaped DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -3808,7 +3808,7 @@ def unstack(self): (the analogue of stack when the columns are not a MultiIndex). Returns: - DataFrame or Series + DataFrame or Series: DataFrame or Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -3866,7 +3866,7 @@ def index(self): dtype=object) Returns: - The index labels of the DataFrame. + Index: The index object of the DataFrame. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -4115,7 +4115,7 @@ def dot(self, other): The other object to compute the matrix product with. Returns: - Series or DataFrame + Series or DataFrame: If `other` is a Series, return the matrix product between self and other as a Series. If other is a DataFrame, return the matrix product of self and other in a DataFrame. diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 1b751ed83b6..6b8dd1d64d6 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -90,7 +90,7 @@ def index(self): dtype=object) Returns: - The index labels of the Series. + Index: The index object of the Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -138,7 +138,7 @@ def transpose(self) -> Series: Return the transpose, which is by definition self. Returns: - Series + Series: Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -593,7 +593,7 @@ def corr(self, other, method="pearson", min_periods=None) -> float: 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 + 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") @@ -611,7 +611,7 @@ def diff(self) -> Series: values. Returns: - {klass}: First differences of the Series. + Series: First differences of the Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1317,7 +1317,7 @@ def le(self, other) -> Series: other: Series, or scalar value Returns: - bigframes.series.Series. The result of the comparison. + bigframes.series.Series: The result of the comparison. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1331,7 +1331,7 @@ def lt(self, other) -> Series: Args: other (Series, or scalar value): - Returns: + Returns: bigframes.series.Series: The result of the operation. """ @@ -1588,7 +1588,7 @@ def divmod(self, other) -> Series: other: Series, or scalar value Returns: - 2-Tuple of Series. The result of the operation. The result is always + 2-Tuple of Series: The result of the operation. The result is always consistent with (floordiv, mod) (though pandas may not). """ @@ -1603,7 +1603,7 @@ def rdivmod(self, other) -> Series: other: Series, or scalar value Returns: - 2-Tuple of Series. The result of the operation. The result is always + 2-Tuple of Series: The result of the operation. The result is always consistent with (rfloordiv, rmod) (though pandas may not). """ @@ -1650,7 +1650,7 @@ def max( Returns: - scalar or scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1664,7 +1664,7 @@ def min( of the ``numpy.ndarray`` method ``argmin``. Returns: - scalar or scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1692,7 +1692,7 @@ def var( Normalized by N-1 by default. Returns: - scalar or Series (if level specified) + scalar or Series (if level specified): Variance. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1702,7 +1702,7 @@ def sum(self): This is equivalent to the method ``numpy.sum``. Returns: - scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1710,7 +1710,7 @@ def mean(self): """Return the mean of the values over the requested axis. Returns: - scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1723,7 +1723,7 @@ def median(self, *, exact: bool = False): one. Note: ``exact=True`` not yet supported. Returns: - scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1731,7 +1731,7 @@ def prod(self): """Return the product of the values over the requested axis. Returns: - scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1741,7 +1741,7 @@ def skew(self): Normalized by N-1. Returns: - scalar + scalar: Scalar. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -1932,7 +1932,7 @@ def clip(self): 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. + Series: Series. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -2083,7 +2083,7 @@ def is_monotonic_increasing(self) -> bool: Return boolean if values in the object are monotonically increasing. Returns: - bool + bool: Boolean. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -2093,7 +2093,7 @@ def is_monotonic_decreasing(self) -> bool: Return boolean if values in the object are monotonically decreasing. Returns: - bool + bool: Boolean. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)