//  [python]
def build_split_centralized_dataset(clientdata,
                                    preprocess_fn,
                                    to_embedding_id,
                                    sample_client_ids,
                                    split_by='date'):
  """Creates centralized dataset for both training and testing.

  The following functions are created to flat_map to avoid loading all
  dataset tensor into the memory.

  Args:
    clientdata: tff client datasets object.
    preprocess_fn: A data preprocessing function.
    to_embedding_id: An encoding function maps client id string to id number.
    sample_client_ids: Client ids for validation.
    split_by: setup of client data split, choose 'data' or 'random'.

  Returns:
    A training set for fine-tuning, a split training set with smaller size, a
    test set.
  """

  # The number of split for train/test sets
  # Default to 50%/50%, that is the test size is the last 50%
  # We further have a smaller split, which is the first 1/num_small_split.
  num_split = 2
  num_small_split = 4

  # Splits by date
  def _create_dataset_with_id_train(client_id):
    # add_in must be created in the _create_dataset function.
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = sort_by_date_pipe(
        clientdata._create_dataset(client_id)).map(add_id)
    total_size = client_ds.reduce(0, lambda x, _: x + 1)
    num_elements_half = tf.cast((total_size - 1) / num_split, dtype=tf.int64)
    return client_ds.take(num_elements_half)

  def _create_dataset_with_id_train_sp(client_id):
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = sort_by_date_pipe(
        clientdata._create_dataset(client_id)).map(add_id)
    total_size = client_ds.reduce(0, lambda x, _: x + 1)
    num_elements_ten_percent = tf.cast(
        (total_size - 1) / num_small_split, dtype=tf.int64)
    return client_ds.take(num_elements_ten_percent)

  def _create_dataset_with_id_test(client_id):
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = sort_by_date_pipe(
        clientdata._create_dataset(client_id)).map(add_id)
    total_size = client_ds.reduce(0, lambda x, _: x + 1)
    num_elements_half = tf.cast((total_size - 1) / num_split, dtype=tf.int64)
    return client_ds.skip(num_elements_half)

  # split by random
  def _random_split_create_dataset_with_id_train(client_id):
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = clientdata._create_dataset(client_id)
    client_ds = random_split_pipe(client_ds, num_split, 0).map(add_id)
    return client_ds

  def _random_split_create_dataset_with_id_train_sp(client_id):
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = clientdata._create_dataset(client_id)
    client_ds = random_split_pipe(
        random_split_pipe(client_ds, num_split, 0), num_small_split,
        0).map(add_id)
    return client_ds

  def _random_split_create_dataset_with_id_test(client_id):
    def add_id(x):
      x['client_id'] = to_embedding_id(client_id)
      return x

    # pylint: disable=protected-access
    client_ds = clientdata._create_dataset(client_id)
    client_ds = random_split_pipe(client_ds, num_split, 1).map(add_id)
    return client_ds

  nested_dataset = tf.data.Dataset.from_tensor_slices(sample_client_ids)

  if split_by == 'date':
    train_fn = _create_dataset_with_id_train
    split_fn = _create_dataset_with_id_train_sp
    test_fn = _create_dataset_with_id_test

  elif split_by == 'random':
    train_fn = _random_split_create_dataset_with_id_train
    split_fn = _random_split_create_dataset_with_id_train_sp
    test_fn = _random_split_create_dataset_with_id_test

  per_train = nested_dataset.flat_map(train_fn)
  per_train_split = nested_dataset.flat_map(split_fn)
  per_test = nested_dataset.flat_map(test_fn)

  per_train = preprocess_fn(per_train).cache()
  per_train_split = preprocess_fn(per_train_split).cache()
  per_test = preprocess_fn(per_test).cache()

  return per_train, per_train_split, per_test