data frame: How to round the values in some columns in Databricks Pyspark

Pyspark python in databricks

I have the following dataframe already created in Databricks.
The values in some columns should be rounded to integer only, which means 4.5 round to 5, 2.3 round to 2

df1:

the values in Number and Quantity columns need to be rounded. the other columns should be kept the same.

I expect the df1 could be changed to like below

BTW: 2.5 should be rounded to 3, not 2

Thanks.

You can use the round function to achieve the desired solution.

from pyspark.sql.functions import round

df = spark.createDataFrame(
        [["A", 1.236, "USA", 9.5, "PC"], ["B", 2.8988, "France",58.12999, "PC"]],
        ["Name", "Number", "Location", "Quantity", "UoM"])

df.withColumn("Number", round("Number")).withColumn("Quantity", round("Quantity")).show()

This outputs the columns in floating points but you can cast it to integer if required.

+----+------+--------+--------+---+
|Name|Number|Location|Quantity|UoM|
+----+------+--------+--------+---+
|   A|   1.0|     USA|    11.0| PC|
|   B|   3.0|  France|    58.0| PC|
+----+------+--------+--------+---+

Assuming based on DataBricks documentation that you are referring to a pandas DataFrame, you can use apply on the dataframe.

import pandas as pd

df = pd.DataFrame([[1.1, "a", 2.22222222222], [1.5, "b", 5.5555555555], [1.7, "c", 6.666666]], columns=["0", "1", "2"])
print(temp)

"""
     0  1         2
0  1.1  a  2.222222
1  1.5  b  5.555556
2  1.7  c  6.666666
"""

df[["0", "2"]] = df[["0", "2"]].apply(round).astype("int32")
print(temp)

"""
   0  1  2
0  1  a  2
1  2  b  6
2  2  c  7
"""

On top of using .apply on the specific columns that you want to change, you would notice that I also call .astype on the results. This is because even though round would return an integer since no ndigits is supplied, the DataFrame has those columns registered as a float and would have the values casted as a float. Converting it with .astype("int32") would turn them back into integer values. You could change the type depending on need as well, such as int64 for very large values.

I am not sure why previous post talked about Pandas.

For pyspark, you can create a custom function and then use it :

def custom_round(column):
    return F.when(F.col(column) - F.floor(F.col(column)) >= 0.5, F.ceil(F.col(column))).otherwise(F.floor(F.col(column)))

df = df.withColumn("RoundedQuantity", custom_round_udf(F.col("Quantity")))