mbejarano89
Databricks Partner

I found the answer just by trying several things, although I do not understand exactly what the problem was. All I had to do was to cache the input data before fitting the model:

assemble=VectorAssembler(inputCols=columns_input, outputCol='features')
assembled_data=assemble.transform(df)
assembled_data = assembled_data.cache()
KMeans_algo=KMeans(featuresCol='features', k=number_of_clusters)            
KMeans_fit=KMeans_algo.fit(assembled_data)