python - How to iterate over a CSV file and create an object from the first element and the rest using Pandas?

Question

I have the following code:

   def _load_data_set(self, dataset_file):
        data_frame = pandas.read_csv(dataset_file)
        table = data_frame.values.tolist()

        rows = len(table)
        patients = []
        for i in range(rows):
            first = table[i][0]
            rest = table[i]
            rest.pop(0)
            p = PatientDataSet(first, rest)
            patients.append(p)
        return patients

This code is basically, iterates over a CSV file (with a header) and for each row it splits the first place and the rest and creates a list of PatientDataSet object.

The input: CSV file with header. The output: List of PatientDataSet objects.

Although it works, I really don't like how I implemented it because I pop the first column and the code looks really ugly. Is it possible to suggest how would be better to do it?

question from:https://stackoverflow.com/questions/65865472/how-to-iterate-over-a-csv-file-and-create-an-object-from-the-first-element-and-t

Answer 1

I'd do it like this:

def _load_data_set(self, dataset_file):
    df = pandas.read_csv(dataset_file, index_col=0)
    return [PatientDataSet(first, rest.tolist()) for first, rest in df.iterrows()]

Passing index_col=0 to read_csv() naturally separates the first column from the rest.

iterrows() gives you the index and the values for each row, which is exactly what you need.

You may be able to remove .tolist() if your PatientDataSet does not actually need a list but can accept rest directly as a Pandas Series. That would be better, to avoid an unnecessary conversion.