Merge branch 'main' into shuowei-anywidget-sort-multi-column

# Conflicts:
#	bigframes/display/html.py

Author: shuoweil

bigframes/core/log_adapter.py

@@ -174,7 +174,8 @@ def wrapper(*args, **kwargs):
             full_method_name = f"{base_name.lower()}-{api_method_name}"
             # Track directly called methods
             if len(_call_stack) == 0:
-                add_api_method(full_method_name)
+                session = _find_session(*args, **kwargs)
+                add_api_method(full_method_name, session=session)
 
             _call_stack.append(full_method_name)
 
@@ -220,7 +221,8 @@ def wrapped(*args, **kwargs):
             full_property_name = f"{class_name.lower()}-{property_name.lower()}"
 
             if len(_call_stack) == 0:
-                add_api_method(full_property_name)
+                session = _find_session(*args, **kwargs)
+                add_api_method(full_property_name, session=session)
 
             _call_stack.append(full_property_name)
             try:
@@ -250,25 +252,41 @@ def wrapper(func):
     return wrapper
 
 
-def add_api_method(api_method_name):
+def add_api_method(api_method_name, session=None):
     global _lock
     global _api_methods
-    with _lock:
-        # Push the method to the front of the _api_methods list
-        _api_methods.insert(0, api_method_name.replace("<", "").replace(">", ""))
-        # Keep the list length within the maximum limit (adjust MAX_LABELS_COUNT as needed)
-        _api_methods = _api_methods[:MAX_LABELS_COUNT]
 
+    clean_method_name = api_method_name.replace("<", "").replace(">", "")
+
+    if session is not None and _is_session_initialized(session):
+        with session._api_methods_lock:
+            session._api_methods.insert(0, clean_method_name)
+            session._api_methods = session._api_methods[:MAX_LABELS_COUNT]
+    else:
+        with _lock:
+            # Push the method to the front of the _api_methods list
+            _api_methods.insert(0, clean_method_name)
+            # Keep the list length within the maximum limit (adjust MAX_LABELS_COUNT as needed)
+            _api_methods = _api_methods[:MAX_LABELS_COUNT]
 
-def get_and_reset_api_methods(dry_run: bool = False):
+
+def get_and_reset_api_methods(dry_run: bool = False, session=None):
     global _lock
+    methods = []
+
+    if session is not None and _is_session_initialized(session):
+        with session._api_methods_lock:
+            methods.extend(session._api_methods)
+            if not dry_run:
+                session._api_methods.clear()
+
     with _lock:
-        previous_api_methods = list(_api_methods)
+        methods.extend(_api_methods)
 
         # dry_run might not make a job resource, so only reset the log on real queries.
         if not dry_run:
             _api_methods.clear()
-    return previous_api_methods
+    return methods
 
 
 def _get_bq_client(*args, **kwargs):
@@ -283,3 +301,36 @@ def _get_bq_client(*args, **kwargs):
             return kwargv._block.session.bqclient
 
     return None
+
+
+def _is_session_initialized(session):
+    """Return True if fully initialized.
+
+    Because the method logger could get called before Session.__init__ has a
+    chance to run, we use the globals in that case.
+    """
+    return hasattr(session, "_api_methods_lock") and hasattr(session, "_api_methods")
+
+
+def _find_session(*args, **kwargs):
+    # This function cannot import Session at the top level because Session
+    # imports log_adapter.
+    from bigframes.session import Session
+
+    session = args[0] if args else None
+    if (
+        session is not None
+        and isinstance(session, Session)
+        and _is_session_initialized(session)
+    ):
+        return session
+
+    session = kwargs.get("session")
+    if (
+        session is not None
+        and isinstance(session, Session)
+        and _is_session_initialized(session)
+    ):
+        return session
+
+    return None

bigframes/display/html.py

@@ -48,60 +48,62 @@ def render_html(
     orderable_columns: list[str] | None = None,
 ) -> str:
     """Render a pandas DataFrame to HTML with specific styling."""
