wehub-resource-sync
29 lines
697 B
Python
29 lines
697 B
Python
from typing import Callable, Iterable, List
|
|
|
|
import numpy as np
|
|
|
|
from ray.experimental.raysort.types import BlockInfo
|
|
|
|
|
|
def get_boundaries(num_parts: int) -> List[int]:
|
|
return [0] * num_parts
|
|
|
|
|
|
def sort_and_partition(part: np.ndarray, boundaries: List[int]) -> List[BlockInfo]:
|
|
N = len(boundaries)
|
|
offset = 0
|
|
size = int(np.ceil(part.size / N))
|
|
blocks = []
|
|
for _ in range(N):
|
|
blocks.append((offset, size))
|
|
offset += size
|
|
return blocks
|
|
|
|
|
|
def merge_partitions(
|
|
num_blocks: int, get_block: Callable[[int, int], np.ndarray]
|
|
) -> Iterable[memoryview]:
|
|
blocks = [get_block(i, 0) for i in range(num_blocks)]
|
|
for block in blocks:
|
|
yield block
|