chore: Make graph visualization more scalable

We now only include query results directly in the html when the query results are less than 100 KB. For larger query results, we store only the reference to the destination table in the HTML, and have the python code re-read the query results from the destination table during the callback.

Also, added a hard limit of 5 MB in the query result size, beyond which, graph visualization is not supported altogether.

Author: ericfe-google

bigquery_magics/bigquery.py

@@ -117,7 +117,6 @@
 import IPython  # type: ignore
 from IPython.core import magic_arguments  # type: ignore
 from IPython.core.getipython import get_ipython
-from google.api_core import client_info
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
 from google.cloud.bigquery import exceptions
@@ -126,13 +125,11 @@
 from google.cloud.bigquery.job import QueryJobConfig
 import pandas
 
-from bigquery_magics import environment
 from bigquery_magics import line_arg_parser as lap
 import bigquery_magics._versions_helpers
 import bigquery_magics.config
 import bigquery_magics.graph_server as graph_server
 import bigquery_magics.pyformat
-import bigquery_magics.version
 
 try:
     from google.cloud import bigquery_storage  # type: ignore
@@ -147,24 +144,6 @@
 context = bigquery_magics.config.context
 
 
-def _get_user_agent():
-    identities = [
-        f"ipython-{IPython.__version__}",
-        f"bigquery-magics/{bigquery_magics.version.__version__}",
-    ]
-
-    if environment.is_vscode():
-        identities.append("vscode")
-        if environment.is_vscode_google_cloud_code_extension_installed():
-            identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
-    elif environment.is_jupyter():
-        identities.append("jupyter")
-        if environment.is_jupyter_bigquery_plugin_installed():
-            identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
-
-    return " ".join(identities)
-
-
 def _handle_error(error, destination_var=None):
     """Process a query execution error.
 
@@ -565,23 +544,7 @@ def _query_with_pandas(query: str, params: List[Any], args: Any):
 
 
 def _create_clients(args: Any) -> Tuple[bigquery.Client, Any]:
-    bigquery_client_options = copy.deepcopy(context.bigquery_client_options)
-    if args.bigquery_api_endpoint:
-        if isinstance(bigquery_client_options, dict):
-            bigquery_client_options["api_endpoint"] = args.bigquery_api_endpoint
-        else:
-            bigquery_client_options.api_endpoint = args.bigquery_api_endpoint
-
-    bq_client = bigquery.Client(
-        project=args.project or context.project,
-        credentials=context.credentials,
-        default_query_job_config=context.default_query_job_config,
-        client_info=client_info.ClientInfo(user_agent=_get_user_agent()),
-        client_options=bigquery_client_options,
-        location=args.location,
-    )
-    if context._connection:
-        bq_client._connection = context._connection
+    bq_client = graph_server.create_bq_client(args)
 
     # Check and instantiate bq storage client
     if args.use_bqstorage_api is not None:
@@ -633,8 +596,9 @@ def _handle_result(result, args):
 
 
 def _colab_query_callback(query: str, params: str):
+    parsed_params = json.loads(params)
     return IPython.core.display.JSON(
-        graph_server.convert_graph_data(query_results=json.loads(params))
+        graph_server.convert_graph_data(query_results=parsed_params["query_result"])
     )
 
 
@@ -663,7 +627,51 @@ def _colab_node_expansion_callback(request: dict, params_str: str):
 singleton_server_thread: threading.Thread = None
 
 
-def _add_graph_widget(query_result):
+MAX_GRAPH_VISUALIZATION_SIZE = 5000000
+MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100000
+
+
+def _estimate_json_size(df: pandas.DataFrame) -> int:
+    """Approximates the length of df.to_json(orient='records')
+    without materializing the string.
+    """
+    num_rows, num_cols = df.shape
+    if num_rows == 0:
+        return 2  # "[]"
+
+    # 1. Key overhead: "column_name": (repeated for every row)
+    # Includes quotes, colon, and comma separator per field
+    key_overhead = sum(len(f'"{col}":') + 1 for col in df.columns) * num_rows
+
+    # 2. Row structural overhead: { } per row and [ ] for the list
+    # Plus commas between rows (num_rows - 1)
+    structural_overhead = (2 * num_rows) + 2 + (num_rows - 1)
+
+    # 3. Value lengths
+    total_val_len = 0
+    for col in df.columns:
+        series = df[col]
+
+        if pandas.api.types.is_numeric_dtype(series):
+            # Numeric values (no quotes). Sample for average length to save memory.
+            sample_size = min(len(series), 1000)
+            avg_len = series.sample(sample_size).astype(str).str.len().mean()
+            total_val_len += avg_len * num_rows
+
+        elif pandas.api.types.is_bool_dtype(series):
+            # true (4) or false (5)
+            total_val_len += series.map({True: 4, False: 5}).sum()
+
+        else:
+            # Strings/Objects: "value" + quotes (2) + rough escaping factor
+            # .str.len() is relatively memory-efficient
+            val_chars = series.astype(str).str.len().sum()
+            total_val_len += val_chars + (2 * num_rows)
+
+    return int(key_overhead + structural_overhead + total_val_len)
+
+
+def _add_graph_widget(query_result: pandas.DataFrame, query_job: Any, args: Any):
     try:
         from spanner_graphs.graph_visualization import generate_visualization_html
     except ImportError as err:
@@ -698,10 +706,36 @@ def _add_graph_widget(query_result):
         port = graph_server.graph_server.port
 
     # Create html to invoke the graph server
+    args_dict = {
+        "bigquery_api_endpoint": args.bigquery_api_endpoint,
+        "project": args.project,
+        "location": args.location,
+    }
+
+    estimated_size = _estimate_json_size(query_result)
+    if estimated_size > MAX_GRAPH_VISUALIZATION_SIZE:
+        IPython.display.display(
+            IPython.core.display.HTML(
+                "<big><b>Error:</b> The query result is too large for graph visualization.</big>"
+            )
+        )
+        return
+
+    table_dict = {
+        "projectId": query_job.configuration.destination.project,
+        "datasetId": query_job.configuration.destination.dataset_id,
+        "tableId": query_job.configuration.destination.table_id,
+    }
+
+    params_dict = {"destination_table": table_dict, "args": args_dict}
+    if estimated_size < MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE:
+        params_dict["query_result"] = json.loads(query_result.to_json())
+    
+    params_str = json.dumps(params_dict)
     html_content = generate_visualization_html(
         query="placeholder query",
         port=port,
-        params=query_result.to_json().replace("\\", "\\\\").replace('"', '\\"'),
+        params=params_str.replace("\\", "\\\\").replace('"', '\\"'),
     )
     IPython.display.display(IPython.core.display.HTML(html_content))
 
@@ -810,7 +844,7 @@ def _make_bq_query(
         result = result.to_dataframe(**dataframe_kwargs)
 
     if args.graph and _supports_graph_widget(result):
-        _add_graph_widget(result)
+        _add_graph_widget(result, query_job, args)
     return _handle_result(result, args)
 
 
@@ -904,7 +938,9 @@ def _make_bqstorage_client(client, client_options):
 
     return client._ensure_bqstorage_client(
         client_options=client_options,
-        client_info=gapic_client_info.ClientInfo(user_agent=_get_user_agent()),
+        client_info=gapic_client_info.ClientInfo(
+            user_agent=graph_server._get_user_agent()
+        ),
     )
 
 

bigquery_magics/graph_server.py

@@ -13,12 +13,72 @@
 # limitations under the License.
 
 import atexit
+import copy
 import http.server
 import json
 import socketserver
 import threading
 from typing import Any, Dict, List
 
+import IPython  # type: ignore
+from google.api_core import client_info
+from google.cloud import bigquery
+
+from bigquery_magics import environment
+import bigquery_magics.config
+import bigquery_magics.version
+
+
+context = bigquery_magics.config.context
+
+
+def _get_user_agent():
+    identities = [
+        f"ipython-{IPython.__version__}",
+        f"bigquery-magics/{bigquery_magics.version.__version__}",
+    ]
+
+    if environment.is_vscode():
+        identities.append("vscode")
+        if environment.is_vscode_google_cloud_code_extension_installed():
+            identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
+    elif environment.is_jupyter():
+        identities.append("jupyter")
+        if environment.is_jupyter_bigquery_plugin_installed():
+            identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
+
+    return " ".join(identities)
+
+
+def create_bq_client(args: Any):
+    """Creates a BigQuery client.
+
+    Args:
+        args (argparse.Namespace): The arguments passed to the cell magic.
+
+    Returns:
+        google.cloud.bigquery.client.Client: The BigQuery client.
+    """
+    bigquery_client_options = copy.deepcopy(context.bigquery_client_options)
+    if args.bigquery_api_endpoint:
+        if isinstance(bigquery_client_options, dict):
+            bigquery_client_options["api_endpoint"] = args.bigquery_api_endpoint
+        else:
+            bigquery_client_options.api_endpoint = args.bigquery_api_endpoint
+
+    bq_client = bigquery.Client(
+        project=args.project or context.project,
+        credentials=context.credentials,
+        default_query_job_config=context.default_query_job_config,
+        client_info=client_info.ClientInfo(user_agent=_get_user_agent()),
+        client_options=bigquery_client_options,
+        location=args.location,
+    )
+    if context._connection:
+        bq_client._connection = context._connection
+
+    return bq_client
+
 
 def execute_node_expansion(params, request):
     return {"error": "Node expansion not yet implemented"}
@@ -251,7 +311,23 @@ def handle_post_ping(self):
 
     def handle_post_query(self):
         data = self.parse_post_data()
-        response = convert_graph_data(query_results=json.loads(data["params"]))
+        params = json.loads(data["params"])
+
+        query_results = None
+        if "query_result" in params:
+            query_results = params["query_result"]
+        else:
+            class BqArgs:
+                pass
+            args_obj = BqArgs()
+            args_obj.bigquery_api_endpoint = params["args"]["bigquery_api_endpoint"]
+            args_obj.project = params["args"]["project"]
+            args_obj.location = params["args"]["location"]
+            bq_client = create_bq_client(args_obj)
+
+            table_ref = bigquery.TableReference.from_api_repr(params["destination_table"])
+            query_results = json.loads(bq_client.list_rows(table_ref).to_dataframe().to_json())
+        response = convert_graph_data(query_results=query_results)
         self.do_data_response(response)
 
     def handle_post_node_expansion(self):