feat: include bq schema and query string in dry run results (#1752)

* feat: include bq schema and query string in dry run results

* rename key

* fix tests

Author: sycai

bigframes/session/dry_runs.py

@@ -37,6 +37,10 @@ def get_table_stats(table: bigquery.Table) -> pandas.Series:
     index.append("columnDtypes")
     values.append(col_dtypes)
 
+    # Add raw BQ schema
+    index.append("bigquerySchema")
+    values.append(table.schema)
+
     for key in ("numBytes", "numRows", "location", "type"):
         index.append(key)
         values.append(table._properties[key])
@@ -96,8 +100,12 @@ def get_query_stats(
 ) -> pandas.Series:
     """Returns important stats from the query job as a Pandas Series."""
 
-    index = []
-    values = []
+    index: List[Any] = []
+    values: List[Any] = []
+
+    # Add raw BQ schema
+    index.append("bigquerySchema")
+    values.append(query_job.schema)
 
     job_api_repr = copy.deepcopy(query_job._properties)
 
@@ -110,6 +118,8 @@ def get_query_stats(
     configuration = job_api_repr.get("configuration", {})
     index.append("jobType")
     values.append(configuration.get("jobType", None))
+    index.append("dispatchedSql")
+    values.append(configuration.get("query", {}).get("query", None))
 
     query_config = configuration.get("query", {})
     for key in ("destinationTable", "useLegacySql"):

tests/system/small/test_dataframe_io.py

@@ -319,7 +319,8 @@ def test_to_pandas_dry_run(session, scalars_pandas_df_multi_index):
 
     result = bf_df.to_pandas(dry_run=True)
 
-    assert len(result) == 14
+    assert isinstance(result, pd.Series)
+    assert len(result) > 0
 
 
 def test_to_arrow_override_global_option(scalars_df_index):

tests/system/small/test_index_io.py

@@ -11,6 +11,8 @@
 # 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
 
 
@@ -35,7 +37,8 @@ def test_to_pandas_dry_run(scalars_df_index):
 
     result = index.to_pandas(dry_run=True)
 
-    assert len(result) == 14
+    assert isinstance(result, pd.Series)
+    assert len(result) > 0
 
 
 def test_to_numpy_override_global_option(scalars_df_index):

tests/system/small/test_series.py

@@ -4604,4 +4604,5 @@ def test_series_to_pandas_dry_run(scalars_df_index):
 
     result = bf_series.to_pandas(dry_run=True)
 
-    assert len(result) == 14
+    assert isinstance(result, pd.Series)
+    assert len(result) > 0

tests/system/small/test_session.py

@@ -1899,9 +1899,11 @@ def _assert_query_dry_run_stats_are_valid(result: pd.Series):
             "columnDtypes",
             "indexLevel",
             "indexDtypes",
+            "bigquerySchema",
             "projectId",
             "location",
             "jobType",
+            "dispatchedSql",
             "destinationTable",
             "useLegacySql",
             "referencedTables",
@@ -1922,6 +1924,7 @@ def _assert_table_dry_run_stats_are_valid(result: pd.Series):
             "isQuery",
             "columnCount",
             "columnDtypes",
+            "bigquerySchema",
             "numBytes",
             "numRows",
             "location",