util.py 文件源码

def make_ragged_index(columns):
    """Make an index to hold data in a ragged array.

    Args:
        columns: A list of [N, _]-shaped numpy arrays of varying size, where
            N is the number of rows.

    Returns:
        A [len(columns)+1]-shaped array of begin,end positions of each column.
    """
    ragged_index = np.zeros([len(columns) + 1], dtype=np.int32)
    ragged_index[0] = 0
    for v, column in enumerate(columns):
        ragged_index[v + 1] = ragged_index[v] + column.shape[-1]
    ragged_index.flags.writeable = False
    return ragged_index