GroupBy with ApplyInPandas in PySpark - how to implement a UDF correctly?

I am trying to use a PandasUDF in PySpark to find the 'longest unique tail' in a hierarchy.

For example, if my input is:

1.2
1.2.3

then the longest tail is '1.2.3'

I may also have multiple unique sets, for example:

1.2
1.2.3
5.6.7
5.6

in which case the output should be:

1.2.3
5.6.7

The approach I am using is:

• sort the input so that like rows are listed so that if a preceding row is 'contained' in a following row I can filter it out and return the longest unique rows only.

example input:

1.2.3
5.6.7
5.6
1.2

sorted becomes:

1.2
1.2.3
5.6
5.6.7

when I filter line on line, my output should be

1.2.3
5.6.7

I have tried two approaches. First is to write a function that loops through a DF sent into it as follows:

def getLongestTail(key, pdf) -> pd.DataFrame:
    sortedData = pdf.sort_values(by='value')
    for i in range(len(sortedData)-1):
        if sortedData.index(i+1).loc['value'].startswith(sortedData.loc['value']):
        sortedData.index(i+1) = False
        return pd.DataFrame(sortedData)

Second is to use a lambda function inline

def getLongestTail(pdf) -> pd.DataFrame:
    pdf = pdf.sort
    return (lambda x: pdf.shift(1).loc['value'].startswith(pdf.loc['value']))

I have also tried to decorate as follows:

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)

Here is my overall code:

import pyspark
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

from pyspark.sql.functions import pandas_udf, PandasUDFType
import pandas as pd
from pyspark.sql.types import *

simpleData = [
      ('A', '1.2.3'),
      ('A', '1.2'),
      ('B', '9.8'),
      ('A', '5.6.7.8'),
      ('B', '9'),
      ('B', '9.8.7'),
      ('A', '5')]

    schema = StructType([
   StructField("letter", StringType()),
   StructField("value", StringType())
])

def getLongestTail(pdf) -> pd.DataFrame:
    pdf = pdf.sort
    return pd.DataFrame((lambda x: pdf.loc['value'].startswith(pdf.shift(1).loc['value'])))


spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame(data=simpleData, schema = schema)

df_result = df.groupby('letter').applyInPandas(getLongestTail, schema=schema).show()

The errors being shown in my Jupyter notebook are showing worker crashed and errors relating to Py4JJavaError. a I am sure there is something basic I am missing - any comments appreciated.

Thank you.

===

error:

---------------------------------------------------------------------------
Py4JJavaError                             Traceback (most recent call last)
/tmp/ipykernel_34305/1009949605.py in <module>
      3 # df_grouped.show()
      4 
----> 5 df_result = df.groupby('letter').applyInPandas(getLongestTailL, schema=schema).show()
      6 

    324             value = OUTPUT_CONVERTER[type](answer[2:], gateway_client)
    325             if answer[1] == REFERENCE_TYPE:
--> 326                 raise Py4JJavaError(
    327                     "An error occurred while calling {0}{1}{2}.\n".
    328                     format(target_id, ".", name), value)