def METHOD_NAME(self, data_source):
    """Returns a dict mapping filenames to the number of records in that file.
    Args:
      data_source: A DataSource describing the data to be written.
    Returns:
      A dict mapping filename to number of records in that file.
    """
    filenames_to_num_records = {}
    partitions = data_source.partitions
    record_per_partition = data_source.num_records // partitions
    remaining = data_source.num_records % partitions
    partition_length = len(str(partitions))
    for file_index in range(partitions):
        # Add leading zero to index_name e.g. 0001
        index_name = str(file_index).zfill(partition_length)
        filename = f"part-{index_name}.{self.extension}"
        num_records = record_per_partition
        if remaining and file_index < remaining:
            num_records = num_records + 1
        filenames_to_num_records[filename] = num_records
    return filenames_to_num_records