return some results in SCRIPT if the last statement is a SELECT

Author: tswast

bigframes/session/_io/bigquery/read_gbq_query.py

@@ -32,6 +32,28 @@
 import bigframes.session
 
 
+def should_return_query_results(query_job: bigquery.QueryJob) -> bool:
+    """Returns True if query_job is the kind of query we expect results from.
+
+    If the query was DDL or DML, return some job metadata. See
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobStatistics2.FIELDS.statement_type
+    for possible statement types. Note that destination table does exist
+    for some DDL operations such as CREATE VIEW, but we don't want to
+    read from that. See internal issue b/444282709.
+    """
+
+    if query_job.statement_type == "SELECT":
+        return True
+
+    if query_job.statement_type == "SCRIPT":
+        # Try to determine if the last statement is a SELECT. Alternatively, we
+        # could do a jobs.list request using query_job as the parent job and
+        # try to determine the statement type of the last child job.
+        return query_job.destination != query_job.ddl_target_table
+
+    return False
+
+
 def create_dataframe_from_query_job_stats(
     query_job: Optional[bigquery.QueryJob], *, session: bigframes.session.Session
 ) -> dataframe.DataFrame:

bigframes/session/loader.py

@@ -1045,29 +1045,30 @@ def read_gbq_query(
         # local node. Likely there are a wide range of sizes in which it
         # makes sense to download the results beyond the first page, even if
         # there is a job and destination table available.
-        if (
-            rows is not None
-            and destination is None
-            and (
-                query_job_for_metrics is None
-                or query_job_for_metrics.statement_type == "SELECT"
-            )
-        ):
+        if query_job_for_metrics is None and rows is not None:
             return bf_read_gbq_query.create_dataframe_from_row_iterator(
                 rows,
                 session=self._session,
                 index_col=index_col,
                 columns=columns,
             )
 
+        # We already checked rows, so if there's no destination table, then
+        # there are no results to return.
+        if destination is None:
+            return bf_read_gbq_query.create_dataframe_from_query_job_stats(
+                query_job_for_metrics,
+                session=self._session,
+            )
+
         # If the query was DDL or DML, return some job metadata. See
         # https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobStatistics2.FIELDS.statement_type
         # for possible statement types. Note that destination table does exist
         # for some DDL operations such as CREATE VIEW, but we don't want to
         # read from that. See internal issue b/444282709.
-        if destination is None or (
+        if (
             query_job_for_metrics is not None
-            and query_job_for_metrics.statement_type != "SELECT"
+            and not bf_read_gbq_query.should_return_query_results(query_job_for_metrics)
         ):
             return bf_read_gbq_query.create_dataframe_from_query_job_stats(
                 query_job_for_metrics,