如何用Python创建类提取DataFrame中的月份与时间差？

首先准备示例数据，模拟包含time_1和time_2时间列的DataFrame：

import pandas as pd

# 构造示例数据
data = {
    'time_1': pd.date_range(start='2023-01-01', periods=5, freq='M'),
    'time_2': pd.date_range(start='2023-02-15', periods=5, freq='M')
}
df = pd.DataFrame(data)

接下来是实现时间特征提取的类，包含提取月份、计算时间差两个核心功能，同时加入列名校验避免出错：

class TimeFeatureExtractor:
    def __init__(self, df):
        self.df = df
        # 初始化时校验必要列是否存在
        self._validate_columns()
    
    def _validate_columns(self):
        required_cols = ['time_1', 'time_2']
        missing_cols = [col for col in required_cols if col not in self.df.columns]
        if missing_cols:
            raise ValueError(f"DataFrame缺少必要列: {', '.join(missing_cols)}")
    
    def extract_month(self, column='time_1'):
        """提取指定时间列的月份值，返回Series"""
        if column not in self.df.columns:
            raise ValueError(f"DataFrame中不存在列: {column}")
        # 确保列是datetime类型，不是的话尝试转换
        if not pd.api.types.is_datetime64_any_dtype(self.df[column]):
            self.df[column] = pd.to_datetime(self.df[column])
        return self.df[column].dt.month
    
    def calculate_time_diff_seconds(self):
        """计算time_2与time_1的时间差，返回总秒数的Series"""
        # 确保两列都是datetime类型
        for col in ['time_1', 'time_2']:
            if not pd.api.types.is_datetime64_any_dtype(self.df[col]):
                self.df[col] = pd.to_datetime(self.df[col])
        time_diff = self.df['time_2'] - self.df['time_1']
        return time_diff.dt.total_seconds()
    
    def add_features_to_df(self):
        """将提取的月份和时间差直接添加到原DataFrame中"""
        self.df['month_from_time1'] = self.extract_month()
        self.df['time_diff_seconds'] = self.calculate_time_diff_seconds()
        return self.df

类的使用示例 #

# 实例化提取器
extractor = TimeFeatureExtractor(df)

# 单独提取月份
month_series = extractor.extract_month()
print("提取的月份值:\n", month_series)

# 单独计算时间差（秒）
time_diff_series = extractor.calculate_time_diff_seconds()
print("时间差（秒）:\n", time_diff_series)

# 直接将特征添加到原DataFrame
updated_df = extractor.add_features_to_df()
print("添加特征后的DataFrame:\n", updated_df)

关键细节说明 #

• _validate_columns是私有方法，用于初始化时检查DataFrame是否包含time_1和time_2，避免后续方法报错
• extract_month支持指定时间列，默认提取time_1的月份，同时会自动尝试将非datetime类型的列转换为datetime格式
• calculate_time_diff_seconds内部做了类型校验，确保时间差计算的正确性
• add_features_to_df方法可以一键将两个特征添加到原DataFrame，方便后续分析

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