Question

我编写了一个函数，该函数接受文件路径列表，然后将它们连接成一个大数据框。我想包含一个参数，该参数采用用户感兴趣的列名列表。

如果用户决定过滤列，则数据框必须始终包含“类别”列，但我希望默认值为返回所有列。我似乎不太清楚如何从数据框中选择列。

这是我的函数，中间插有一些伪代码，以解释我在说什么。

def combine_all_data(data_files, columns_needed=ALL):
    dataframes = map(pd.read_csv, data_files)

    if columns_needed != ALL
        columns_needed = ['category'] + columns_needed

    df = pd.concat(dataframes, sort=False)[columns_needed]
    return df

Answer 1

如果它是ALL，您不知道如何实现，则可以尝试以下操作：

def combine_all_data(data_files, columns_needed=None):
    kwargs= dict()
    if columns_needed is not None:
        if 'category' not in columns_needed:
            columns_needed= ['category'] + columns_needed
        kwargs['usecols']= columns_needed
    dataframes = [pd.read_csv(data_file, **kwargs) for data_file in data_files]
    return pd.concat(dataframes, sort=False)

这样做的好处是，您需要较少的内存，因为您不想看到的列在读取过程中已被跳过。另外，您返回一个完整的数据帧，而不是一个数据片段。因此，您可以不受限制地使用它。

Answer 2

read_csv有一个usecols参数：

def combine_all_data(data_files, columns_needed='ALL'):
   if needed_columns != 'ALL':
      if not 'category' in columns_needed:
         columns_needed.append('category')
      return pd.concat([pd.read_csv(x, usecols=columns_needed) for x 
      in data_files], sort=False)
   else: 
      return pd.concat([pd.read_csv(x) for x in data_files], sort=False)

（可选）从Pandas数据框中选择列

2 个答案: