基于 pyspark 中的列值的舍入

Question

I need to round off summary_measure_value based on reading_precision value

from pyspark.sql.functions import *
import pyspark.sql.functions as F
from pyspark.sql import *

df = spark.createDataFrame(
[(123, 2897402, 43.25, 2),
(124, 2897402, 49.25, 0),
(125, 2897402, 43.25, 2), 
(126, 2897402, 48.75, 0)]
, ['model_id','lab_test_id','summary_measure_value','reading_precision'])


partition_by_reading = [
    "model_id",
    "lab_test_id"
]
df.withColumn(
        "reading_value",
        round(avg("summary_measure_value").over(
                    Window.partitionBy(partition_by_reading))
                ,col("reading_precision"))).show()

I'm getting TypeError: 'Column' object is not callable

Answer 1

The pyspark round function expects a constant value for the scale/precision. You may have better luck creating a custom udf that applies the rounding logic.

I've included an example below from a test I've done based on your shared example :

udf_round = F.udf(lambda val,precision: round(val,precision))
df.withColumn(
        "reading_value",
        udf_round(F.avg("summary_measure_value").over(
                    Window.partitionBy(partition_by_reading))
                ,F.col("reading_precision"))).show()

NB.round in the udf above refers to the built-in python round function.

result:

model_id	lab_test_id	summary_measure_value	reading_precision	reading_value
125	2897402	43.25	2	43.25
124	2897402	49.25	0	49.0
123	2897402	43.25	2	43.25
126	2897402	48.75	0	49.0

Edit 1:

If you have erroneous data the following udf is more robust

@udf
def udf_round(val,precision) -> float:
    try:
        # ensure that value is a float and precision is an integer
        return float(round(float(val),int(precision)))
    except:
        return val # return val if there are any errors

Answer 2

Out of curiosity, applied a pandas_udf iterating multiple series. Seems faster than I thought

from pyspark.sql.functions import *
import pyspark.sql.functions as F
from pyspark.sql import *
from typing import Iterator, Tuple
from pyspark.sql.functions import struct, col


@pandas_udf('double')
def round_series(iterator: Iterator[Tuple[pd.Series, pd.DataFrame]]) -> Iterator[pd.Series]:
    return(a.round(b) for a,b in iterator)

df1=df.withColumn("reading_value",F.avg("summary_measure_value").over( Window.partitionBy("model_id", "lab_test_id"))).withColumn("reading_value",round_series('summary_measure_value','reading_precision')).show()

+--------+-----------+---------------------+-----------------+-------------+
|model_id|lab_test_id|summary_measure_value|reading_precision|reading_value|
+--------+-----------+---------------------+-----------------+-------------+
|     123|    2897402|                43.25|                2|        43.25|
|     124|    2897402|                49.11|                1|         49.1|
|     125|    2897402|                43.25|                2|        43.25|
|     126|    2897402|                48.75|                0|         49.0|
+--------+-----------+---------------------+-----------------+-------------+