| 1213 | |
| 1214 | |
| 1215 | class _SetIndexPost(Blockwise): |
| 1216 | _parameters = [ |
| 1217 | "frame", |
| 1218 | "index_name", |
| 1219 | "drop", |
| 1220 | "set_name", |
| 1221 | "column_dtype", |
| 1222 | "computed_divisions", |
| 1223 | "user_divisions", |
| 1224 | ] |
| 1225 | _is_length_preserving = True |
| 1226 | |
| 1227 | @property |
| 1228 | def _args(self) -> list: |
| 1229 | return self.operands[:5] |
| 1230 | |
| 1231 | @staticmethod |
| 1232 | def operation(df, index_name, drop, set_name, column_dtype): |
| 1233 | df = df.set_index(set_name, drop=drop) |
| 1234 | df.index.name = index_name |
| 1235 | df.columns = df.columns.astype(column_dtype) |
| 1236 | return df |
| 1237 | |
| 1238 | def _get_culled_divisions(self, divisions): |
| 1239 | if self.frame.npartitions < len(divisions) - 1: |
| 1240 | part_filter = list(self.frame.find_operations(PartitionsFiltered)) |
| 1241 | if len(part_filter) > 0: |
| 1242 | return tuple( |
| 1243 | [ |
| 1244 | div |
| 1245 | for i, div in enumerate(divisions) |
| 1246 | if i in part_filter[0]._partitions |
| 1247 | ] |
| 1248 | + [divisions[-1]] |
| 1249 | ) |
| 1250 | else: |
| 1251 | return self.frame.divisions |
| 1252 | |
| 1253 | return divisions |
| 1254 | |
| 1255 | def _divisions(self): |
| 1256 | if self.operand("user_divisions") is not None: |
| 1257 | return self._get_culled_divisions(self.operand("user_divisions")) |
| 1258 | assert self.computed_divisions is not None |
| 1259 | return self._get_culled_divisions(self.computed_divisions[0]) |
| 1260 | |
| 1261 | |
| 1262 | class SortIndexBlockwise(Blockwise): |