Skip to content

perf: defer query in read_gbq with wildcard tables #1661

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 33 commits into from
May 6, 2025
Merged

Conversation

tswast
Copy link
Collaborator

@tswast tswast commented Apr 27, 2025

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 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 405773140 🦕

@tswast tswast requested review from a team as code owners April 27, 2025 03:09
@tswast tswast requested a review from drylks-work April 27, 2025 03:09
@product-auto-label product-auto-label bot added size: m Pull request size is medium. api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. labels Apr 27, 2025
@tswast tswast requested review from Genesis929 and removed request for drylks-work April 27, 2025 03:10
@tswast
Copy link
Collaborator Author

tswast commented Apr 28, 2025

Failures look like real failures.

        if not 200 <= response.status_code < 300:
>           raise exceptions.from_http_response(response)
E           google.api_core.exceptions.BadRequest: 400 GET https://bigquery.googleapis.com/bigquery/v2/projects/python-docs-samples-tests/queries/32f42306-e95f-48bc-a2fb-56761aec5476?maxResults=0&location=US&prettyPrint=false: Invalid field name "_TABLE_SUFFIX". Field names are not allowed to start with the (case-insensitive) prefixes _PARTITION, _TABLE_, _FILE_, _ROW_TIMESTAMP, __ROOT__ and _COLIDENTIFIER
E           
E           Location: US
E           Job ID: 32f42306-e95f-48bc-a2fb-56761aec5476

While we can query such fields, it looks like we can't materialize them.

@GarrettWu GarrettWu removed their assignment Apr 28, 2025
Genesis929
Genesis929 previously approved these changes Apr 28, 2025
@Genesis929 Genesis929 self-requested a review April 28, 2025 20:50
@Genesis929 Genesis929 dismissed their stale review April 28, 2025 20:52

Seems tests are affected

tswast added 2 commits April 28, 2025 16:57
…g pseudocolumns

Fixes this code sample:

import bigframes.pandas as bpd
df = bpd.read_gbq("bigquery-public-data.google_analytics_sample.ga_sessions_*")
df[df["_TABLE_SUFFIX"] == "20161204"].peek()
@tswast tswast added the do not merge Indicates a pull request not ready for merge, due to either quality or timing. label Apr 28, 2025
@tswast
Copy link
Collaborator Author

tswast commented Apr 28, 2025

Added do not merge. Need to make sure this is compatible with to_gbq() and cached().

@tswast
Copy link
Collaborator Author

tswast commented Apr 28, 2025

From the notebook tests:

File /[tmpfs/src/github/python-bigquery-dataframes/bigframes/core/nodes.py:711](https://cs.corp.google.com/piper///depot/google3/tmpfs/src/github/python-bigquery-dataframes/bigframes/core/nodes.py?l=711), in GbqTable.from_table(table, columns)
    708 @staticmethod
    709 def from_table(table: bq.Table, columns: Sequence[str] = ()) -> GbqTable:
    710     # Subsetting fields with columns can reduce cost of row-hash default ordering
--> 711     table_schema = bigframes.core.tools.bigquery.get_schema_and_pseudocolumns(table)
    713     if columns:
    714         schema = tuple(item for item in table_schema if item.name in columns)

AttributeError: module 'bigframes.core.tools' has no attribute 'bigquery'

Looks like we're missing some imports too.

@tswast tswast removed the do not merge Indicates a pull request not ready for merge, due to either quality or timing. label Apr 29, 2025
@tswast
Copy link
Collaborator Author

tswast commented Apr 29, 2025

Tested the failing samples tests locally. I think my latest commits solve the issue of not being able to materialize to _TABLE_SUFFIX as the column name.

@tswast
Copy link
Collaborator Author

tswast commented Apr 29, 2025

e2e and notebook failures are for the same:

E               google.api_core.exceptions.BadRequest: 400 'FOR SYSTEM_TIME AS OF' expression for table 'bigframes-load-testing.bigframes_testing.penguins_dcdc3525965d3bf2805a055ee80a0ae7' evaluates to a TIMESTAMP value in the future: 2025-04-29 15:16:24.477885 UTC.; reason: invalidQuery, location: query, message: 'FOR SYSTEM_TIME AS OF' expression for table 'bigframes-load-testing.bigframes_testing.penguins_dcdc3525965d3bf2805a055ee80a0ae7' evaluates to a TIMESTAMP value in the future: 2025-04-29 15:16:24.477885 UTC.

I don't think these relate to this change, but I do recall BQML having a hard time with time travel. Potentially we're missing a force cache somewhere now?

@tswast
Copy link
Collaborator Author

tswast commented Apr 29, 2025

Notebook failures appear to be flakes, as they are related to remote functions and succeeded in 3.10 but not 3.11.

nox > * notebook-3.10: success
nox > * notebook-3.11: failed

Comment on lines 368 to 370
sql, schema, ordering_info = self.compiler.compile_raw(
self.logical_plan(array_value.node)
self.logical_plan(renamed.node)
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe there is a better way, where the compiler just does not return invalid sql ever?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that's possible. This is valid SQL before the renaming. It's just restricted as to where we can materialize.

Comment on lines 376 to 383
cached_replacement = (
renamed.as_cached(
cache_table=self.bqclient.get_table(tmp_table),
ordering=ordering_info,
)
.rename_columns(dict(zip(new_col_ids, prev_col_ids)))
.node
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can apply the renaming as part of the scan definition of the cache node rather than as a selection node (via rename_columns) on top of it. Should be more robust that way

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried that in 5b0d0a0

The problem I was having is that the physical schema often contains more columns than the array value. If we rename there, then we lose track of which columns belong to which id.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a "renames" argument to as_cached, which I does allow me to use the scan list to rename columns instead. Indeed that does seem more robust. I think some of these renames might have been pruned out in recent commits.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moving this change to #1667

@tswast tswast requested a review from TrevorBergeron April 30, 2025 14:26
@tswast
Copy link
Collaborator Author

tswast commented Apr 30, 2025

e2e failure might indicate a real issue FAILED tests/system/large/operations/test_semantics.py::test_sim_join[has_score_column]

>           raise exceptions.from_http_response(response)
E           google.api_core.exceptions.BadRequest: 400 GET [https://bigquery.googleapis.com/bigquery/v2/projects/bigframes-load-testing/queries/e2e2cede-621f-47fc-9fbc-4077c0579eaf?maxResults=0&location=US&prettyPrint=false](https://www.google.com/url?q=https://bigquery.googleapis.com/bigquery/v2/projects/bigframes-load-testing/queries/e2e2cede-621f-47fc-9fbc-4077c0579eaf?maxResults%3D0%26location%3DUS%26prettyPrint%3Dfalse&sa=D): Duplicate column names in the result are not supported when a destination table is present. Found duplicate(s): creatures

@tswast
Copy link
Collaborator Author

tswast commented Apr 30, 2025

Different set of e2e failures this time:

FAILED tests/system/large/blob/test_function.py::test_blob_pdf_chunk[True-expected0]
FAILED tests/system/large/blob/test_function.py::test_blob_image_normalize_to_series
FAILED tests/system/large/blob/test_function.py::test_blob_image_blur_to_series
FAILED tests/system/large/blob/test_function.py::test_blob_image_normalize_to_folder
FAILED tests/system/large/blob/test_function.py::test_blob_image_blur_to_folder
FAILED tests/system/large/blob/test_function.py::test_blob_image_normalize_to_bq
FAILED tests/system/large/blob/test_function.py::test_blob_image_blur_to_bq
FAILED tests/system/large/blob/test_function.py::test_blob_image_resize_to_series
ERROR tests/system/large/blob/test_function.py::test_blob_pdf_extract[True-expected0]

Maybe it was just an LLM-related flake last time?

@tswast
Copy link
Collaborator Author

tswast commented May 5, 2025

I'll split to defer part from the pseudocolumn work into #1689

I think the defer step can be done without out that to make this PR a lot smaller.

@product-auto-label product-auto-label bot added size: s Pull request size is small. and removed size: l Pull request size is large. labels May 5, 2025
@tswast tswast merged commit 5c125c9 into main May 6, 2025
23 of 24 checks passed
@tswast tswast deleted the b405773140-wildcard branch May 6, 2025 17:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. size: s Pull request size is small.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants