diff --git a/python/xorbits/_mars/dataframe/base/pivot.py b/python/xorbits/_mars/dataframe/base/pivot.py index 0135e7ab1..eff428971 100644 --- a/python/xorbits/_mars/dataframe/base/pivot.py +++ b/python/xorbits/_mars/dataframe/base/pivot.py @@ -67,11 +67,14 @@ def execute_combine(cls, ctx: Union[dict, Context], op: "DataFramePivot"): input_data = ctx[op.inputs[0].key] out = op.outputs[0] - for dtype in op.output_columns: - if dtype not in input_data.dtypes.index: - input_data[dtype] = np.nan + dtypes = [ + dtype for dtype in op.output_columns if dtype not in input_data.dtypes.index + ] + nan_data = pd.DataFrame( + {col: [np.nan] for col in dtypes}, index=input_data.index + ) - ctx[out.key] = input_data + ctx[out.key] = pd.concat([input_data, nan_data], axis=1) @classmethod def execute(cls, ctx: Union[dict, Context], op: "DataFramePivot"):