如何去除DataFrame转Pandas列表后整数列的末尾.0？

问题 #

我的机器学习模型要求输入为String和bigint类型的两列，准备输入数据的代码如下，预期输入格式为{'data': [['12.5', 12]]}，但将Spark DataFrame转为Pandas再转成列表后，bigint列的数据末尾出现了.0，请问如何去除？

现有代码 #

input_df = inference_df
test_df = input_df.select("avg_value", input_df.avg_value.alias("value"))
test_df = test_df.withColumn("avg_value", test_df.avg_value.cast("String"))
test_df = test_df.withColumn("value", f.round(test_df["value"], 0).cast("bigint"))
test_pddf = test_df.toPandas()
data = {"data": test_pddf.values.tolist()}

print(data)

当前输出 #

{'data': [['3341.650014545455', 3342.0], ['7020.98821', 7021.0], ['161.76136000000002', 162.0], ['160.047379', 160.0], ['11832.535650000002', 11833.0], ['4714.418000000001', 4714.0], ['5294.925936470588', 5295.0], ['9330.04428', 9330.0], [None, nan], ['1.5845243856846047', 2.0], ['29.144277718733324', 29.0], [None, nan]]}

期望输出 #

{'data': [['3341.650014545455', 3342], ['7020.98821', 7021], ['161.76136000000002', 162], ['160.047379', 160], ['11832.535650000002', 11833], ['4714.418000000001', 4714], ['5294.925936470588', 5295], ['9330.04428', 9330], [None, nan], ['1.5845243856846047', 2], ['29.144277718733324', 29], [None, nan]]}

解决方案 #

出现.0的核心原因是：Spark的bigint类型列包含null值时，转Pandas会被自动识别为float64类型（Pandas原生int类型不支持null）。以下是三种可行的解决方法：

方法1：转Pandas前用Spark填充空值（允许修改空值时使用） #

如果业务允许将value列的null值替换为特定整数（比如0），可以先填充再转类型，这样Pandas会识别为int类型：

from pyspark.sql import functions as f

input_df = inference_df
test_df = input_df.select("avg_value", input_df.avg_value.alias("value"))
test_df = test_df.withColumn("avg_value", test_df.avg_value.cast("String"))
# 用0填充空值后转bigint
test_df = test_df.withColumn("value", f.coalesce(f.round(test_df["value"], 0).cast("bigint"), f.lit(0)))
test_pddf = test_df.toPandas()
data = {"data": test_pddf.values.tolist()}

print(data)

方法2：转Pandas后修改为支持null的整数类型 #

保留空值的前提下，将Pandas列转为Int64类型（Pandas提供的支持null的整数类型）：

input_df = inference_df
test_df = input_df.select("avg_value", input_df.avg_value.alias("value"))
test_df = test_df.withColumn("avg_value", test_df.avg_value.cast("String"))
test_df = test_df.withColumn("value", f.round(test_df["value"], 0).cast("bigint"))
test_pddf = test_df.toPandas()
# 转换为支持空值的Int64类型
test_pddf['value'] = test_pddf['value'].astype('Int64')
data = {"data": test_pddf.values.tolist()}

print(data)

方法3：生成列表时手动转换数值 #

遍历列表，将非空的float数值转为整数，保留空值：

import pandas as pd
from pyspark.sql import functions as f

input_df = inference_df
test_df = input_df.select("avg_value", input_df.avg_value.alias("value"))
test_df = test_df.withColumn("avg_value", test_df.avg_value.cast("String"))
test_df = test_df.withColumn("value", f.round(test_df["value"], 0).cast("bigint"))
test_pddf = test_df.toPandas()

# 遍历处理每一行数据
data_list = []
for str_val, num_val in test_pddf.values.tolist():
    if not pd.isna(num_val):
        num_val = int(num_val)
    data_list.append([str_val, num_val])
data = {"data": data_list}

print(data)

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