有没有一种方法可以表示要存储在 Spark DF 中的复数?

Is there a way to represent complex numbers to store in Spark DF?

我有 ndarray 个数据类型为 numpy.complex128 的值。当我尝试使用这些值创建 Spark DF 时,出现错误:

UserWarning: createDataFrame attempted Arrow optimization because 'spark.sql.execution.arrow.enabled' is set to true; however, failed by the reason below:
Unsupported numpy type 15
Attempting non-optimization as 'spark.sql.execution.arrow.fallback.enabled' is set to true.
TypeError: not supported type: <class 'complex'>

有没有人运行遇到过这样的情况?我如何表示这些复数,记住我最终将需要检索它们?

复数就是一对浮点数。如果你有一个形状为 (n1, n2, ..., nZ) 并输入 complex128 的 numpy 数组,你可以 view 它作为一个形状为 (n1, n2, ..., 2 * nZ) 的数组并输入 float64:

>>> a = np.linspace(0.+1.j, 1.+0j, 12).reshape(3, 4)
>>> a.shape
(3, 4)
>>> a.dtype
dtype('complex128')

>>> b = a.view(np.float64)
>>> b.shape
(3, 8)
>>> b.dtype
np.float64

实部和虚部占据数组的所有其他元素。您可以验证数据不会因查看为兼容的数据类型而发生变化:

>>> (b[:, ::2] == a.real).all()
True
>>> (b[:, 1::2] == a.imag).all()
True

该操作非常便宜:在相同数据上创建具有不同步幅的新数组对象。反序列化时,您可以简单地重新设置形状数组 (n1, n2, ..., 2 * nZ) 并将 float64 键入形状 (n1, n2, ..., nZ) 之一并键入 complex128:

>>> a2 = b.view(np.complex128)
>>> a2.shape
(3, 4)
>>> a2.dtype
dtype('complex128')