refactor: resue pandas struct.explode()

Author: shuoweil

bigframes/display/_flatten.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, cast
+from typing import cast
 
 import pandas as pd
 import pyarrow as pa
@@ -39,54 +39,12 @@ def flatten_nested_data(
         nested_originated_columns,
     ) = _classify_columns(result_df)
 
-    # Flatten ARRAY of STRUCT columns
-    def update_array_columns(col_name: str, new_col_names: list[str]) -> None:
-        array_columns.remove(col_name)
-        array_columns.extend(new_col_names)
-
-    def create_list_series(
-        original_arr: pa.Array, field_arr: pa.Array, index: pd.Index, field: pa.Field
-    ) -> pd.Series:
-        new_list_array = pa.ListArray.from_arrays(
-            original_arr.offsets, field_arr, mask=original_arr.is_null()
-        )
-        return pd.Series(
-            new_list_array.to_pylist(),
-            dtype=pd.ArrowDtype(pa.list_(field.type)),
-            index=index,
-        )
-
-    result_df = _flatten_and_replace_columns(
-        result_df,
-        array_of_struct_columns,
-        nested_originated_columns,
-        get_struct_type=lambda t: t.value_type,
-        get_field_values=lambda arr: arr.values.flatten(),
-        create_series=create_list_series,
-        update_metadata=update_array_columns,
+    result_df, array_columns = _flatten_array_of_struct_columns(
+        result_df, array_of_struct_columns, array_columns, nested_originated_columns
     )
 
-    # Flatten regular STRUCT columns
-    def update_clear_on_continuation(col_name: str, new_col_names: list[str]) -> None:
-        clear_on_continuation_cols.extend(new_col_names)
-
-    def create_struct_series(
-        original_arr: pa.Array, field_arr: pa.Array, index: pd.Index, field: pa.Field
-    ) -> pd.Series:
-        return pd.Series(
-            field_arr.to_pylist(),
-            dtype=pd.ArrowDtype(field.type),
-            index=index,
-        )
-
-    result_df = _flatten_and_replace_columns(
-        result_df,
-        struct_columns,
-        nested_originated_columns,
-        get_struct_type=lambda t: t,
-        get_field_values=lambda arr: arr.flatten(),
-        create_series=create_struct_series,
-        update_metadata=update_clear_on_continuation,
+    result_df, clear_on_continuation_cols = _flatten_struct_columns(
+        result_df, struct_columns, clear_on_continuation_cols, nested_originated_columns
     )
 
     # Now handle ARRAY columns (including the newly created ones from ARRAY of STRUCT)
@@ -146,36 +104,45 @@ def _classify_columns(
     )
 
 
-def _flatten_and_replace_columns(
+def _flatten_array_of_struct_columns(
     dataframe: pd.DataFrame,
-    columns: list[str],
+    array_of_struct_columns: list[str],
+    array_columns: list[str],
     nested_originated_columns: set[str],
-    get_struct_type: Callable[[pa.DataType], pa.DataType],
-    get_field_values: Callable[[pa.Array], list[pa.Array]],
-    create_series: Callable[[pa.Array, pa.Array, pd.Index, pa.Field], pd.Series],
-    update_metadata: Callable[[str, list[str]], None],
-) -> pd.DataFrame:
-    """Generic helper to flatten structure-like columns and replace them in the DataFrame."""
+) -> tuple[pd.DataFrame, list[str]]:
+    """Flatten ARRAY of STRUCT columns into separate array columns for each field."""
     result_df = dataframe.copy()
-    for col_name in columns:
+    for col_name in array_of_struct_columns:
         col_data = result_df[col_name]
         pa_type = cast(pd.ArrowDtype, col_data.dtype).pyarrow_dtype
-        struct_type = get_struct_type(pa_type)
+        struct_type = pa_type.value_type
 
+        # Use PyArrow to reshape the list<struct> into multiple list<field> arrays
         arrow_array = pa.array(col_data)
-        flattened_fields = get_field_values(arrow_array)
+        offsets = arrow_array.offsets
+        values = arrow_array.values  # StructArray
+        flattened_fields = values.flatten()  # List[Array]
 
         new_cols_to_add = {}
-        new_col_names = []
+        new_array_col_names = []
 
+        # Create new columns for each struct field
         for field_idx in range(struct_type.num_fields):
             field = struct_type.field(field_idx)
             new_col_name = f"{col_name}.{field.name}"
             nested_originated_columns.add(new_col_name)
-            new_col_names.append(new_col_name)
+            new_array_col_names.append(new_col_name)
 
-            new_cols_to_add[new_col_name] = create_series(
-                arrow_array, flattened_fields[field_idx], result_df.index, field
+            # Reconstruct ListArray for this field
+            # Use mask=arrow_array.is_null() to preserve nulls from the original list
+            new_list_array = pa.ListArray.from_arrays(
+                offsets, flattened_fields[field_idx], mask=arrow_array.is_null()
+            )
+
+            new_cols_to_add[new_col_name] = pd.Series(
+                new_list_array.to_pylist(),
+                dtype=pd.ArrowDtype(pa.list_(field.type)),
+                index=result_df.index,
             )
 
         col_idx = result_df.columns.to_list().index(col_name)
@@ -190,9 +157,11 @@ def _flatten_and_replace_columns(
             axis=1,
         )
 
-        update_metadata(col_name, new_col_names)
-
-    return result_df
+        # Update array_columns list
+        array_columns.remove(col_name)
+        # Add the new array columns
+        array_columns.extend(new_array_col_names)
+    return result_df, array_columns
 
 
 def _explode_array_columns(
@@ -271,3 +240,39 @@ def _explode_array_columns(
         return exploded_df, array_row_groups
     else:
         return dataframe, array_row_groups
+
+
+def _flatten_struct_columns(
+    dataframe: pd.DataFrame,
+    struct_columns: list[str],
+    clear_on_continuation_cols: list[str],
+    nested_originated_columns: set[str],
+) -> tuple[pd.DataFrame, list[str]]:
+    """Flatten regular STRUCT columns using pandas accessor."""
+    result_df = dataframe.copy()
+    for col_name in struct_columns:
+        # Use pandas struct accessor to explode the struct column into a DataFrame of its fields
+        exploded_struct = result_df[col_name].struct.explode()
+
+        # Rename columns to 'parent.child' format
+        exploded_struct.columns = [
+            f"{col_name}.{sub_col}" for sub_col in exploded_struct.columns
+        ]
+
+        # Update metadata
+        for new_col in exploded_struct.columns:
+            nested_originated_columns.add(new_col)
+            clear_on_continuation_cols.append(new_col)
+
+        # Replace the original struct column with the new field columns
+        col_idx = result_df.columns.to_list().index(col_name)
+        result_df = pd.concat(
+            [
+                result_df.iloc[:, :col_idx],
+                exploded_struct,
+                result_df.iloc[:, col_idx + 1 :],
+            ],
+            axis=1,
+        )
+
+    return result_df, clear_on_continuation_cols