From a42f37a638c42a66853dd08645ce422256ab194a Mon Sep 17 00:00:00 2001
From: Hyun Goo Kang <hyungoo@twitter.com>
Date: Thu, 15 Mar 2012 02:13:56 -0700
Subject: [PATCH] Make DifferenceQuery obey page size - It will prevent
 resources from being wasted. - It will fix the flaky DifferenceQuerySpec
 test.

---
 .../scala/com/twitter/flockdb/ResultWindow.scala  | 15 ++++++++++-----
 .../twitter/flockdb/queries/DifferenceQuery.scala |  2 +-
 .../flockdb/unit/IntersectionQuerySpec.scala      |  3 +--
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/main/scala/com/twitter/flockdb/ResultWindow.scala b/src/main/scala/com/twitter/flockdb/ResultWindow.scala
index 33154604..f49158bc 100644
--- a/src/main/scala/com/twitter/flockdb/ResultWindow.scala
+++ b/src/main/scala/com/twitter/flockdb/ResultWindow.scala
@@ -58,14 +58,17 @@ class ResultWindow[T](val data: ResultWindowRows[T], val inNextCursor: Cursor, v
   val nextCursor = if (nextChanged && !page.isEmpty) page(page.size - 1).cursor else inNextCursor
   val prevCursor = if (prevChanged && !page.isEmpty) page(0).cursor.reverse else inPrevCursor
 
-  def ++(other: ResultWindow[T]) = {
+  def ++(other: ResultWindow[T]) = concat(other)
+
+  def concat(other: ResultWindow[T], newCount: Int = count) = {
     if (cursor < Cursor.Start) {
-      new ResultWindow(new ResultWindowRows(other.page ++ page), nextCursor, other.prevCursor, count, cursor)
+      new ResultWindow(new ResultWindowRows(other.page ++ page), nextCursor, other.prevCursor, newCount, cursor)
     } else {
-      new ResultWindow(new ResultWindowRows(page ++ other.page), other.nextCursor, prevCursor, count, cursor)
+      new ResultWindow(new ResultWindowRows(page ++ other.page), other.nextCursor, prevCursor, newCount, cursor)
     }
   }
 
+
   def merge(other: ResultWindow[T]) = {
     val newPage = Sorting.stableSort((Set((page ++ other.page): _*)).toSeq)
     val newNextCursor = if (nextCursor == Cursor.End && other.nextCursor == Cursor.End) Cursor.End else newPage(newPage.size - 1).cursor
@@ -73,12 +76,14 @@ class ResultWindow[T](val data: ResultWindowRows[T], val inNextCursor: Cursor, v
     new ResultWindow(new ResultWindowRows(newPage), newNextCursor, newPrevCursor, count, cursor)
   }
 
-  def --(values: Seq[T]) = {
+  def --(values: Seq[T]) = diff(values)
+
+  def diff(values: Seq[T], newCount: Int = count) = {
     val rejects = Set(values: _*)
     val newPage = page.filter { row => !rejects.contains(row.id) }
     val newNextCursor = if (nextCursor == Cursor.End || newPage.size == 0) Cursor.End else newPage(newPage.size - 1).cursor
     val newPrevCursor = if (prevCursor == Cursor.End || newPage.size == 0) Cursor.End else newPage(0).cursor.reverse
-    new ResultWindow(new ResultWindowRows(newPage), newNextCursor, newPrevCursor, count, cursor)
+    new ResultWindow(new ResultWindowRows(newPage), newNextCursor, newPrevCursor, newCount, cursor)
   }
 
   def length = page.length
diff --git a/src/main/scala/com/twitter/flockdb/queries/DifferenceQuery.scala b/src/main/scala/com/twitter/flockdb/queries/DifferenceQuery.scala
index 9c37bb9a..a718d436 100644
--- a/src/main/scala/com/twitter/flockdb/queries/DifferenceQuery.scala
+++ b/src/main/scala/com/twitter/flockdb/queries/DifferenceQuery.scala
@@ -51,7 +51,7 @@ class DifferenceQuery(query1: QueryTree, query2: QueryTree, averageIntersectionP
   private def pageDifference(internalPageSize: Int, count: Int, cursor: Cursor) = {
     val results = query1.selectPageByDestinationId(internalPageSize, cursor)
     val rejects = query2.selectWhereIn(results.view)
-    results -- rejects
+    results.diff(rejects, count)
   }
 
   override def toString =
diff --git a/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala b/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala
index 87ffd320..9bb8bcf0 100644
--- a/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala
+++ b/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala
@@ -29,11 +29,10 @@ object IntersectionQuerySpec extends ConfiguredSpecification with JMocker {
     val query1 = new queries.SeqQuery(List(1,2,3,4,5,6,7,8,9,10))
     val query2 = new queries.SeqQuery(List(1,2,3,4,11))
     val queryConfig = config.intersectionQuery
-    queryConfig.averageIntersectionProportion = 1.0
 
     "sizeEstimate" in {
       val intersectionQuery = queryConfig.intersect(query1, query2)
-      intersectionQuery.sizeEstimate() mustEqual 5
+      intersectionQuery.sizeEstimate() mustEqual (5 * queryConfig.averageIntersectionProportion).toInt
     }
 
     "selectWhereIn" in {