如何在Psycopg3的COPY语句中将NaN转换为null？

背景：Psycopg2的实现方式 #

在Psycopg2中，可通过以下自定义适配器将Python的NaN值转换为PostgreSQL的NULL：

__REGISTERED = False

def _nan_to_null(f,
        _NULL=psycopg2.extensions.AsIs('NULL'),
        _Float=psycopg2.extensions.Float):
    if not np.isnan(f):
        return _Float(f)
    return _NULL

def register_nan_adapter():
    global __REGISTERED
    if not __REGISTERED:
        print('Register nan to null adapter for psycopg2...')
        psycopg2.extensions.register_adapter(float, _nan_to_null)
        __REGISTERED = True
    else:
        print('nan to null adapter for psycopg2 is already registered!')

问题：Psycopg3中的适配难题 #

尝试在Psycopg3中通过自定义FloatDumper实现相同功能时，最初的代码无法生效：

class NullNan(FloatDumper):
    def dump(self, elem):
        if np.isnan(elem):
            return b"NULL"
        else:
            return super().dump(elem)

connection.adapters.register_dumper(float, NullNan)

后续调整为返回None后，普通的插入、更新操作可正常工作，但在执行COPY语句时触发异常：

问题代码 #

class NullNan(FloatDumper):
    def dump(self, elem):
        if np.isnan(elem):
            return None
        else:
            return super().dump(elem)

with cursor.copy(f"COPY _r ({col_names_str}) FROM STDIN") as copy:
                for index, row in df.iterrows():
                    rec = row.values.tolist()
                    copy.write_row(rec)

报错信息 #

psycopg.errors.QueryCanceled: COPY from stdin failed: error from Python: TypeError - expected string or bytes-like object

解决方案 #

方案1：修改自定义Dumper适配COPY场景 #

COPY命令对NULL的处理逻辑和普通SQL操作不同，需要返回PostgreSQL COPY格式专用的b'\\N'来表示NULL。修改后的Dumper代码如下：

from psycopg.types.numeric import FloatDumper
import numpy as np

class NullNan(FloatDumper):
    def dump(self, elem):
        if np.isnan(elem):
            # 普通SQL操作返回None，COPY操作返回专用的NULL标识
            if self.ctx.format is None:
                return None
            else:
                return b'\\N'
        else:
            return super().dump(elem)

# 注册自定义Dumper
connection.adapters.register_dumper(float, NullNan)

方案2：提前处理行数据 #

在生成COPY的行数据时，直接将NaN替换为None，跳过Dumper的适配逻辑：

import numpy as np

with cursor.copy(f"COPY _r ({col_names_str}) FROM STDIN") as copy:
    for index, row in df.iterrows():
        # 遍历行数据，将所有NaN替换为None
        rec = [
            None if isinstance(x, float) and np.isnan(x) 
            else x 
            for x in row.values.tolist()
        ]
        copy.write_row(rec)

两种方案均可解决COPY语句的报错问题，同时保证插入、更新操作正常将NaN转换为数据库NULL。

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