def _get_split_key(keys, num_splits):
  """Given a list of keys and a number of splits find the keys to split on.

  Args:
    keys: the list of keys.
    num_splits: the number of splits.

  Returns:
    A list of keys to split on.

  """

  # If the number of keys is less than the number of splits, we are limited
  # in the number of splits we can make.
  if not keys or (len(keys) < (num_splits - 1)):
    return keys

  # Calculate the number of keys per split. This should be KEYS_PER_SPLIT,
  # but may be less if there are not KEYS_PER_SPLIT * (numSplits - 1) scatter
  # entities.
  #
  # Consider the following dataset, where - represents an entity and
  # * represents an entity that is returned as a scatter entity:
  # ||---*-----*----*-----*-----*------*----*----||
  # If we want 4 splits in this data, the optimal split would look like:
  # ||---*-----*----*-----*-----*------*----*----||
  #            |          |            |
  # The scatter keys in the last region are not useful to us, so we never
  # request them:
  # ||---*-----*----*-----*-----*------*---------||
  #            |          |            |
  # With 6 scatter keys we want to set scatter points at indexes: 1, 3, 5.
  #
  # We keep this as a float so that any "fractional" keys per split get
  # distributed throughout the splits and don't make the last split
  # significantly larger than the rest.

  num_keys_per_split = max(1.0, float(len(keys)) / (num_splits - 1))

  split_keys = []

  # Grab the last sample for each split, otherwise the first split will be too
  # small.
  for i in range(1, num_splits):
    split_index = int(round(i * num_keys_per_split) - 1)
    split_keys.append(keys[split_index])

  return split_keys