Fix: expand dict transformation

marimo/_plugins/ui/_impl/dataframes/transforms/handlers.py

@@ -497,7 +497,81 @@ def handle_explode_columns(
     def handle_expand_dict(
         df: DataFrame, transform: ExpandDictTransform
     ) -> DataFrame:
-        return df.explode(transform.column_id)
+        collected_df, undo = collect_and_preserve_type(df)
+        native_df = collected_df.to_native()
+
+        # Keep pandas handling fully pandas-native so mixed/object columns in
+        # unrelated fields do not trigger Arrow coercion errors.
+        if nw.dependencies.is_pandas_dataframe(native_df):
+            import math
+
+            import pandas as pd
+
+            result_df = native_df.copy()
+
+            def normalise_empty_dict(value: Any) -> Any:
+                if value is None:
+                    return {}
+                if isinstance(value, float) and math.isnan(value):
+                    return {}
+                return value
+
+            # Keep expansion shallow and replace top-level null/nan entries so
+            # pandas and other backends agree on expand-dict behaviour.
+            expanded = pd.json_normalize(
+                result_df.pop(transform.column_id).map(normalise_empty_dict),
+                # type: ignore[arg-type]
+                max_level=0,
+            )
+            duplicate_columns = sorted(
+                set(result_df.columns) & set(expanded.columns)
+            )
+            if duplicate_columns:
+                raise nw.exceptions.InvalidOperationError(
+                    "Cannot expand dict because it would duplicate existing "
+                    f"columns: {duplicate_columns}"
+                )
+            expanded.index = result_df.index
+            return undo(nw.from_native(result_df.join(expanded)))
+
+        schema = collected_df.collect_schema()
+        dtype = schema.get(transform.column_id)
+        fields = getattr(dtype, "fields", None)
+        if fields is not None:
+            field_names = [field.name for field in fields]
+            columns = schema.names()
+            column_index = columns.index(transform.column_id)
+            expanded_columns = (
+                columns[:column_index]
+                + field_names
+                + columns[column_index + 1 :]
+            )
+            duplicate_columns = sorted(
+                (set(columns) - {transform.column_id}) & set(field_names)
+            )
+            if duplicate_columns:
+                raise nw.exceptions.InvalidOperationError(
+                    "Cannot expand dict because it would duplicate existing "
+                    f"columns: {duplicate_columns}"
+                )
+            return undo(
+                collected_df.with_columns(
+                    [
+                        nw.col(transform.column_id)
+                        .struct.field(field_name)
+                        .alias(field_name)
+                        for field_name in field_names
+                    ]
+                )
+                .drop(transform.column_id)
+                .select(expanded_columns)
+            )
+
+        raise nw.exceptions.InvalidOperationError(
+            "Expand dict requires a struct-like column with named fields on "
+            f"this backend; got column '{transform.column_id}' with dtype "
+            f"{dtype!r}."
+        )
 
     @staticmethod
     def handle_unique(df: DataFrame, transform: UniqueTransform) -> DataFrame:

marimo/_plugins/ui/_impl/dataframes/transforms/print_code.py

@@ -222,8 +222,11 @@ def generate_where_clause(df_name: str, where: FilterCondition) -> str:
 
     elif transform.type == TransformType.EXPAND_DICT:
         column_id = _as_literal(transform.column_id)
-        args = f"{df_name}.pop({column_id}).values.tolist()"
-        return f"{df_name}.join(pd.DataFrame({args}))"
+        return (
+            f"{df_name}.join("
+            f"pd.json_normalize({df_name}.pop({column_id}).map(lambda value: {{}} if value is None or (isinstance(value, float) and value != value) else value), max_level=0).set_axis({df_name}.index, axis=0)"
+            f")"
+        )
 
     elif transform.type == TransformType.UNIQUE:
         column_ids = transform.column_ids
@@ -465,7 +468,7 @@ def generate_where_clause_polars(where: FilterCondition) -> str:
 
     elif transform.type == TransformType.EXPAND_DICT:
         column_id = _as_literal(transform.column_id)
-        return f"{df_name}.hstack(pl.DataFrame({df_name}.select({column_id}).to_series().to_list())).drop({column_id})"
+        return f"{df_name}.unnest({column_id})"
 
     elif transform.type == TransformType.UNIQUE:
         column_ids = transform.column_ids

tests/_plugins/ui/_impl/dataframes/test_handlers.py

@@ -47,8 +47,8 @@
 
 pytest.importorskip("ibis")
 pd = pytest.importorskip("pandas")
-pytest.importorskip("polars")
 pytest.importorskip("pyarrow")
+pytest.importorskip("polars")
 
 
 def apply(df: DataFrameType, transform: Transform) -> DataFrameType:
@@ -86,7 +86,10 @@ def assert_frame_equal(a: DataFrameType, b: DataFrameType) -> None:
 
 
 def assert_frame_equal_with_nans(
-    a: DataFrameType, b: DataFrameType, allow_nan_equals_zero: bool = False
+    a: DataFrameType,
+    b: DataFrameType,
+    allow_nan_equals_zero: bool = False,
+    allow_none_equals_nan: bool = False,
 ) -> None:
     """
     Assert two dataframes are equal, treating NaNs in the same locations as equal.
@@ -97,6 +100,9 @@ def assert_frame_equal_with_nans(
         allow_nan_equals_zero: If True, treat NaN and 0.0 as equivalent values.
             This is useful for pivot operations where missing aggregations may
             be filled with 0.0 or NaN depending on the backend.
+        allow_none_equals_nan: If True, treat None and NaN as equivalent
+            missing values. This is useful when different backends materialise
+            missing numeric values differently.
     """
     import math
 
@@ -137,7 +143,25 @@ def assert_frame_equal_with_nans(
                     or val_b == 0.0
                 )
             )
