Skip to content

Commit 46307b2

Browse files
Pengsrowen
authored andcommitted
[SPARK-21401][ML][MLLIB] add poll function for BoundedPriorityQueue
## What changes were proposed in this pull request? The most of BoundedPriorityQueue usages in ML/MLLIB are: Get the value of BoundedPriorityQueue, then sort it. For example, in Word2Vec: pq.toSeq.sortBy(-_._2) in ALS, pq.toArray.sorted() The test results show using pq.poll is much faster than sort the value. It is good to add the poll function for BoundedPriorityQueue. ## How was this patch tested? The existing UT Author: Peng <peng.meng@intel.com> Author: Peng Meng <peng.meng@intel.com> Closes #18620 from mpjlu/add-poll.
1 parent ae253e5 commit 46307b2

File tree

2 files changed

+55
-0
lines changed

2 files changed

+55
-0
lines changed

core/src/main/scala/org/apache/spark/util/BoundedPriorityQueue.scala

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ private[spark] class BoundedPriorityQueue[A](maxSize: Int)(implicit ord: Orderin
5151
this
5252
}
5353

54+
def poll(): A = {
55+
underlying.poll()
56+
}
57+
5458
override def +=(elem1: A, elem2: A, elems: A*): this.type = {
5559
this += elem1 += elem2 ++= elems
5660
}
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.spark.util
19+
20+
import org.apache.spark.SparkFunSuite
21+
22+
class BoundedPriorityQueueSuite extends SparkFunSuite {
23+
test("BoundedPriorityQueue poll test") {
24+
val pq = new BoundedPriorityQueue[Double](4)
25+
26+
pq += 0.1
27+
pq += 1.5
28+
pq += 1.0
29+
pq += 0.3
30+
pq += 0.01
31+
32+
assert(pq.isEmpty == false)
33+
assert(pq.poll() == 0.1)
34+
assert(pq.poll() == 0.3)
35+
assert(pq.poll() == 1.0)
36+
assert(pq.poll() == 1.5)
37+
assert(pq.isEmpty == true)
38+
39+
val pq2 = new BoundedPriorityQueue[(Int, Double)](4)(Ordering.by(_._2))
40+
pq2 += 1 -> 0.5
41+
pq2 += 5 -> 0.1
42+
pq2 += 3 -> 0.3
43+
pq2 += 4 -> 0.2
44+
pq2 += 1 -> 0.4
45+
46+
assert(pq2.poll()._2 == 0.2)
47+
assert(pq2.poll()._2 == 0.3)
48+
assert(pq2.poll()._2 == 0.4)
49+
assert(pq2.poll()._2 == 0.5)
50+
}
51+
}

0 commit comments

Comments
 (0)