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