title: Pyspark/Databricks - org.apache.spark.SparkException: Job aborted due to stage failure:
body: <p>I'm trying to use Linear Regression on a simple dataframe with <strong>one feature</strong> and <strong>one label</strong> using Python pyspark in Databricks. However, I'm running into some issues with stage failure. I've reviewed many similar problems, but most of them are in Scala or are out of the scope of what I'm doing here.</p>
  
  <p><strong>Versions:</strong></p>
  
  <p>Notebook: Databricks 
  5.3 (includes Apache Spark 2.4.0, Scala 2.11)
  Python version: 2</p>
  
  <p><strong>Here's what I've done:</strong></p>
  
  <ol>
  <li>Original dataframe looks like this:</li>
  </ol>
  
  <pre><code>    df_red = df_extra.select('cca3', 'class', 'device_id').groupby('cca3').pivot('class').count()
  
      display(df_red)
  </code></pre>
  
  <p><a href="https://i.stack.imgur.com/IfhOq.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/IfhOq.png" alt="enter image description here"></a></p>
  
  <p>I want the 'mac' column as my label and the 'other' column as my single feature.</p>
  
  <p>2.Drop column 'cca3' and create label/feature</p>
  
  <pre><code>features = ['other']
  lr_data = df_red.drop('cca3').select(col('mac').alias('label'), *features)
  display(lr_data)
  </code></pre>
  
  <p><a href="https://i.stack.imgur.com/I4C60.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/I4C60.png" alt="enter image description here"></a></p>
  
  <ol start="3">
  <li>Create Vector Assembler and remove null values in dataframe</li>
  </ol>
  
  <pre><code>assembler = VectorAssembler(inputCols = features, outputCol = "features")
  output = assembler.transform(lr_data)
  new_lr_data = output.select("label", "features").where(col('label').isNotNull())
  new_lr_data.show()
  </code></pre>
  
  <p><a href="https://i.stack.imgur.com/oC7eI.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/oC7eI.png" alt="enter image description here"></a></p>
  
  <ol start="4">
  <li>Linear Regression Model Fit:</li>
  </ol>
  
  <pre><code># Fit the model
  lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
  lrModel = lr.fit(new_lr_data)
  
  # Print the coefficients and intercept for linear regression
  print("Coefficients: %s" % str(lrModel.coefficients))
  print("Intercept: %s" % str(lrModel.intercept))
  
  # Summarize the model over the training set and print out some metrics
  trainingSummary = lrModel.summary
  #print("numIterations: %d" % trainingSummary.totalIterations)
  #print("objectiveHistory: %s" % str(trainingSummary.objectiveHistory))
  #trainingSummary.residuals.show()
  #print("RMSE: %f" % trainingSummary.rootMeanSquaredError)
  #print("r2: %f" % trainingSummary.r2)
  </code></pre>
  
  <p>At this point I get the error below:</p>
  
  <blockquote>
    <p>org.apache.spark.SparkException: Job aborted due to stage failure:
    Task 73 in stage 979.0 failed 1 times, most recent failure: Lost task
    73.0 in stage 979.0 (TID 32624, localhost, executor driver): org.apache.spark.SparkException: Failed to execute user defined
    function($anonfun$4:
    (struct&lt;other_double_VectorAssembler_a2059b1f0691:double&gt;) =&gt;
    struct&lt;type:tinyint,size:int,indices:array&lt;int&gt;,values:array&lt;double&gt;&gt;)</p>
  </blockquote>
  
  <p>What causes the above error to occur within Databricks? Could this be because I only used one feature as opposed to many features (usually the case)?</p>
  
  <p>Any help is much appreciated!</p>