如何比较含Polars Date类型列的Polars DataFrame对象？

测试包含Polars原生pl.Date类型列的DataFrame等价性时，polars.testing.assert_frame_equal会抛出PanicException: not implemented异常，但使用Python标准库datetime.date类型则可正常比较。

复现代码 #

正常运行的示例（datetime.date） #

import datetime as dt
import polars as pl
from polars.testing import assert_frame_equal

assert_frame_equal(
    pl.DataFrame({"foo": [1], "bar": [dt.date(2000, 1, 1)]}),
    pl.DataFrame({"foo": [1], "bar": [dt.date(2000, 1, 1)]})
)

抛出异常的示例（pl.Date） #

import polars as pl
from polars.testing import assert_frame_equal

assert_frame_equal(
    pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]}),
    pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
)

有三种可行方法实现Polars Date类型DataFrame的等价比较：

• 方法1：使用DataFrame原生的equals方法
  Polars的DataFrame.equals()方法可以正确处理原生Date类型的比较，无需依赖测试库的断言函数：

  import polars as pl
  
  df1 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  df2 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  
  assert df1.equals(df2)
  

• 方法2：自定义列比较函数传入assert_frame_equal
  利用assert_frame_equal的check_equal参数，传入针对Date类型的比较逻辑：

  import polars as pl
  from polars.testing import assert_frame_equal
  
  def custom_col_compare(left: pl.Series, right: pl.Series) -> bool:
      if left.dtype == pl.Date and right.dtype == pl.Date:
          return (left == right).all()
      return left.equals(right)
  
  df1 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  df2 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  
  assert_frame_equal(df1, df2, check_equal=custom_col_compare)
  

• 方法3：转换为Python标准日期类型后比较
  将Polars Date列转换为datetime.date类型后再使用assert_frame_equal：

  import polars as pl
  from polars.testing import assert_frame_equal
  
  df1 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  df2 = pl.DataFrame({"foo": [1], "bar": [pl.Date(2000, 1, 1)]})
  
  # 转换列类型
  df1_converted = df1.with_columns(pl.col("bar").dt.to_python_datetime().cast(pl.Date))
  df2_converted = df2.with_columns(pl.col("bar").dt.to_python_datetime().cast(pl.Date))
  
  assert_frame_equal(df1_converted, df2_converted)
  

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