Python Excel 小技巧：通用筛选函数

案例目的

第一个函数从指定文件路径读取 Excel 文件并转换为DataFrame（指定读取某一个sheet，可以选择指定其中一列为行索引和第一行为列索引），第二个函数使用灵活的条件筛选DataFrame，第三个函数将筛选后的多个dataframe储存到一个.xlsx文件的不同sheet中。

示例数据

!+&	idx	Market	CURR	PM	TERMANT	……
*	1	JP	USD	1	0	…
*	1	CHINA	EUR	0	0	…
*	1	USA	USD	1	0	…
*	2	JP	JPY	1	0	…
*	3	USA	CNY	1	1	…
*	4	CHINA	CNY	0	0	…
*	5	JP	USD	1	1	…
*	6	JP	JPY	0	0	…

假定数据来源是字典

import pandas as pd
import csv
from io import StringIO

def csv_to_dataframe(csv_data=None, file_path=None, delimiter='\t'):
    """
    将CSV数据转换为DataFrame
    
    参数:
    csv_data: CSV格式的字符串数据
    file_path: CSV文件路径
    delimiter: 分隔符，默认为制表符
    
    返回:
    pandas DataFrame对象
    """
    if csv_data:
        # 从字符串读取数据
        df = pd.read_csv(StringIO(csv_data), delimiter=delimiter)
    elif file_path:
        # 从文件读取数据
        df = pd.read_csv(file_path, delimiter=delimiter)
    else:
        raise ValueError("必须提供csv_data或file_path参数")
    
    return df

def filter_dataframe(df, **kwargs):
    """
    根据多列条件筛选DataFrame
    
    参数:
    df: 要筛选的DataFrame
    **kwargs: 筛选条件，格式为 列名=值 或 列名=[值1, 值2, ...]
    
    返回:
    筛选后的DataFrame
    """
    # 初始化筛选条件
    condition = pd.Series([True] * len(df))
    
    # 处理每个筛选条件
    for column, values in kwargs.items():
        if column not in df.columns:
            raise ValueError(f"列名 '{column}' 不存在于DataFrame中")
        
        # 如果值是列表，使用isin方法
        if isinstance(values, list):
            condition = condition & df[column].isin(values)
        # 如果是单个值，使用等号比较
        else:
            condition = condition & (df[column] == values)
    
    # 应用筛选条件
    filtered_df = df[condition]
    
    return filtered_df

def filter_dataframe(df, **kwargs):
    """
    根据多列条件筛选DataFrame
    Filter a pandas DataFrame by multiple column conditions.
    参数Parameters:
    df: 要筛选的DataFrame pd.DataFrame
    **kwargs: 筛选条件，格式为 列名=值 或 列名=[值1, 值2, ...]
    **kwargs : Any
    Filtering criteria supplied as keyword arguments in the form
    column=value or column=[value1, value2, ...].
    返回:
    筛选后的DataFrame
    """
    if df is None or df.empty:
        print("Warning: DataFrame is empty or not provided; returning empty DataFrame.")
        return pd.DataFrame()  # 返回空DataFrame而不是None，保持类型一致性
    
    # 初始化筛选条件 Initialize a Boolean mask with all True
    condition = pd.Series([True] * len(df), index=df.index)
    
    # 处理每个筛选条件 Build the combined Boolean mask
    for column, values in kwargs.items():
        if column not in df.columns:
            print(f"Warning: Column '{column}' not found in DataFrame!")
            continue
        
        # 如果值是列表，使用isin方法
        if isinstance(values, list):
            condition = condition & df[column].isin(values)
        # 如果是单个值，使用等号比较
        else:
            condition = condition & (df[column] == values)
    
    # 应用筛选条件
    # filtered_df = df[condition]
    
    # print(f"筛选后数据形状: {filtered_df.shape}")
    return df[condition]

# 示例使用
if __name__ == "__main__":
    # 示例数据
    csv_data = """!+&	idx	Market	CURR	PM	TERMANT
*	1	JP	USD	1	0
*	1	CHINA	EUD	0	0
*	1	USA	USD	1	0
*	2	JP	JPY	1	0
*	3	USA	CNY	1	1
*	4	CHINA	CNY	0	0
*	5	JP	USD	1	1
*	6	JP	JPY	0	0"""
    
    # 使用第一个函数将CSV数据转换为DataFrame
    df = csv_to_dataframe(csv_data=csv_data, delimiter='\t')
    print("原始数据:")
    print(df)
    print("\n" + "="*50 + "\n")
    
    # 示例1: 筛选出Market为JP且CURR为JPY的所有行
    result1 = filter_dataframe(df, Market='JP', CURR='JPY')
    print("示例1 - Market为JP且CURR为JPY:")
    print(result1)
    print("\n" + "="*50 + "\n")
    
    # 示例2: 筛选出Market为日本和美国，CURR为USD和CNY
    result2 = filter_dataframe(df, Market=['JP', 'USA'], CURR=['USD', 'CNY'])
    print("示例2 - Market为JP或USA，CURR为USD或CNY:")
    print(result2)
    print("\n" + "="*50 + "\n")
    
    # 示例3: 筛选出Market为JP，CURR为USD、CNY和EUD
    result3 = filter_dataframe(df, Market='JP', CURR=['USD', 'CNY', 'EUD'])
    print("示例3 - Market为JP，CURR为USD、CNY或EUD:")
    print(result3)

