Python按指定列拆分CSV异常：选中列1却按列3分组

AI 小北

2026-8-9

问题

我编写了一个Python批处理脚本，流程为导入CSV数据、清洗数据、生成报表，最后按指定类别列将CSV拆分为多个小文件，要求按该列同类值分组并生成对应CSV文件。但目前选中列1进行拆分时，脚本却按列3执行分组操作。以下是使用pandas库实现的代码：

import pandas as pd

def print_column_options(df):
  # Print a numbered list of the columns in the DataFrame
  for i, column_name in enumerate(df.columns, start=1):
    print(f"{i}. {column_name}")

def get_column_name(df):
  num_columns = len(df.columns)
  column_number = int(input("Enter the number of the column to use for splitting the file: "))
  if column_number < 1 or column_number > num_columns:
    print(f"Error: columns are effed")
    return None
    return df.columns[column_number - 1]

def split_csv(filename, column_name, output_path="./"):
  # Read the input CSV file into a pandas DataFrame
  df = pd.read_csv(filename)

  # Split the DataFrame into multiple DataFrames based on the values in the specified column
  split_dfs = df.groupby(by=column_name, sort=False)

  # Save each split DataFrame to a separate CSV file
  for value, df in split_dfs:
    df.to_csv(output_path + value + ".csv", index=False)

# Prompt the user for the file path of the input CSV file
filename = "subscripts/output-cleaned.csv"

# Read the input CSV file into a pandas DataFrame
df = pd.read_csv(filename)

# Print the available column options
print("Column options:")
print_column_options(df)

# Prompt the user to select a column by number
column_name = get_column_name(df)

# Split the CSV file into multiple files based on the values in the specified column
split_csv(filename, column_name, output_path="subscripts/csvs-by-category/")

# Output a message indicating that the script has finished processing the data
print("Finished processing the data.")

我已多次重构脚本，此前使用标准csv库无法完成导出，转而使用pandas库，但问题仍未解决。

错误原因

核心问题在get_column_name函数的缩进错误：

当输入合法的列编号时，return df.columns[column_number - 1]语句被放在return None之后，属于不可达代码，函数会默认返回None。
当split_csv接收column_name=None时，pandas的groupby触发未定义行为，最终表现为错误使用列3进行分组。

修复后的代码

import pandas as pd

def print_column_options(df):
    # 打印带编号的列列表
    for i, column_name in enumerate(df.columns, start=1):
        print(f"{i}. {column_name}")

def get_column_name(df):
    num_columns = len(df.columns)
    try:
        column_number = int(input("Enter the number of the column to use for splitting the file: "))
        if 1 <= column_number <= num_columns:
            return df.columns[column_number - 1]
        else:
            print(f"Error: Column number must be between 1 and {num_columns}")
            return None
    except ValueError:
        print("Error: Please enter a valid integer")
        return None

def split_csv(filename, column_name, output_path="./"):
    if not column_name:
        print("Cannot split: No valid column selected")
        return
    
    df = pd.read_csv(filename)
    # 按指定列分组
    split_dfs = df.groupby(by=column_name, sort=False)
    
    # 保存每个分组到独立CSV
    for value, group_df in split_dfs:
        # 处理文件名中的非法字符
        safe_value = str(value).replace('/', '_').replace('\\', '_').replace(':', '_')
        output_file = f"{output_path}{safe_value}.csv"
        group_df.to_csv(output_file, index=False)

filename = "subscripts/output-cleaned.csv"

# 处理文件不存在的情况
try:
    df = pd.read_csv(filename)
except FileNotFoundError:
    print(f"Error: File {filename} not found")
    exit()

print("Column options:")
print_column_options(df)

column_name = get_column_name(df)

split_csv(filename, column_name, output_path="subscripts/csvs-by-category/")

print("Finished processing the data.")

修复说明

修正缩进错误：将return df.columns[column_number - 1]移到if块外部，确保合法输入时返回正确列名。
增强输入验证：增加整数转换异常捕获、列范围检查，提供清晰的错误提示。
增加前置检查：在split_csv中先验证column_name是否有效，避免无效操作。
文件名安全处理：替换分组值中的非法路径字符，防止保存失败。
文件读取容错：捕获文件不存在的异常，避免脚本崩溃。

内容的提问来源于stack exchange，提问作者SpikeyToes