Feat: implement support for bigframes

[z3z1ma]

I PR'ed the upstream lib so that they have no upper bound on their SQLGlot pin. This means the library is now fully compatible with SQLMesh. This PR adds support. Instead of bundling the extra with bigquery, I have temporarily added it to its own extra so it is opt-in. Once they cut a version release, we can add it with a proper >= version specifier.

This enables many novel uses such as using Gemini / LLM capabilities at scale more effectively, using ML models (training , predicting, etc), and so on leveraging BigQuery's compute through a dataframe interface wrapped intuitively by sqlmesh's python model system.

https://cloud.google.com/bigquery/docs/use-bigquery-dataframes#pandas-examples

https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/getting_started/ml_fundamentals_bq_dataframes.ipynb

I have tested this locally. I was able to crunch through 35G in 15 seconds with no local computation. Its very nice.

The following example requires no seeds or pre-loaded data as it leverages BQ's open datasets. You can run this with or without the remote function.

"""Working example to get your feet wet"""
import typing as t
from datetime import datetime

from bigframes.pandas import DataFrame

from sqlmesh import ExecutionContext, model


def get_bucket(num: int):
    if not num:
        return "NA"
    boundary = 10
    return "at_or_above_10" if num >= boundary else "below_10"


@model(
    "mart.wiki",
    columns={
        "title": "text",
        "views": "int",
        "bucket": "text",
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    execution_time: datetime,
    **kwargs: t.Any,
) -> DataFrame:
    # Create a remote function to be used in the Bigframe DataFrame
    remote_get_bucket = context.bigframe.remote_function([int], str)(get_bucket)

    # Returns the Bigframe DataFrame handle, no data is computed locally
    df = context.bigframe.read_gbq("bigquery-samples.wikipedia_pageviews.200809h")

    df = (
        # This runs entirely on the BigQuery engine lazily
        df[df.title.str.contains(r"[Gg]oogle")]
        .groupby(["title"], as_index=False)["views"]
        .sum(numeric_only=True)
        .sort_values("views", ascending=False)
    )

    return df.assign(bucket=df["views"].apply(remote_get_bucket))

Support was simple given the way the query_factory works and similar art with Snowpark. Very straightforward. Good work on that.

You can even do interesting things like include a read from a gcs bucket in BQ engine:

filepath_or_buffer = "gs://cloud-samples-data/bigquery/us-states/us-states.csv"
df_from_gcs = bpd.read_csv(filepath_or_buffer)

[georgesittas]

Awesome, thanks for driving this! LGTM, just one question to make sure the imports are safe. Will let others chime in here too.

[z3z1ma]

FYI @tobymao this PR is good to merge whenever you all want to, the test failure is completely unrelated I think.

[georgesittas]

FYI @tobymao this PR is good to merge whenever you all want to, the test failure is completely unrelated I think.

Yeah, this test is flaky for some reason, no need to worry about it.