使用SQLAlchemy加载Dask DataFrame失败：索引列类型异常

尝试使用dd.read_sql_query通过SQLAlchemy加载Dask DataFrame，定义的表中balance_date列在代码里设为DateTime类型，但数据库中实际是DATE类型：

class test_loans(Base):
      __tablename__ = 'test_loans'
      annual_income = Column(Float)
      balance = Column(Float)
      balance_date = Column(DateTime)  # 代码中设为DateTime类型
      cust_segment = Column(String)
      total_amount_paid = Column(Float)
      the_key = Column(Integer)
      __table_args__ = (PrimaryKeyConstraint(the_key),)

执行以下代码时失败：

stmt = select([ test_loans.balance_date, test_loans.total_amount_paid ]) 
ddf = dd.read_sql_query(stmt, con=con, index_col='balance_date', npartitions=3)

报错信息：

TypeError: Provided index column is of type "object".  If divisions is
not provided the index column type must be numeric or datetime.

疑问：该如何解决？这是否是一个缺陷？

方法1：显式转换查询中日期列的类型 #

在SQLAlchemy查询语句里，将balance_date强制转换为DateTime类型，确保Dask能识别为日期类型而非object：

from sqlalchemy import func

stmt = select([
    func.cast(test_loans.balance_date, DateTime).label('balance_date'),
    test_loans.total_amount_paid
])
ddf = dd.read_sql_query(stmt, con=con, index_col='balance_date', npartitions=3)

方法2：手动指定divisions参数 #

若不想修改查询，可手动计算日期范围并传入divisions参数，跳过Dask自动检测索引类型的步骤：

# 查询日期的最小和最大值
min_date = con.execute(select(func.min(test_loans.balance_date))).scalar()
max_date = con.execute(select(func.max(test_loans.balance_date))).scalar()

# 生成对应3个分区的divisions
from datetime import timedelta
delta = (max_date - min_date) // 3
divisions = [min_date, min_date + delta, min_date + delta*2, max_date]

# 加载数据
stmt = select([ test_loans.balance_date, test_loans.total_amount_paid ]) 
ddf = dd.read_sql_query(stmt, con=con, index_col='balance_date', npartitions=3, divisions=divisions)

方法3：修正模型列类型匹配数据库 #

既然数据库中balance_date是DATE类型，将模型里的对应列改为Date类型，保持与数据库一致：

from sqlalchemy import Date

class test_loans(Base):
      __tablename__ = 'test_loans'
      annual_income = Column(Float)
      balance = Column(Float)
      balance_date = Column(Date)  # 改为Date类型匹配数据库
      cust_segment = Column(String)
      total_amount_paid = Column(Float)
      the_key = Column(Integer)
      __table_args__ = (PrimaryKeyConstraint(the_key),)

修改后执行原查询，Dask可正确识别日期类型。

关于是否是缺陷 #

这不算缺陷，原因是数据库DATE类型与SQLAlchemy DateTime类型映射时，返回的对象可能是Python的date对象而非datetime对象，Dask自动检测时将date类型归为object，从而触发报错。只需让索引列明确为datetime类型，或手动指定divisions即可解决问题。

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