diff --git a/src/main/scala/com/databricks/spark/sql/perf/mllib/feature/VectorAssembler.scala b/src/main/scala/com/databricks/spark/sql/perf/mllib/feature/VectorAssembler.scala
index 914abec..7e862b5 100644
--- a/src/main/scala/com/databricks/spark/sql/perf/mllib/feature/VectorAssembler.scala
+++ b/src/main/scala/com/databricks/spark/sql/perf/mllib/feature/VectorAssembler.scala
@@ -34,7 +34,6 @@ object VectorAssembler extends BenchmarkAlgorithm with TestFromTraining {
       df = df.withColumn(colName, sliceVec(col("features"), lit(fromIndex), lit(untilIndex)))
     }
     df.drop(col("features"))
-    df
   }
 
   override def getPipelineStage(ctx: MLBenchContext): PipelineStage = {