<jupyter_start><jupyter_text>`udf()` function and sql types` The `pyspark.sql.functions.udf()` function is a very important function. It allows us to transfer a **user defined function** to a **`pyspark.sql.functions`** function which can act on columns of a DataFrame. It makes data framsformation much more flexible.Using `udf()` could be tricky. The key is to understand how to define the `returnType` parameter.<jupyter_code>from pyspark.sql.types import *
from pyspark.sql.functions import udf
mtcars = spark.read.csv('../../data/mtcars.csv', inferSchema=True, header=True)
mtcars = mtcars.withColumnRenamed('_c0', 'model')
mtcars.show(5)<jupyter_output>+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+
|            model| mpg|cyl| disp| hp|drat|   wt| qsec| vs| am|gear|carb|
+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+
|        Mazda RX4|21.0|  6|160.0|110| 3.9| 2.62|16.46|  0|  1|   4|   4|
|    Mazda RX4 Wag|21.0|  6|160.0|110| 3.9|2.875|17.02|  0|  1|   4|   4|
|       Datsun 710|22.8|  4|108.0| 93|3.85| 2.32|18.61|  1|  1|   4|   1|
|   Hornet 4 Drive|21.4|  6|258.0|110|3.08|3.215|19.44|  1|  0|   3|   1|
|Hornet Sportabout|18.7|  8|360.0|175|3.15| 3.44|17.02|  0|  0|   3|   2|
+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+
only showing top 5 rows<jupyter_text>**The structure of the schema passed to `returnType` has to match the data structure of the return value from the user defined function**. **Case 1**: divide **disp** by **hp** and put the result to a new columnThe user defined function returns a float value.<jupyter_code>def disp_by_hp(disp, hp):
    return(disp/hp)
disp_by_hp_udf = udf(disp_by_hp, returnType=FloatType())
all_original_cols = [eval('mtcars.' + x) for x in mtcars.columns]
all_original_cols
disp_by_hp_col = disp_by_hp_udf(mtcars.disp, mtcars.hp)
disp_by_hp_col
all_new_cols = all_original_cols + [disp_by_hp_col]
all_new_cols
mtcars.select(all_new_cols).show()<jupyter_output>+-------------------+----+---+-----+---+----+-----+-----+---+---+----+----+--------------------+
|              model| mpg|cyl| disp| hp|drat|   wt| qsec| vs| am|gear|carb|disp_by_hp(disp, hp)|
+-------------------+----+---+-----+---+----+-----+-----+---+---+----+----+--------------------+
|          Mazda RX4|21.0|  6|160.0|110| 3.9| 2.62|16.46|  0|  1|   4|   4|           1.4545455|
|      Mazda RX4 Wag|21.0|  6|160.0|110| 3.9|2.875|17.02|  0|  1|   4|   4|           1.4545455|
|         Datsun 710|22.8|  4|108.0| 93|3.85| 2.32|18.61|  1|  1|   4|   1|           1.1612903|
|     Hornet 4 Drive|21.4|  6|258.0|110|3.08|3.215|19.44|  1|  0|   3|   1|           2.3454545|
|  <NAME>|18.7|  8|360.0|175|3.15| 3.44|17.02|  0|  0|   3|   2|            2.057143|
|            Valiant|18.1|  6|225.0|105|2.76| 3.46|20.22|  1|  0|   3|   1|            2.142857|
|         Duster 360|14.3|  8|360.0|245|3.21| 3.57|15.84|  0|  0|   3|   4|           1.4693878|
|          Merc 240D|24.4|  4|[...]<jupyter_text>**case 2**: create an array column that contain **disp** and **hp** values<jupyter_code># define function
def merge_two_columns(col1, col2):
    return([float(col1), float(col2)])

# convert user defined function into an udf function (sql function)
array_merge_two_columns_udf = udf(merge_two_columns, returnType=ArrayType(FloatType()))
array_col = array_merge_two_columns_udf(mtcars.disp, mtcars.hp)
array_col
all_new_cols = all_original_cols + [array_col]
all_new_cols
mtcars.select(all_new_cols).show(5, truncate=False)<jupyter_output>+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+---------------------------+
|model            |mpg |cyl|disp |hp |drat|wt   |qsec |vs |am |gear|carb|merge_two_columns(disp, hp)|
+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+---------------------------+
|Mazda RX4        |21.0|6  |160.0|110|3.9 |2.62 |16.46|0  |1  |4   |4   |[160.0, 110.0]             |
|Mazda RX4 Wag    |21.0|6  |160.0|110|3.9 |2.875|17.02|0  |1  |4   |4   |[160.0, 110.0]             |
|Datsun 710       |22.8|4  |108.0|93 |3.85|2.32 |18.61|1  |1  |4   |1   |[108.0, 93.0]              |
|Hornet 4 Drive   |21.4|6  |258.0|110|3.08|3.215|19.44|1  |0  |3   |1   |[258.0, 110.0]             |
|Hornet Sportabout|18.7|8  |360.0|175|3.15|3.44 |17.02|0  |0  |3   |2   |[360.0, 175.0]             |
+-----------------+----+---+-----+---+----+-----+-----+---+---+----+----+---------------------------+
only showing top 5 rows<jupyter_text>`ArrayType` vs. `StructType` Both `ArrayType` and `StructType` can be used to build `returnType` for a list. The difference is: 1. `ArrayType` requires all elements in the list have the same `elementType`, while `StructType` can have different `elementTypes`.2. `StructType` represents a `Row` object. **Define an `ArrayType` with elementType being `FloatType`.**<jupyter_code># define function
def merge_two_columns(col1, col2):
    return([float(col1), float(col2)])
array_type = ArrayType(FloatType())
array_merge_two_columns_udf = udf(merge_two_columns, returnType=array_type)<jupyter_output><empty_output><jupyter_text>**Define a `StructType` with one elementType being `StringType` and the other being `FloatType`.**<jupyter_code># define function
def merge_two_columns(col1, col2):
    return([str(col1), float(col2)])
struct_type = StructType([
    StructField('f1', StringType()),
    StructField('f2', FloatType())
])
struct_merge_two_columns_udf = udf(merge_two_columns, returnType=struct_type)<jupyter_output><empty_output><jupyter_text>**array column** expression: both values are float type values<jupyter_code>array_col = array_merge_two_columns_udf(mtcars.hp, mtcars.disp)
array_col<jupyter_output><empty_output><jupyter_text>**struct column** expression: first value is a string and the second value is a float type value.<jupyter_code>struct_col = struct_merge_two_columns_udf(mtcars.model, mtcars.disp)
struct_col<jupyter_output><empty_output><jupyter_text>**Results**<jupyter_code>mtcars.select(array_col, struct_col).show(truncate=False)<jupyter_output>+---------------------------+------------------------------+
|merge_two_columns(hp, disp)|merge_two_columns(model, disp)|
+---------------------------+------------------------------+
|[110.0, 160.0]             |[Mazda RX4, 160.0]            |
|[110.0, 160.0]             |[Mazda RX4 Wag, 160.0]        |
|[93.0, 108.0]              |[Datsun 710, 108.0]           |
|[110.0, 258.0]             |[Hornet 4 Drive, 258.0]       |
|[175.0, 360.0]             |[Hornet Sportabout, 360.0]    |
|[105.0, 225.0]             |[Valiant, 225.0]              |
|[245.0, 360.0]             |[Duster 360, 360.0]           |
|[62.0, 146.7]              |[Merc 240D, 146.7]            |
|[95.0, 140.8]              |[Merc 230, 140.8]             |
|[123.0, 167.6]             |[Merc 280, 167.6]             |
|[123.0, 167.6]             |[Merc 280C, 167.6]            |
|[180.0, 275.8]             |[Merc 450SE, 275.8]           |
|[180.0, 275.8]             |[Merc 450SL, 275.8]           |
|[180.0, 275.8]         [...]
{'fraction_non_alphanumeric': 0.2559255079006772, 'fraction_numerical': 0.10637697516930023, 'mean_word_length': 2.1478685612788633, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '127601', 'license': 'MIT', 'n_tokens_mistral': 3166, 'n_tokens_neox': 2758, 'n_words': 566, 'path': 'notebooks/04-miscellaneous/.ipynb_checkpoints/user-defined-sql-function (udf)-checkpoint.ipynb', 'repo_name': 'peymanyousefianasml/learning-apache-spark', 'size': '7099'}