| 1352 | |
| 1353 | |
| 1354 | class RenameFrame(Elemwise): |
| 1355 | _parameters = ["frame", "columns"] |
| 1356 | |
| 1357 | @functools.cached_property |
| 1358 | def unique_partition_mapping_columns_from_shuffle(self): |
| 1359 | result = set() |
| 1360 | columns = self.operand("columns") |
| 1361 | for elem in self.frame.unique_partition_mapping_columns_from_shuffle: |
| 1362 | if isinstance(elem, tuple): |
| 1363 | subset = self.frame._meta[list(elem)].rename(columns=columns) |
| 1364 | result.add(tuple(subset.columns)) |
| 1365 | else: |
| 1366 | # scalar |
| 1367 | subset = self.frame._meta[[elem]] |
| 1368 | result.add(subset.columns[0]) |
| 1369 | return result |
| 1370 | |
| 1371 | @staticmethod |
| 1372 | def operation(df, columns): |
| 1373 | return df.rename(columns=columns) |
| 1374 | |
| 1375 | def _simplify_up(self, parent, dependents): |
| 1376 | if isinstance(parent, Projection) and isinstance( |
| 1377 | self.operand("columns"), Mapping |
| 1378 | ): |
| 1379 | reverse_mapping = {val: key for key, val in self.operand("columns").items()} |
| 1380 | |
| 1381 | columns = determine_column_projection(self, parent, dependents) |
| 1382 | columns = _convert_to_list(columns) |
| 1383 | frame_columns = set(self.frame.columns) |
| 1384 | columns = [ |
| 1385 | ( |
| 1386 | reverse_mapping[col] |
| 1387 | if col in reverse_mapping and reverse_mapping[col] in frame_columns |
| 1388 | else col |
| 1389 | ) |
| 1390 | for col in columns |
| 1391 | ] |
| 1392 | columns = [col for col in self.frame.columns if col in columns] |
| 1393 | if columns == self.frame.columns: |
| 1394 | return |
| 1395 | |
| 1396 | return type(parent)( |
| 1397 | type(self)(self.frame[columns], *self.operands[1:]), |
| 1398 | *parent.operands[1:], |
| 1399 | ) |
| 1400 | |
| 1401 | |
| 1402 | class ColumnsSetter(RenameFrame): |