code: def get_sources_from_dataset(p, dataset, mode):
    
  
    import apache_beam as beam
    import csv
    from google.datalab.ml import CsvDataSet, BigQueryDataSet
  
    check_dataset(dataset, mode)
    if type(dataset) is CsvDataSet:
      source_list = []
      for ii, input_path in enumerate(dataset.files):
        source_list.append(p | 'Read from Csv %d (%s)' % (ii, mode) >>
                           beam.io.ReadFromText(input_path, strip_trailing_newlines=True))
      return (source_list |
              'Flatten Sources (%s)' % mode >>
              beam.Flatten() |
              'Create Dict from Csv (%s)' % mode >>
              beam.Map(lambda line: csv.DictReader([line], fieldnames=['image_url',
                                                                       'label']).next()))
    elif type(dataset) is BigQueryDataSet:
      bq_source = (beam.io.BigQuerySource(table=dataset.table) if dataset.table is not None else
                   beam.io.BigQuerySource(query=dataset.query))
      return p | 'Read source from BigQuery (%s)' % mode >> beam.io.Read(bq_source)
    else:
      raise ValueError('Invalid DataSet. Expect CsvDataSet or BigQueryDataSet')
docs: get pcollection from dataset.