feat(display): improve display of nested data types

Author: shuoweil

bigframes/display/html.py

@@ -30,12 +30,13 @@
 
 def _flatten_nested_data(
     dataframe: pd.DataFrame,
-) -> Tuple[pd.DataFrame, Dict[str, List[int]]]:
+) -> Tuple[pd.DataFrame, Dict[str, List[int]], List[str]]:
     """Flatten nested STRUCT and ARRAY columns for display."""
     if dataframe.empty:
-        return dataframe.copy(), {}
+        return dataframe.copy(), {}, []
 
     result_df = dataframe.copy()
+    initial_columns = list(result_df.columns)
 
     # Attempt to parse JSON-like strings into structs or arrays
     for col_name in result_df.columns:
@@ -60,22 +61,25 @@ def _flatten_nested_data(
     struct_columns: List[str] = []
     array_columns: List[str] = []
     array_of_struct_columns: List[str] = []
+    clear_on_continuation_cols: List[str] = []
 
     for col_name_raw, col_data in result_df.items():
         col_name = str(col_name_raw)
-        # Fix: Use isinstance for proper type narrowing
         dtype = col_data.dtype
         if isinstance(dtype, pd.ArrowDtype):
             pa_type = dtype.pyarrow_dtype
             if pa.types.is_struct(pa_type):
                 struct_columns.append(col_name)
             elif pa.types.is_list(pa_type):
                 array_columns.append(col_name)
-                # Check if it's an ARRAY of STRUCT
                 if hasattr(pa_type, "value_type") and pa.types.is_struct(
                     pa_type.value_type
                 ):
                     array_of_struct_columns.append(col_name)
+            else:
+                clear_on_continuation_cols.append(col_name)
+        elif col_name in initial_columns:
+            clear_on_continuation_cols.append(col_name)
 
     # Handle ARRAY of STRUCT columns first
     for col_name in array_of_struct_columns:
@@ -123,6 +127,7 @@ def _flatten_nested_data(
         for field_idx in range(pa_type.num_fields):
             field = pa_type.field(field_idx)
             new_col_name = f"{col_name}.{field.name}"
+            clear_on_continuation_cols.append(new_col_name)
 
             regular_field_values: List[Any] = []
             for val in col_data:
@@ -137,36 +142,39 @@ def _flatten_nested_data(
 
     # Now handle ARRAY columns (including the newly created ones from ARRAY of STRUCT)
     if not array_columns:
-        return result_df, array_row_groups
+        return result_df, array_row_groups, clear_on_continuation_cols
 
-    # Find the maximum length of all array columns
-    max_array_length = 0
-    for col_name in array_columns:
-        col_data = result_df[col_name]
-        for val in col_data:
-            if val is not None and not (
-                isinstance(val, list) and len(val) == 1 and pd.isna(val[0])
-            ):
-                max_array_length = max(max_array_length, len(val))
-
-    # Create exploded rows
+    # Find the maximum length of array columns within each row
     exploded_rows = []
     for orig_idx, row in result_df.iterrows():
-        # Get array values for this row
         array_values = {}
+        max_len_in_row = 0
+        non_na_array_found = False
+
         for col_name in array_columns:
             val = row[col_name]
-            if val is None or (
+            if val is not None and not (
                 isinstance(val, list) and len(val) == 1 and pd.isna(val[0])
             ):
-                array_values[col_name] = [pd.NA] * max_array_length
+                array_values[col_name] = list(val)
+                max_len_in_row = max(max_len_in_row, len(val))
+                non_na_array_found = True
             else:
-                array_len = len(val)
-                padded_val = list(val) + [pd.NA] * (max_array_length - array_len)
-                array_values[col_name] = padded_val
+                array_values[col_name] = []
+
+        if not non_na_array_found:
+            new_row = row.copy()
+            for col_name in array_columns:
+                new_row[f"{col_name}"] = pd.NA
+            exploded_rows.append(new_row)
+            orig_key = str(orig_idx)
+            if orig_key not in array_row_groups:
+                array_row_groups[orig_key] = []
+            array_row_groups[orig_key].append(len(exploded_rows) - 1)
+            continue
 
-        # Create one row per array element
-        for array_idx in range(max_array_length):
+        # Create one row per array element, up to max_len_in_row
+        for array_idx in range(max_len_in_row):
             new_row = row.copy()
 
             # Remove array columns from the row copy
@@ -175,7 +183,10 @@ def _flatten_nested_data(
 
             # Add the specific array element for this index
             for col_name in array_columns:
-                new_row[f"{col_name}"] = array_values[col_name][array_idx]
+                if array_idx < len(array_values.get(col_name, [])):
+                    new_row[f"{col_name}"] = array_values[col_name][array_idx]
+                else:
+                    new_row[f"{col_name}"] = pd.NA
 
             exploded_rows.append(new_row)
 
@@ -187,9 +198,9 @@ def _flatten_nested_data(
 
     if exploded_rows:
         exploded_df = pd.DataFrame(exploded_rows)
-        return exploded_df, array_row_groups
+        return exploded_df, array_row_groups, clear_on_continuation_cols
     else:
-        return result_df, array_row_groups
+        return result_df, array_row_groups, clear_on_continuation_cols
 
 
 def _is_dtype_numeric(dtype) -> bool:
@@ -204,7 +215,11 @@ def render_html(
 ) -> str:
     """Render a pandas DataFrame to HTML with specific styling and nested data support."""
     # Flatten nested data first
-    flattened_df, array_row_groups = _flatten_nested_data(dataframe)
+    (
+        flattened_df,
+        array_row_groups,
+        clear_on_continuation,
+    ) = _flatten_nested_data(dataframe)
 
     classes = "dataframe table table-striped table-hover"
     table_html = [f'<table border="1" class="{classes}" id="{table_id}">']
@@ -226,10 +241,12 @@ def render_html(
         # Determine if this is an array continuation row
         row_class = ""
         orig_row_idx = None
+        is_continuation = False
         for orig_key, row_indices in array_row_groups.items():
             if i in row_indices and row_indices[0] != i:
                 row_class = "array-continuation"
                 orig_row_idx = orig_key
+                is_continuation = True
                 break
 
         if row_class:
@@ -241,14 +258,18 @@ def render_html(
 
         row = flattened_df.iloc[i]
         for col_name, value in row.items():
+            col_name_str = str(col_name)
+            if is_continuation and col_name_str in clear_on_continuation:
+                table_html.append('      <td style="padding: 0.5em;"></td>')
+                continue
             dtype = flattened_df.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)
             )
 
             if pandas.api.types.is_scalar(value) and pd.isna(value):
-                table_html.append('        <em style="color: gray;">&lt;NA&gt;</em>')
+                table_html.append("        ")
             else:
                 if isinstance(value, float):
                     formatted_value = f"{value:.{precision}f}"