-            if not (val_a == val_b or both_nan or nan_or_zero_match):
+            # Useful for expand dict operations where None and nan are equal
+            none_nan_match = allow_none_equals_nan and (
+                (
+                    val_a is None
+                    and isinstance(val_b, float)
+                    and math.isnan(val_b)
+                )
+                or (
+                    val_b is None
+                    and isinstance(val_a, float)
+                    and math.isnan(val_a)
+                )
+            )
+            if not (
+                val_a == val_b
+                or both_nan
+                or nan_or_zero_match
+                or none_nan_match
+            ):
                 raise AssertionError(
                     f"DataFrame values differ at column '{col}', row {idx}: {val_a} != {val_b}"
                 )
@@ -1733,18 +1757,15 @@ def test_explode_columns(df: DataFrameType) -> None:
         assert nw_result.columns == ["A", "B", "C"]
 
     @staticmethod
-    @pytest.mark.skip(
-        reason="Dict/struct expansion not supported uniformly across backends"
-    )
     @pytest.mark.parametrize(
         ("df", "expected"),
         list(
             zip(
                 create_test_dataframes(
-                    {"A": [{"foo": 1, "bar": "hello"}], "B": [1]}
+                    {"A": [{"foo": 1, "bar": "hello"}, None], "B": [1, 2]},
                 ),
                 create_test_dataframes(
-                    {"B": [1], "foo": [1], "bar": ["hello"]}
+                    {"B": [1, 2], "foo": [1, None], "bar": ["hello", None]},
                 ),
                 strict=False,
             )
@@ -1760,11 +1781,82 @@ def test_expand_dict(df: DataFrameType, expected: DataFrameType) -> None:
         nw_expected = collect_df(expected)
         result_cols = sorted(nw_result.columns)
         expected_cols = sorted(nw_expected.columns)
-        assert_frame_equal(
+        assert_frame_equal_with_nans(
             nw_expected.select(expected_cols),
             nw_result.select(result_cols),
+            allow_none_equals_nan=True,
         )
 
+    @staticmethod
+    @pytest.mark.parametrize(
+        ("df", "expected"),
+        list(
+            zip(
+                create_test_dataframes(
+                    {
+                        "A": [
+                            {"foo": 1, "nested": {"x": 2}},
+                            None,
+                        ],
+                        "B": [1, 2],
+                    },
+                    include=["pandas", "polars"],
+                ),
+                create_test_dataframes(
+                    {
+                        "B": [1, 2],
+                        "foo": [1, None],
+                        "nested": [{"x": 2}, None],
+                    },
+                    include=["pandas", "polars"],
+                ),
+                strict=False,
+            )
+        ),
+    )
+    def test_expand_dict_nested_dicts(
+        df: DataFrameType, expected: DataFrameType
+    ) -> None:
+        transform = ExpandDictTransform(
+            type=TransformType.EXPAND_DICT, column_id="A"
+        )
+        result = apply(df, transform)
+        nw_result = collect_df(result)
+        nw_expected = collect_df(expected)
+        result_cols = sorted(nw_result.columns)
+        expected_cols = sorted(nw_expected.columns)
+        assert_frame_equal_with_nans(
+            nw_expected.select(expected_cols),
+            nw_result.select(result_cols),
+            allow_none_equals_nan=True,
+        )
+
+    @pytest.mark.xfail(
+        reason=(
+            "Duplicate-column errors from expand-dict are not yet normalised "
+            "across all optional backends."
+        )
+    )
+    @staticmethod
+    @pytest.mark.parametrize(
+        "df",
+        create_test_dataframes(
+            {
+                "A": [{"B": 1}, {"B": 2}],
+                "B": [10, 20],
+            },
+            strict=False,
+        ),
+    )
+    def test_expand_dict_duplicate_columns_raises_xfail(
+        df: DataFrameType,
+    ) -> None:
+        transform = ExpandDictTransform(
+            type=TransformType.EXPAND_DICT, column_id="A"
+        )
+        with pytest.raises(nw.exceptions.InvalidOperationError):
+            apply(df, transform)
+
     @staticmethod
     @pytest.mark.parametrize(
         (