-    classes = "dataframe table table-striped table-hover"
-    table_html = [f'<table border="1" class="{classes}" id="{table_id}">']
-    precision = options.display.precision
     orderable_columns = orderable_columns or []
+    classes = "dataframe table table-striped table-hover"
+    table_html_parts = [f'<table border="1" class="{classes}" id="{table_id}">']
+    table_html_parts.append(_render_table_header(dataframe, orderable_columns))
+    table_html_parts.append(_render_table_body(dataframe))
+    table_html_parts.append("</table>")
+    return "".join(table_html_parts)
 
-    # Render table head
-    table_html.append("  <thead>")
-    table_html.append('    <tr style="text-align: left;">')
+
+def _render_table_header(dataframe: pd.DataFrame, orderable_columns: list[str]) -> str:
+    """Render the header of the HTML table."""
+    header_parts = ["  <thead>", "    <tr>"]
     for col in dataframe.columns:
         th_classes = []
         if col in orderable_columns:
             th_classes.append("sortable")
         class_str = f'class="{" ".join(th_classes)}"' if th_classes else ""
-        header_div = (
-            '<div style="resize: horizontal; overflow: auto; '
-            "box-sizing: border-box; width: 100%; height: 100%; "
-            'padding: 0.5em;">'
-            f"{html.escape(str(col))}"
-            "</div>"
-        )
-        table_html.append(
-            f'      <th style="text-align: left;" {class_str}>{header_div}</th>'
+        header_parts.append(
+            f'      <th {class_str}><div class="bf-header-content">'
+            f"{html.escape(str(col))}</div></th>"
         )
-    table_html.append("    </tr>")
-    table_html.append("  </thead>")
+    header_parts.extend(["    </tr>", "  </thead>"])
+    return "\n".join(header_parts)
+
+
+def _render_table_body(dataframe: pd.DataFrame) -> str:
+    """Render the body of the HTML table."""
+    body_parts = ["  <tbody>"]
+    precision = options.display.precision
 
-    # Render table body
-    table_html.append("  <tbody>")
     for i in range(len(dataframe)):
-        table_html.append("    <tr>")
+        body_parts.append("    <tr>")
         row = dataframe.iloc[i]
         for col_name, value in row.items():
             dtype = dataframe.dtypes.loc[col_name]  # type: ignore
             align = "right" if _is_dtype_numeric(dtype) else "left"
-            table_html.append(
-                '      <td style="text-align: {}; padding: 0.5em;">'.format(align)
-            )
 
             # TODO(b/438181139): Consider semi-exploding ARRAY/STRUCT columns
             # into multiple rows/columns like the BQ UI does.
             if pandas.api.types.is_scalar(value) and pd.isna(value):
-                table_html.append('        <em style="color: gray;">&lt;NA&gt;</em>')
+                body_parts.append(
+                    f'      <td class="cell-align-{align}">'
+                    '<em class="null-value">&lt;NA&gt;</em></td>'
+                )
             else:
                 if isinstance(value, float):
-                    formatted_value = f"{value:.{precision}f}"
-                    table_html.append(f"        {html.escape(formatted_value)}")
+                    cell_content = f"{value:.{precision}f}"
                 else:
-                    table_html.append(f"        {html.escape(str(value))}")
-            table_html.append("      </td>")
-        table_html.append("    </tr>")
-    table_html.append("  </tbody>")
-    table_html.append("</table>")
-
-    return "\n".join(table_html)
+                    cell_content = str(value)
+                body_parts.append(
+                    f'      <td class="cell-align-{align}">'
+                    f"{html.escape(cell_content)}</td>"
+                )
+        body_parts.append("    </tr>")
+    body_parts.append("  </tbody>")
+    return "\n".join(body_parts)
 
 
 def _obj_ref_rt_to_html(obj_ref_rt: str) -> str: