Staged task-based shuffle implementation
| 367 | |
| 368 | |
| 369 | class TaskShuffle(SimpleShuffle): |
| 370 | """Staged task-based shuffle implementation""" |
| 371 | |
| 372 | @functools.cached_property |
| 373 | def _meta(self): |
| 374 | meta = self.frame._meta |
| 375 | if self.ignore_index: |
| 376 | meta = meta.reset_index(drop=True) |
| 377 | return meta |
| 378 | |
| 379 | def _layer(self): |
| 380 | max_branch = (self.options or {}).get("max_branch", None) or 32 |
| 381 | npartitions_input = self.frame.npartitions |
| 382 | if len(self._partitions) <= max_branch or npartitions_input <= max_branch: |
| 383 | # We are creating a small number of output partitions, |
| 384 | # or starting with a small number of input partitions. |
| 385 | # No need for staged shuffling. Staged shuffling will |
| 386 | # sometimes require extra work/communication in this case. |
| 387 | return super()._layer() |
| 388 | |
| 389 | # Calculate number of stages and splits per stage |
| 390 | npartitions = self.npartitions_out |
| 391 | stages = int(math.ceil(math.log(npartitions_input) / math.log(max_branch))) |
| 392 | if stages > 1: |
| 393 | nsplits = int(math.ceil(npartitions_input ** (1 / stages))) |
| 394 | else: |
| 395 | nsplits = npartitions_input |
| 396 | |
| 397 | # Construct global data-movement plan |
| 398 | inputs = [ |
| 399 | tuple(digit(i, j, nsplits) for j in range(stages)) |
| 400 | for i in range(nsplits**stages) |
| 401 | ] |
| 402 | inp_part_map = {inp: i for i, inp in enumerate(inputs)} |
| 403 | parts_out = range(len(inputs)) |
| 404 | |
| 405 | # Build graph |
| 406 | dsk = {} |
| 407 | name = self.frame._name |
| 408 | meta_input = make_meta(self.frame._meta) |
| 409 | for stage in range(stages): |
| 410 | # Define names |
| 411 | name_input = name |
| 412 | if stage == (stages - 1) and npartitions == npartitions_input: |
| 413 | name = self._name |
| 414 | parts_out = self._partitions |
| 415 | _filter = parts_out if self._filtered else None |
| 416 | else: |
| 417 | name = f"stage-{stage}-{self._name}" |
| 418 | _filter = None |
| 419 | |
| 420 | shuffle_group_name = f"group-{name}" |
| 421 | split_name = f"split-{name}" |
| 422 | |
| 423 | for global_part, part in enumerate(parts_out): |
| 424 | out = inputs[part] |
| 425 | |
| 426 | _concat_list = [] # get_item tasks to concat for this output partition |