| 3244 | |
| 3245 | |
| 3246 | def test_apply_infer_columns(): |
| 3247 | df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [10, 20, 30, 40]}) |
| 3248 | ddf = dd.from_pandas(df, npartitions=2) |
| 3249 | |
| 3250 | def return_df(x): |
| 3251 | # will create new DataFrame which columns is ['sum', 'mean'] |
| 3252 | return pd.Series([x.sum(), x.mean()], index=["sum", "mean"]) |
| 3253 | |
| 3254 | # DataFrame to completely different DataFrame |
| 3255 | with warnings.catch_warnings(): |
| 3256 | warnings.simplefilter("ignore", UserWarning) |
| 3257 | result = ddf.apply(return_df, axis=1) |
| 3258 | assert isinstance(result, dd.DataFrame) |
| 3259 | tm.assert_index_equal(result.columns, pd.Index(["sum", "mean"])) |
| 3260 | assert_eq(result, df.apply(return_df, axis=1)) |
| 3261 | |
| 3262 | # DataFrame to Series |
| 3263 | with warnings.catch_warnings(): |
| 3264 | warnings.simplefilter("ignore", UserWarning) |
| 3265 | result = ddf.apply(lambda x: 1, axis=1) |
| 3266 | assert isinstance(result, dd.Series) |
| 3267 | assert result.name is None |
| 3268 | assert_eq(result, df.apply(lambda x: 1, axis=1)) |
| 3269 | |
| 3270 | def return_df2(x): |
| 3271 | return pd.Series([x * 2, x * 3], index=["x2", "x3"]) |
| 3272 | |
| 3273 | # Series to completely different DataFrame |
| 3274 | with warnings.catch_warnings(): |
| 3275 | warnings.simplefilter("ignore", UserWarning) |
| 3276 | result = ddf.x.apply(return_df2) |
| 3277 | assert isinstance(result, dd.DataFrame) |
| 3278 | tm.assert_index_equal(result.columns, pd.Index(["x2", "x3"])) |
| 3279 | assert_eq(result, df.x.apply(return_df2)) |
| 3280 | |
| 3281 | # Series to Series |
| 3282 | with warnings.catch_warnings(): |
| 3283 | warnings.simplefilter("ignore", UserWarning) |
| 3284 | result = ddf.x.apply(lambda x: 1) |
| 3285 | assert isinstance(result, dd.Series) |
| 3286 | assert result.name == "x" |
| 3287 | assert_eq(result, df.x.apply(lambda x: 1)) |
| 3288 | |
| 3289 | |
| 3290 | def test_index_time_properties(): |