feat: use dataclass for flatten_nested_data

Author: shuoweil

bigframes/display/_flatten.py

@@ -16,18 +16,44 @@
 
 from __future__ import annotations
 
+import dataclasses
 from typing import cast
 
 import pandas as pd
 import pyarrow as pa
 
 
+@dataclasses.dataclass(frozen=True)
+class FlattenResult:
+    """The result of flattening a DataFrame."""
+
+    dataframe: pd.DataFrame
+    """The flattened DataFrame."""
+
+    row_groups: dict[str, list[int]]
+    """
+    A mapping from original row index to the new row indices that were created
+    from it.
+    """
+
+    cleared_on_continuation: list[str]
+    """A list of column names that should be cleared on continuation rows."""
+
+    nested_columns: set[str]
+    """A set of column names that were created from nested data."""
+
+
 def flatten_nested_data(
     dataframe: pd.DataFrame,
-) -> tuple[pd.DataFrame, dict[str, list[int]], list[str], set[str]]:
+) -> FlattenResult:
     """Flatten nested STRUCT and ARRAY columns for display."""
     if dataframe.empty:
-        return dataframe.copy(), {}, [], set()
+        return FlattenResult(
+            dataframe=dataframe.copy(),
+            row_groups={},
+            cleared_on_continuation=[],
+            nested_columns=set(),
+        )
 
     result_df = dataframe.copy()
 
@@ -49,19 +75,19 @@ def flatten_nested_data(
 
     # Now handle ARRAY columns (including the newly created ones from ARRAY of STRUCT)
     if not array_columns:
-        return (
-            result_df,
-            {},
-            clear_on_continuation_cols,
-            nested_originated_columns,
+        return FlattenResult(
+            dataframe=result_df,
+            row_groups={},
+            cleared_on_continuation=clear_on_continuation_cols,
+            nested_columns=nested_originated_columns,
         )
 
     result_df, array_row_groups = _explode_array_columns(result_df, array_columns)
-    return (
-        result_df,
-        array_row_groups,
-        clear_on_continuation_cols,
-        nested_originated_columns,
+    return FlattenResult(
+        dataframe=result_df,
+        row_groups=array_row_groups,
+        cleared_on_continuation=clear_on_continuation_cols,
+        nested_columns=nested_originated_columns,
     )
 
 

bigframes/display/html.py

@@ -49,23 +49,20 @@ 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,
-        clear_on_continuation,
-        nested_originated_columns,
-    ) = _flatten.flatten_nested_data(dataframe)
+    flatten_result = _flatten.flatten_nested_data(dataframe)
 
     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(flattened_df, orderable_columns))
+    table_html_parts.append(
+        _render_table_header(flatten_result.dataframe, orderable_columns)
+    )
     table_html_parts.append(
         _render_table_body(
-            flattened_df,
-            array_row_groups,
-            clear_on_continuation,
-            nested_originated_columns,
+            flatten_result.dataframe,
+            flatten_result.row_groups,
+            flatten_result.cleared_on_continuation,
+            flatten_result.nested_columns,
         )
     )
     table_html_parts.append("</table>")

tests/unit/display/test_html.py

@@ -165,7 +165,9 @@ def test_flatten_nested_data_flattens_structs():
         }
     )
 
-    flattened, _, _, nested_originated_columns = flatten_nested_data(struct_data)
+    result = flatten_nested_data(struct_data)
+    flattened = result.dataframe
+    nested_originated_columns = result.nested_columns
 
     assert "struct_col.name" in flattened.columns
     assert "struct_col.age" in flattened.columns
@@ -186,7 +188,10 @@ def test_flatten_nested_data_explodes_arrays():
         }
     )
 
-    flattened, groups, _, nested_originated_columns = flatten_nested_data(array_data)
+    result = flatten_nested_data(array_data)
+    flattened = result.dataframe
+    groups = result.row_groups
+    nested_originated_columns = result.nested_columns
 
     assert len(flattened) == 5  # 3 + 2 array elements
     assert "0" in groups  # First original row