title: GroupBy with ApplyInPandas in PySpark - how to implement a UDF correctly?
body: <p>I am trying to use a PandasUDF in <strong>PySpark</strong> to find the 'longest unique tail' in a hierarchy.</p>
  <p>For example, if my input is:</p>
  <pre><code>1.2
  1.2.3
  </code></pre>
  <p>then the longest tail is '1.2.3'</p>
  <p>I may also have multiple unique sets, for example:</p>
  <pre><code>1.2
  1.2.3
  5.6.7
  5.6
  </code></pre>
  <p>in which case the output should be:</p>
  <pre><code>1.2.3
  5.6.7
  </code></pre>
  <p>The approach I am using is:</p>
  <ul>
  <li>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.</li>
  </ul>
  <p>example input:</p>
  <pre><code>1.2.3
  5.6.7
  5.6
  1.2
  </code></pre>
  <p>sorted becomes:</p>
  <pre><code>1.2
  1.2.3
  5.6
  5.6.7
  </code></pre>
  <p>when I filter line on line, my output should be</p>
  <pre><code>1.2.3
  5.6.7
  </code></pre>
  <p>I have tried two approaches.
  First is to write a function that loops through a DF sent into it as follows:</p>
  <pre><code>def getLongestTail(key, pdf) -&gt; 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)
  </code></pre>
  <p>Second is to use a lambda function inline</p>
  <pre><code>def getLongestTail(pdf) -&gt; pd.DataFrame:
      pdf = pdf.sort
      return (lambda x: pdf.shift(1).loc['value'].startswith(pdf.loc['value']))
  </code></pre>
  <p>I have also tried to decorate as follows:</p>
  <p>@pandas_udf(schema, PandasUDFType.GROUPED_MAP)</p>
  <p>Here is my overall code:</p>
  <pre><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(&quot;letter&quot;, StringType()),
     StructField(&quot;value&quot;, StringType())
  ])
  
  def getLongestTail(pdf) -&gt; 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()
  </code></pre>
  <p>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.</p>
  <p>Thank you.</p>
  <p>===</p>
  <p>error:</p>
  <pre><code>---------------------------------------------------------------------------
  Py4JJavaError                             Traceback (most recent call last)
  /tmp/ipykernel_34305/1009949605.py in &lt;module&gt;
        3 # df_grouped.show()
        4 
  ----&gt; 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:
  --&gt; 326                 raise Py4JJavaError(
      327                     &quot;An error occurred while calling {0}{1}{2}.\n&quot;.
      328                     format(target_id, &quot;.&quot;, name), value)
  </code></pre>