Pyspark with error self._sock.recv_into(b) socket.timeout: timed out

Question

The goal is to use a UDF to categorize rows. I am using pyspark on windows.

Using simple functions or operations like filter appear to work.

Any direction on how to address the timeout/socket failure would be helpful (see error below).

There are no nulls in the data.

from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType,StringType

def BreakDown(arr_value):
    start_year = arr_value[0]
    start_month = arr_value[1]
    end_year = arr_value[2]
    end_month = arr_value[3]
    curr_year = arr_value[4]
    curr_month = arr_value[5]
    if   (curr_year == start_year) & (curr_month >= start_month) : return 1
    elif   (curr_year == end_year) & (curr_month <= end_month) : return 1
    elif   (curr_year > start_year) & (curr_year < end_year) : return 1
    else: return 0

    
udfBreakDown = udf(BreakDown, IntegerType())

temp = temp.withColumn('include', udfBreakDown(F.struct('start_year','start_month','end_year','end_month','curr_year','curr_month')))

PythonException: An exception was thrown from the Python worker. Please see the stack trace below. Traceback (most recent call last):
File "E:\spark\spark-3.0.1-bin-hadoop2.7\python\lib\pyspark.zip\pyspark\worker.py", line 585, in main File "E:\spark\spark-3.0.1-bin-hadoop2.7\python\lib\pyspark.zip\pyspark\serializers.py", line 593, in read_int length = stream.read(4) File "C:\ProgramData\Anaconda3\lib\socket.py", line 669, in readinto return self._sock.recv_into(b) socket.timeout: timed out

Answer 1

Always avoid using UDFs when you can use Spark built-in functions. You can rewrite your logic using when function like this:

from pyspark.sql import functions as F

def get_include_col():
    c = F.when((F.col("curr_year") == F.col("start_year")) & (F.col("curr_month") >= F.col("start_month")), F.lit(1)) \
        .when((F.col("curr_year") == F.col("end_year")) & (F.col("curr_month") <= F.col("end_month")), F.lit(1)) \
        .when((F.col("curr_year") > F.col("start_year")) & (F.col("curr_year") < F.col("end_year")), F.lit(1)) \
        .otherwise(F.lit(0))
    return c


temp = temp.withColumn('include', get_include_col())

You can also use functools.reduce to dynamically generate the when expressions without having to tape all of them. For example:

import functools
from pyspark.sql import functions as F

cases = [
    ("curr_year = start_year and curr_month >= start_month", 1),
    ("curr_year = end_year and curr_month <= end_month", 1),
    ("curr_year > start_year and curr_year < end_year", 1)
]

include_col = functools.reduce(
    lambda acc, x: acc.when(F.expr(x[0]), F.lit(x[1])),
    cases,
    F
).otherwise(F.lit(0))

temp = temp.withColumn('include', include_col)