python - DASK:当一方的 key 为 NA 时，合并会抛出错误，而 pd.merge 有效

Question

我有这些示例数据框:

tdf1 = pd.DataFrame([{"id": 1, "val": 4}, {"id": 2, "val": 5}, {"id": 3, "val": 6}, {"id": pd.NA, "val": 7}, {"id": 4, "val": 8}])
tdf2 = pd.DataFrame([{"some_id": 1, "name": "Josh"}, {"some_id": 3, "name": "Jake"}])

pd.merge(tdf1, tdf2, how="left", left_on="id", right_on="some_id").head()

并且合并工作完美。现在，如果我想使用 Dask 进行同样的操作:

dd_tdf1 = dd.from_pandas(tdf1, npartitions=10)
dd_tdf2 = dd.from_pandas(tdf2, npartitions=10)

dd_tdf2.merge(dd_tdf1, left_on="some_id", right_on="id", how="right", npartitions=10).compute(scheduler="threads").head()

我收到以下错误:

File /opt/conda/lib/python3.10/site-packages/pandas/core/reshape/merge.py:1585, in <genexpr>(.0)
   1581         return _get_no_sort_one_missing_indexer(left_n, False)
   1583 # get left & right join labels and num. of levels at each location
   1584 mapped = (
-> 1585     _factorize_keys(left_keys[n], right_keys[n], sort=sort, how=how)
   1586     for n in range(len(left_keys))
   1587 )
   1588 zipped = zip(*mapped)
   1589 llab, rlab, shape = (list(x) for x in zipped)

File /opt/conda/lib/python3.10/site-packages/pandas/core/reshape/merge.py:2313, in _factorize_keys(lk, rk, sort, how)
   2309 if is_integer_dtype(lk.dtype) and is_integer_dtype(rk.dtype):
   2310     # GH#23917 TODO: needs tests for case where lk is integer-dtype
   2311     #  and rk is datetime-dtype
   2312     klass = libhashtable.Int64Factorizer
-> 2313     lk = ensure_int64(np.asarray(lk))
   2314     rk = ensure_int64(np.asarray(rk))
   2316 elif needs_i8_conversion(lk.dtype) and is_dtype_equal(lk.dtype, rk.dtype):
   2317     # GH#23917 TODO: Needs tests for non-matching dtypes

File pandas/_libs/algos_common_helper.pxi:86, in pandas._libs.algos.ensure_int64()

TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NAType'

有没有人知道为什么会发生这种情况？

Answer 1

这扩展了@Michael Delgado 的评论。棘手的是 int64 不支持 pd.NA。你可以看到这个错误是由 pandas 通过尝试引发的:

tdf1['id'] = tdf1['id'].astype('int64')
# TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NAType'

一个修复是确保 id 的 dtype 设置为 Int64(支持 pd.NA):

tdf1['id'] = tdf1['id'].astype('Int64')
tdf2['some_id'] = tdf2['some_id'].astype('Int64')

一旦数据类型相同，您就可以运行 dask 片段并获得所需的合并。