这两个函数提供了灵活的筛选功能：

1. csv_to_dataframe 函数可以从字符串或文件读取CSV数据并转换为DataFrame
2. filter_dataframe 函数接受任意数量的列筛选条件，每个条件可以是单个值或值列表

1.0

import pandas as pd
import numpy as np
import re
from typing import Dict, Callable, Optional, Union, List, Any

def read_excel_to_dataframe(
    file_path: str,
    sheet_name: str,
    index_col: Optional[Union[int, str]] = None,
    header: int = 0,
    usecols: Optional[list] = None,
    dtype: Optional[dict] = None
) -> pd.DataFrame:
    """
    高效读取Excel文件到DataFrame
    
    参数:
        file_path: Excel文件路径
        sheet_name: 要读取的工作表名称
        index_col: 作为行索引的列（索引号或列名）
        header: 作为列名的行号（默认0）
        usecols: 需要读取的列（列表或列范围）
        dtype: 列数据类型映射（减少内存占用）
    
    返回:
        pandas DataFrame对象
    """
    # 使用openpyxl引擎提高大文件读取效率
    return pd.read_excel(
        file_path,
        sheet_name=sheet_name,
        index_col=index_col,
        header=header,
        usecols=usecols,
        dtype=dtype,
        engine='openpyxl'
    )

def filter_dataframe(
    df: pd.DataFrame,
    conditions: Optional[Dict[str, Callable]] = None,
    **kwargs
) -> pd.DataFrame:
    """
    灵活筛选DataFrame，支持函数条件和值匹配
    
    参数:
        df: 输入DataFrame
        conditions: {列名: 条件函数} 字典
        kwargs: 列名=值 或 列名=[值1,值2] 形式的简单匹配
        
    返回:
        筛选后的DataFrame
    """
    if df.empty:
        return df
    
    # 处理kwargs简单匹配条件
    for col, value in kwargs.items():
        if col not in df.columns:
            continue
            
        if isinstance(value, list):
            conditions = conditions or {}
            conditions[col] = lambda x, v=value: x.isin(v)
        else:
            conditions = conditions or {}
            conditions[col] = lambda x, v=value: x == v
    
    # 应用函数条件
    if conditions:
        mask = pd.Series(True, index=df.index)
        for column, condition_func in conditions.items():
            if column in df.columns:
                try:
                    mask &= condition_func(df[column])
                except Exception as e:
                    print(f"Error applying condition to column '{column}': {e}")
        return df.loc[mask]
    
    return df

# v2
def filter_dataframe(
    df: pd.DataFrame,
    conditions: Optional[Dict[str, Callable]] = None,
    **kwargs
) -> pd.DataFrame:
    """
    灵活筛选DataFrame，支持函数条件和值匹配（优化版本）
    
    参数:
        df: 输入DataFrame
        conditions: {列名: 条件函数} 字典，函数应返回布尔序列
        kwargs: 列名=值 或 列名=[值1,值2] 形式的简单匹配
        
    返回:
        筛选后的DataFrame
        
    示例:
        filter_dataframe(df, conditions={
            'Age': lambda x: x >= 30,
            'Salary': lambda x: x > 50000},
            Department='IT', 
            Status=['Active', 'Pending']
        )
    """
    if df.empty:
        return df
    
    # 创建掩码序列
    mask = pd.Series(True, index=df.index)
    
    # 处理kwargs简单匹配条件
    for col, value in kwargs.items():
        if col not in df.columns:
            continue
            
        if isinstance(value, list):
            # 使用isin进行列表匹配
            mask &= df[col].isin(value)
        else:
            # 直接比较
            mask &= (df[col] == value)
    
    # 应用函数条件
    if conditions:
        for column, condition_func in conditions.items():
            if column in df.columns:
                try:
                    # 应用条件函数
                    column_mask = condition_func(df[column])
                    # 确保返回的是布尔序列
                    if isinstance(column_mask, pd.Series) and column_mask.dtype == bool:
                        mask &= column_mask
                    else:
                        raise ValueError(f"条件函数对于列 '{column}' 未返回布尔序列")
                except Exception as e:
                    raise ValueError(f"应用条件到列 '{column}' 时出错: {str(e)}")
    
    return df.loc[mask]


