diff --git a/marimo/_plugins/ui/_impl/dataframes/transforms/handlers.py b/marimo/_plugins/ui/_impl/dataframes/transforms/handlers.py index 175fa251462..08b11d8ee20 100644 --- a/marimo/_plugins/ui/_impl/dataframes/transforms/handlers.py +++ b/marimo/_plugins/ui/_impl/dataframes/transforms/handlers.py @@ -497,7 +497,80 @@ 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) + if isinstance(dtype, nw.Struct): + fields = dtype.fields + 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: diff --git a/marimo/_plugins/ui/_impl/dataframes/transforms/print_code.py b/marimo/_plugins/ui/_impl/dataframes/transforms/print_code.py index d27d1d07698..c2461e56c37 100644 --- a/marimo/_plugins/ui/_impl/dataframes/transforms/print_code.py +++ b/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 diff --git a/tests/_plugins/ui/_impl/dataframes/test_handlers.py b/tests/_plugins/ui/_impl/dataframes/test_handlers.py index cf1b60d507e..c6a6b21bd53 100644 --- a/tests/_plugins/ui/_impl/dataframes/test_handlers.py +++ b/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,100 @@ 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, + ) + + @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( + 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( + "df", + create_test_dataframes( + {"A": [1, 2], "B": [10, 20]}, + include=[ + "polars", + "pyarrow", + "ibis", + ], + ), + ) + def test_expand_dict_unsupported_column_raises( + df: DataFrameType, + ) -> None: + transform = ExpandDictTransform( + type=TransformType.EXPAND_DICT, column_id="A" + ) + with pytest.raises( + nw.exceptions.InvalidOperationError, + match="Expand dict requires a struct-like column with named fields", + ): + apply(df, transform) + @staticmethod @pytest.mark.parametrize( ( diff --git a/tests/_plugins/ui/_impl/dataframes/test_print_code.py b/tests/_plugins/ui/_impl/dataframes/test_print_code.py index 57a297b0cb4..50b41eb7996 100644 --- a/tests/_plugins/ui/_impl/dataframes/test_print_code.py +++ b/tests/_plugins/ui/_impl/dataframes/test_print_code.py @@ -525,6 +525,96 @@ def test_print_code_result_matches_actual_transform_pandas( ) +@pytest.mark.skipif( + not DependencyManager.pandas.has(), reason="pandas not installed" +) +def test_print_code_expand_dict_nested_dict_pandas() -> None: + import pandas as pd + + transform = ExpandDictTransform( + type=TransformType.EXPAND_DICT, + column_id="dicts", + ) + transformations = Transformations([transform]) + my_df = pd.DataFrame( + { + "dicts": [{"a": 1, "nested": {"x": 2}}, None], + "other": [10, 20], + } + ) + + pandas_code = python_print_transforms( + "my_df", + list(my_df.columns), + transformations.transforms, + python_print_pandas, + ) + assert pandas_code + + loc = {"pd": pd, "my_df": my_df.copy()} + exec(pandas_code, {}, loc) + code_result = loc["my_df_next"] + + nw_df = nw.from_native(my_df.copy(), eager_only=True).lazy() + result_nw = _apply_transforms( + nw_df, + NarwhalsTransformHandler(), + transformations, + ) + real_result = result_nw.collect().to_native().reset_index(drop=True) + + pd.testing.assert_frame_equal( + code_result.reset_index(drop=True), + real_result, + ) + + assert list(code_result.columns) == ["other", "a", "nested"] + + +@pytest.mark.skipif( + not DependencyManager.pandas.has(), reason="pandas not installed" +) +def test_print_code_expand_dict_nan_pandas() -> None: + import pandas as pd + + transform = ExpandDictTransform( + type=TransformType.EXPAND_DICT, + column_id="dicts", + ) + transformations = Transformations([transform]) + my_df = pd.DataFrame( + { + "dicts": [{"a": 1}, float("nan"), {"a": 3}], + "other": [10, 20, 30], + } + ) + + pandas_code = python_print_transforms( + "my_df", + list(my_df.columns), + transformations.transforms, + python_print_pandas, + ) + assert pandas_code + + loc = {"pd": pd, "my_df": my_df.copy()} + exec(pandas_code, {}, loc) + code_result = loc["my_df_next"] + + nw_df = nw.from_native(my_df.copy(), eager_only=True).lazy() + result_nw = _apply_transforms( + nw_df, + NarwhalsTransformHandler(), + transformations, + ) + real_result = result_nw.collect().to_native().reset_index(drop=True) + + pd.testing.assert_frame_equal( + code_result.reset_index(drop=True), + real_result, + ) + + @given( transform=create_transform_strategy( defined_column_id, @@ -772,6 +862,50 @@ def test_print_code_result_matches_actual_transform_polars( pl_testing.assert_frame_equal(code_result, real_result) +@pytest.mark.skipif( + not DependencyManager.polars.has(), reason="polars not installed" +) +def test_print_code_expand_dict_nested_dict_polars() -> None: + import polars as pl + import polars.testing as pl_testing + + transform = ExpandDictTransform( + type=TransformType.EXPAND_DICT, + column_id="dicts", + ) + transformations = Transformations([transform]) + my_df = pl.DataFrame( + { + "dicts": [{"a": 1, "nested": {"x": 2}}, None], + "other": [10, 20], + } + ) + + polars_code = python_print_transforms( + "my_df", + my_df.columns, + transformations.transforms, + python_print_polars, + ) + assert polars_code + + loc = {"pl": pl, "my_df": my_df.clone()} + exec(polars_code, globals(), loc) + code_result = loc["my_df_next"] + + nw_df = nw.from_native(my_df.clone(), eager_only=True).lazy() + result_nw = _apply_transforms( + nw_df, + NarwhalsTransformHandler(), + transformations, + ) + real_result = result_nw.collect().to_native() + + pl_testing.assert_frame_equal(code_result, real_result) + + assert code_result.columns == ["a", "nested", "other"] + + @given( transform=create_transform_strategy( defined_column_id,