diff --git a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index f56472f1b5..c7c302e5a9 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -49,7 +49,13 @@ " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb index d49a44a780..b9a9c6723b 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb @@ -46,7 +46,13 @@ " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb b/notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb index 32e6cb7924..4c589abf95 100644 --- a/notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_ml_drug_name_generation.ipynb @@ -48,7 +48,13 @@ " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/getting_started/getting_started_bq_dataframes.ipynb b/notebooks/getting_started/getting_started_bq_dataframes.ipynb index e79db455fc..f0cc61d8d2 100644 --- a/notebooks/getting_started/getting_started_bq_dataframes.ipynb +++ b/notebooks/getting_started/getting_started_bq_dataframes.ipynb @@ -49,7 +49,13 @@ " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/getting_started/ml_fundamentals_bq_dataframes.ipynb b/notebooks/getting_started/ml_fundamentals_bq_dataframes.ipynb index e7b69f017b..d8e277d6a8 100644 --- a/notebooks/getting_started/ml_fundamentals_bq_dataframes.ipynb +++ b/notebooks/getting_started/ml_fundamentals_bq_dataframes.ipynb @@ -49,7 +49,13 @@ " \"Vertex\n", " Open in Vertex AI Workbench\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb b/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb index b6c9029a01..c450f563c2 100644 --- a/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb +++ b/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb @@ -18,7 +18,13 @@ " \"GitHub\n", " View on GitHub\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb b/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb index 40efe9d18c..9e44a0a5c1 100644 --- a/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb +++ b/notebooks/visualization/bq_dataframes_covid_line_graphs.ipynb @@ -43,7 +43,13 @@ " \"GitHub\n", " View on GitHub\n", " \n", - " \n", + " \n", + " \n", + " \n", + " \"BQ\n", + " Open in BQ Studio\n", + " \n", + " \n", "" ] }, diff --git a/third_party/bigframes_vendored/tpch/queries/q21.py b/third_party/bigframes_vendored/tpch/queries/q21.py index b844cd6f6b..097a730d43 100644 --- a/third_party/bigframes_vendored/tpch/queries/q21.py +++ b/third_party/bigframes_vendored/tpch/queries/q21.py @@ -1,6 +1,4 @@ -# Contains code from https://github.com/pola-rs/tpch/blob/main/queries/duckdb/q21.py - -import typing +# Contains code from https://github.com/pola-rs/tpch/blob/main/queries/polars/q21.py import bigframes import bigframes.pandas as bpd @@ -24,39 +22,38 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session): index_col=bigframes.enums.DefaultIndexKind.NULL, ) - nation = nation[nation["N_NAME"] == "SAUDI ARABIA"] - orders = orders[orders["O_ORDERSTATUS"] == "F"] - - l1 = lineitem[lineitem["L_RECEIPTDATE"] > lineitem["L_COMMITDATE"]][ - ["L_ORDERKEY", "L_SUPPKEY"] - ] + var1 = "SAUDI ARABIA" - l2 = lineitem.groupby("L_ORDERKEY", as_index=False).agg( - NUNIQUE_COL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="nunique") + q1 = lineitem.groupby("L_ORDERKEY", as_index=False).agg( + N_SUPP_BY_ORDER=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size") ) - l2 = l2[l2["NUNIQUE_COL"] > 1][["L_ORDERKEY"]] + q1 = q1[q1["N_SUPP_BY_ORDER"] > 1] - l3 = l1.groupby("L_ORDERKEY", as_index=False).agg( - NUNIQUE_COL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="nunique") - ) - l3 = l3[l3["NUNIQUE_COL"] == 1][["L_ORDERKEY"]] + lineitem_filtered = lineitem[lineitem["L_RECEIPTDATE"] > lineitem["L_COMMITDATE"]] - l1 = l1.merge(l2, on="L_ORDERKEY", how="inner").merge( - l3, on="L_ORDERKEY", how="inner" + q1 = q1.merge(lineitem_filtered, on="L_ORDERKEY") + + q_final = q1.groupby("L_ORDERKEY", as_index=False).agg( + N_SUPP_BY_ORDER_FINAL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size") ) - merged = supplier.merge(nation, left_on="S_NATIONKEY", right_on="N_NATIONKEY") - merged = merged.merge(l1, left_on="S_SUPPKEY", right_on="L_SUPPKEY") - merged = merged.merge(orders, left_on="L_ORDERKEY", right_on="O_ORDERKEY") + q_final = q_final.merge(q1, on="L_ORDERKEY") + q_final = q_final.merge(supplier, left_on="L_SUPPKEY", right_on="S_SUPPKEY") + q_final = q_final.merge(nation, left_on="S_NATIONKEY", right_on="N_NATIONKEY") + q_final = q_final.merge(orders, left_on="L_ORDERKEY", right_on="O_ORDERKEY") + + q_final = q_final[ + (q_final["N_SUPP_BY_ORDER_FINAL"] == 1) + & (q_final["N_NAME"] == var1) + & (q_final["O_ORDERSTATUS"] == "F") + ] - result = merged.groupby("S_NAME", as_index=False).agg( + q_final = q_final.groupby("S_NAME", as_index=False).agg( NUMWAIT=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size") ) - result = ( - typing.cast(bpd.DataFrame, result) - .sort_values(["NUMWAIT", "S_NAME"], ascending=[False, True]) - .head(100) - ) + q_final = q_final.sort_values( + by=["NUMWAIT", "S_NAME"], ascending=[False, True] + ).head(100) - result.to_gbq() + q_final.to_gbq()