@@ -215,18 +215,31 @@ class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkCon
215215 }
216216
217217 test(" intermediate storage level" ) {
218- val friends = Graphs .friends
219- val expected = Set (Set (" a" , " b" , " c" , " d" , " e" , " f" ), Set (" g" ))
220-
221- val cc = friends.connectedComponents
222- assert(cc.getIntermediateStorageLevel === StorageLevel .MEMORY_AND_DISK )
223-
224- for (storageLevel <- Seq (StorageLevel .DISK_ONLY , StorageLevel .MEMORY_ONLY , StorageLevel .NONE )) {
225- // TODO: it is not trivial to confirm the actual storage level used
226- val components = cc
227- .setIntermediateStorageLevel(storageLevel)
228- .run()
229- assertComponents(components, expected)
218+ // disabling adaptive query execution helps assertComponents
219+ val enabled = spark.conf.getOption(" spark.sql.adaptive.enabled" )
220+ try {
221+ spark.conf.set(" spark.sql.adaptive.enabled" , value = false )
222+
223+ val friends = Graphs .friends
224+ val expected = Set (Set (" a" , " b" , " c" , " d" , " e" , " f" ), Set (" g" ))
225+
226+ val cc = friends.connectedComponents
227+ assert(cc.getIntermediateStorageLevel === StorageLevel .MEMORY_AND_DISK )
228+
229+ for (storageLevel <- Seq (StorageLevel .DISK_ONLY , StorageLevel .MEMORY_ONLY , StorageLevel .NONE )) {
230+ // TODO: it is not trivial to confirm the actual storage level used
231+ val components = cc
232+ .setIntermediateStorageLevel(storageLevel)
233+ .run()
234+ assertComponents(components, expected)
235+ }
236+ } finally {
237+ // restoring earlier conf
238+ if (enabled.isDefined) {
239+ spark.conf.set(" spark.sql.adaptive.enabled" , value = enabled.get)
240+ } else {
241+ spark.conf.unset(" spark.sql.adaptive.enabled" )
242+ }
230243 }
231244 }
232245
0 commit comments