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",