def save_dataframes_to_excel(
    dataframes: Dict[str, pd.DataFrame],
    output_path: str,
    max_sheet_name_length: int = 31,
    max_rows_per_sheet: Optional[int] = None,
    engine: str = "auto"
) -> None:
    """
    将多个DataFrame保存到Excel的不同工作表，自动处理工作表名称和唯一性
    
    参数:
        dataframes: {工作表名称: DataFrame} 的字典
        output_path: 输出文件路径（推荐.xlsx格式）
        max_sheet_name_length: Excel工作表名称最大长度（默认31）
        max_rows_per_sheet: 每个工作表最大行数（超过自动分表）
        engine: 写入引擎（'auto', 'xlsxwriter', 'openpyxl'）
        
    异常处理:
        捕获并报告常见写入错误
    """
    # 自动选择最佳引擎
    if engine == "auto":
        engine = "xlsxwriter" if output_path.endswith(".xlsx") else "openpyxl"
    
    # 处理工作表名称
    def sanitize_sheet_name(name: str) -> str:
        """清理工作表名称使其符合Excel要求"""
        # 移除非法字符
        name = re.sub(r'[\\/*?:\[\]]', '', name)
        # 截断到最大长度
        if len(name) > max_sheet_name_length:
            name = name[:max_sheet_name_length]
        return name
    
    # 确保工作表名称唯一
    sheet_names = {}
    for original_name in dataframes.keys():
        sanitized = sanitize_sheet_name(original_name)
        base_name = sanitized
        counter = 1
        while sanitized in sheet_names.values():
            suffix = f"_{counter}"
            if len(base_name) + len(suffix) > max_sheet_name_length:
                truncate_length = max_sheet_name_length - len(suffix)
                base_name = base_name[:truncate_length]
            sanitized = f"{base_name}{suffix}"
            counter += 1
        sheet_names[original_name] = sanitized
    
    try:
        with pd.ExcelWriter(output_path, engine=engine) as writer:
            for original_name, df in dataframes.items():
                sheet_name = sheet_names[original_name]
                
                # 处理大型数据集
                if max_rows_per_sheet and len(df) > max_rows_per_sheet:
                    # 自动分表
                    for i, start in enumerate(range(0, len(df), max_rows_per_sheet)):
                        end = min(start + max_rows_per_sheet, len(df))
                        chunk = df.iloc[start:end]
                        chunk_sheet_name = f"{sheet_name}_part{i+1}"
                        
                        # 确保分表名称长度合法
                        if len(chunk_sheet_name) > max_sheet_name_length:
                            truncate_length = max_sheet_name_length - len(f"_part{i+1}")
                            chunk_sheet_name = f"{sheet_name[:truncate_length]}_part{i+1}"
                        
                        chunk.to_excel(
                            writer, 
                            sheet_name=chunk_sheet_name, 
                            index=False
                        )
                else:
                    df.to_excel(
                        writer, 
                        sheet_name=sheet_name, 
                        index=False
                    )
    
    except Exception as e:
        # 处理常见异常
        error_type = type(e).__name__
        
        if "Permission denied" in str(e):
            raise PermissionError(f"无法写入文件，请关闭已打开的Excel文件: {output_path}") from e
        elif "No such file or directory" in str(e):
            raise FileNotFoundError(f"路径不存在: {output_path}") from e
        elif "Workbook contains no default style" in str(e):
            raise ValueError("数据包含Excel不支持的样式或格式") from e
        elif "String longer than 32767 characters" in str(e):
            raise ValueError("单元格内容超过Excel最大长度限制(32767字符)") from e
        else:
            raise RuntimeError(f"保存Excel文件时出错: {str(e)}") from e

# 使用示例
if __name__ == "__main__":
    # 读取Excel
    
	# 推荐的数据类型优化映射
	dtype_mapping = {
	    'id': 'int32',          # 整型优化
	    'price': 'float32',     # 浮点优化
	    'category': 'category', # 分类数据
	    'flag': 'bool',         # 布尔类型
	    'date': 'datetime64'    # 日期类型
	}
    df = read_excel_to_dataframe(
        file_path='large_dataset.xlsx',
        sheet_name='SalesData',
        index_col='TransactionID',
        usecols=['TransactionID', 'Product', 'Quantity', 'Price', 'Date'],
        dtype={'Quantity': 'int32', 'Price': 'float32'}
    )
    
	# 1. 使用函数条件
	filtered = filter_dataframe(df, conditions={
	    'Age': lambda x: x >= 30,
	    'Salary': lambda x: x > 50000
	})
	
	# 2. 使用值匹配
	filtered = filter_dataframe(df, Department='IT', Status=['Active', 'Pending'])
	
	# 3. 混合使用
	filtered = filter_dataframe(df, 
	    conditions={'Salary': lambda x: x > 50000},
	    Department='IT'
	)
    
    # 保存结果（多个DataFrame示例）
    save_dataframes_to_excel(
        dataframes={
            'HighValueSales': filtered,
            'SummaryStats': pd.DataFrame(filtered.describe())
        },
        output_path='filtered_results.xlsx'
    )