def partition_by_rows()

in smallpond/io/arrow.py [0:0]

• 21 lines of code
• 6 McCabe index (conditional complexity)

    def partition_by_rows(row_ranges: List["RowRange"], npartition: int) -> List[List["RowRange"]]:
        """Evenly split a list of row ranges into `npartition` partitions."""
        # NOTE: `row_ranges` should not be modified by this function
        row_ranges = copy.deepcopy(row_ranges)
        num_rows: int = sum(row_range.num_rows for row_range in row_ranges)
        num_partitions: int = npartition
        row_range_partitions: List[List[RowRange]] = []
        while num_partitions:
            rows_in_partition = (num_rows + num_partitions - 1) // num_partitions
            num_rows -= rows_in_partition
            num_partitions -= 1
            row_ranges_in_partition = []
            while rows_in_partition:
                current_range = row_ranges[0]
                if current_range.num_rows == 0:
                    row_ranges.pop(0)
                    continue
                taken_range = current_range.take(rows_in_partition)
                row_ranges_in_partition.append(taken_range)
                rows_in_partition -= taken_range.num_rows
            row_range_partitions.append(row_ranges_in_partition)
        assert num_rows == 0 and num_partitions == 0
        return row_range_partitions