def yield_batch(iterable, batch_size, num_tensors=1):
  """Generator that yields batches of a DataFrame iterator.

  Args:
    :iterable: Spark partition iterator.
    :batch_size: number of items to retrieve per invocation.
    :num_tensors: number of tensors (columns) expected in each item.

  Returns:
    An array of ``num_tensors`` arrays, each of length `batch_size`
  """
  tensors = [[] for i in range(num_tensors)]
  for item in iterable:
    if item is None:
      break
    for i in range(num_tensors):
      tmp = str(item[i]) if type(item[i]) is bytearray else item[i]
      tensors[i].append(tmp)
    if len(tensors[0]) >= batch_size:
      yield tensors
      tensors = [[] for i in range(num_tensors)]
  if len(tensors[0]) > 0:
      yield tensors